diff --git a/pyproject.toml b/pyproject.toml
index 340abb385..6036ad96c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,6 @@ requires = [
     "packaging>=24.2",
     "setuptools>=77.0.3,<80.0.0",
     "setuptools-scm>=8.0",
-    "torch == 2.7.0",
     "wheel",
     "jinja2",
 ]
diff --git a/requirements/build.txt b/requirements/build.txt
index 528cd3b53..d0198efb4 100644
--- a/requirements/build.txt
+++ b/requirements/build.txt
@@ -4,7 +4,6 @@ ninja
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8
-torch==2.7.0
 wheel
 jinja2>=3.1.6
 regex
diff --git a/requirements/common.txt b/requirements/common.txt
index 8bc0be777..93e3f7ac6 100644
--- a/requirements/common.txt
+++ b/requirements/common.txt
@@ -2,12 +2,12 @@ regex # Replace re for higher-performance regex matching
 cachetools
 psutil
 sentencepiece  # Required for LLaMA tokenizer.
-numpy
+numpy==1.26.1
 requests >= 2.26.0
 tqdm
 blake3
 py-cpuinfo
-transformers >= 4.51.1
+transformers == 4.52.4
 huggingface-hub[hf_xet] >= 0.33.0  # Required for Xet downloads.
 tokenizers >= 0.21.1  # Required for fast incremental detokenization.
 protobuf # Required by LlamaTokenizer.
diff --git a/requirements/cpu-build.txt b/requirements/cpu-build.txt
index 37f072202..046cc61cb 100644
--- a/requirements/cpu-build.txt
+++ b/requirements/cpu-build.txt
@@ -1,12 +1,8 @@
-# Temporarily used for x86 CPU backend to avoid performance regression of torch>2.6.0+cpu,
-# see https://github.com/pytorch/pytorch/pull/151218
 cmake>=3.26.1
 ninja
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0
 setuptools-scm>=8
---extra-index-url https://download.pytorch.org/whl/cpu
-torch==2.6.0+cpu
 wheel
 jinja2>=3.1.6
 regex
diff --git a/requirements/cpu.txt b/requirements/cpu.txt
index df3a33935..b30cd7d5e 100644
--- a/requirements/cpu.txt
+++ b/requirements/cpu.txt
@@ -7,23 +7,10 @@ numba == 0.61.2; python_version > '3.9'
 # Dependencies for CPUs
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0
---extra-index-url https://download.pytorch.org/whl/cpu
-torch==2.6.0+cpu; platform_machine == "x86_64" # torch>2.6.0+cpu has performance regression on x86 platform, see https://github.com/pytorch/pytorch/pull/151218
-torch==2.7.0; platform_system == "Darwin"
-torch==2.7.0; platform_machine == "ppc64le" or platform_machine == "aarch64"
 
-# required for the image processor of minicpm-o-2_6, this must be updated alongside torch
-torchaudio; platform_machine != "ppc64le" and platform_machine != "s390x"
-torchaudio==2.7.0; platform_machine == "ppc64le"
 
-# required for the image processor of phi3v, this must be updated alongside torch
-torchvision; platform_machine != "ppc64le" and platform_machine != "s390x"
-torchvision==0.22.0; platform_machine == "ppc64le"
 datasets # for benchmark scripts
 
-# Intel Extension for PyTorch, only for x86_64 CPUs
 intel-openmp==2024.2.1; platform_machine == "x86_64"
-intel_extension_for_pytorch==2.6.0; platform_machine == "x86_64" # torch>2.6.0+cpu has performance regression on x86 platform, see https://github.com/pytorch/pytorch/pull/151218
 py-libnuma; platform_system != "Darwin"
 psutil; platform_system != "Darwin"
-triton==3.2.0; platform_machine == "x86_64" # Triton is required for torch 2.6+cpu, as it is imported in torch.compile.
diff --git a/requirements/cuda.txt b/requirements/cuda.txt
index a71d9728f..f36780514 100644
--- a/requirements/cuda.txt
+++ b/requirements/cuda.txt
@@ -6,9 +6,4 @@ numba == 0.61.2; python_version > '3.9'
 
 # Dependencies for NVIDIA GPUs
 ray[cgraph]>=2.43.0, !=2.44.* # Ray Compiled Graph, required for pipeline parallelism in V1.
-torch==2.7.0
-torchaudio==2.7.0
-# These must be updated alongside torch
-torchvision==0.22.0 # Required for phi3v processor. See https://github.com/pytorch/vision?tab=readme-ov-file#installation for corresponding version
 # https://github.com/facebookresearch/xformers/releases/tag/v0.0.30
-xformers==0.0.30; platform_system == 'Linux' and platform_machine == 'x86_64'  # Requires PyTorch >= 2.7
diff --git a/requirements/neuron.txt b/requirements/neuron.txt
index 7df478edd..a1cbd0d9e 100644
--- a/requirements/neuron.txt
+++ b/requirements/neuron.txt
@@ -4,6 +4,4 @@
 # Dependencies for Neuron devices
 packaging>=24.2
 setuptools>=77.0.3,<80.0.0
-torch-neuronx >= 2.5.0
 neuronx-cc>=2.0.0a0
-torchvision # Required for Llama3.2 multimodal image preprocessing
diff --git a/requirements/rocm-build.txt b/requirements/rocm-build.txt
index 94201543c..b6dfd41f7 100644
--- a/requirements/rocm-build.txt
+++ b/requirements/rocm-build.txt
@@ -1,10 +1,6 @@
 # Common dependencies
 -r common.txt
 
---extra-index-url https://download.pytorch.org/whl/rocm6.2.4
-torch==2.7.0
-torchvision==0.22.0
-torchaudio==2.7.0
 
 triton==3.2
 cmake>=3.26.1,<4
diff --git a/requirements/test.txt b/requirements/test.txt
index f6f599df7..68e1d278f 100644
--- a/requirements/test.txt
+++ b/requirements/test.txt
@@ -1,5 +1,4 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test.in -o requirements/test.txt --index-strategy unsafe-best-match --torch-backend cu128
 absl-py==2.1.0
     # via rouge-score
 accelerate==1.0.1
@@ -129,10 +128,8 @@ einops==0.8.0
     #   -r requirements/test.in
     #   encodec
     #   mamba-ssm
-    #   vector-quantize-pytorch
     #   vocos
 einx==0.3.0
-    # via vector-quantize-pytorch
 email-validator==2.2.0
     # via pydantic
 encodec==0.1.1
@@ -153,7 +150,6 @@ filelock==3.16.1
     #   datasets
     #   huggingface-hub
     #   ray
-    #   torch
     #   transformers
     #   virtualenv
 fonttools==4.54.1
@@ -173,7 +169,6 @@ fsspec==2024.9.0
     #   evaluate
     #   fastparquet
     #   huggingface-hub
-    #   torch
 genai-perf==0.0.8
     # via -r requirements/test.in
 genson==1.3.0
@@ -244,7 +239,6 @@ isort==5.13.2
 jinja2==3.1.6
     # via
     #   datamodel-code-generator
-    #   torch
 jiwer==3.0.5
     # via -r requirements/test.in
 jmespath==1.0.1
@@ -328,7 +322,6 @@ multiprocess==0.70.16
 mypy-extensions==1.0.0
     # via black
 networkx==3.2.1
-    # via torch
 ninja==1.11.1.3
     # via mamba-ssm
 nltk==3.9.1
@@ -373,7 +366,6 @@ numpy==1.26.4
     #   soxr
     #   statsmodels
     #   tensorizer
-    #   torchvision
     #   transformers
     #   tritonclient
     #   vocos
@@ -381,39 +373,25 @@ nvidia-cublas-cu12==12.8.3.14
     # via
     #   nvidia-cudnn-cu12
     #   nvidia-cusolver-cu12
-    #   torch
 nvidia-cuda-cupti-cu12==12.8.57
-    # via torch
 nvidia-cuda-nvrtc-cu12==12.8.61
-    # via torch
 nvidia-cuda-runtime-cu12==12.8.57
-    # via torch
 nvidia-cudnn-cu12==9.7.1.26
-    # via torch
 nvidia-cufft-cu12==11.3.3.41
-    # via torch
 nvidia-cufile-cu12==1.13.0.11
-    # via torch
 nvidia-curand-cu12==10.3.9.55
-    # via torch
 nvidia-cusolver-cu12==11.7.2.55
-    # via torch
 nvidia-cusparse-cu12==12.5.7.53
     # via
     #   nvidia-cusolver-cu12
-    #   torch
 nvidia-cusparselt-cu12==0.6.3
-    # via torch
 nvidia-nccl-cu12==2.26.2
-    # via torch
 nvidia-nvjitlink-cu12==12.8.61
     # via
     #   nvidia-cufft-cu12
     #   nvidia-cusolver-cu12
     #   nvidia-cusparse-cu12
-    #   torch
 nvidia-nvtx-cu12==12.8.55
-    # via torch
 opencensus==0.11.4
     # via ray
 opencensus-context==0.1.3
@@ -466,7 +444,6 @@ pillow==10.4.0
     #   matplotlib
     #   mistral-common
     #   sentence-transformers
-    #   torchvision
 platformdirs==4.3.6
     # via
     #   black
@@ -690,7 +667,6 @@ setuptools==77.0.3
     # via
     #   mamba-ssm
     #   pytablewriter
-    #   torch
     #   triton
 shellingham==1.5.4
     # via typer
@@ -728,7 +704,6 @@ statsmodels==0.14.4
 sympy==1.13.3
     # via
     #   einx
-    #   torch
 tabledata==1.3.3
     # via pytablewriter
 tabulate==0.9.0
@@ -757,7 +732,6 @@ tomli==2.2.1
     # via schemathesis
 tomli-w==1.2.0
     # via schemathesis
-torch==2.7.0+cu128
     # via
     #   -r requirements/test.in
     #   accelerate
@@ -772,16 +746,11 @@ torch==2.7.0+cu128
     #   sentence-transformers
     #   tensorizer
     #   timm
-    #   torchaudio
-    #   torchvision
-    #   vector-quantize-pytorch
     #   vocos
-torchaudio==2.7.0+cu128
     # via
     #   -r requirements/test.in
     #   encodec
     #   vocos
-torchvision==0.22.0+cu128
     # via
     #   -r requirements/test.in
     #   timm
@@ -812,7 +781,6 @@ transformers==4.52.4
 transformers-stream-generator==0.0.5
     # via -r requirements/test.in
 triton==3.3.0
-    # via torch
 tritonclient==2.51.0
     # via
     #   -r requirements/test.in
@@ -835,7 +803,6 @@ typing-extensions==4.12.2
     #   pqdm
     #   pydantic
     #   pydantic-core
-    #   torch
     #   typer
     #   typing-inspection
 typing-inspection==0.4.1
@@ -851,7 +818,6 @@ urllib3==2.2.3
     #   requests
     #   responses
     #   tritonclient
-vector-quantize-pytorch==1.21.2
     # via -r requirements/test.in
 virtualenv==20.31.2
     # via ray
diff --git a/requirements/tpu.txt b/requirements/tpu.txt
index 2b5fd8941..b9179169b 100644
--- a/requirements/tpu.txt
+++ b/requirements/tpu.txt
@@ -11,16 +11,9 @@ ray[default]
 ray[data]
 setuptools==78.1.0
 
-# Install torch_xla
 --pre
---extra-index-url https://download.pytorch.org/whl/nightly/cpu
 --find-links https://storage.googleapis.com/libtpu-wheels/index.html
 --find-links https://storage.googleapis.com/libtpu-releases/index.html
 --find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
 --find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
-torch==2.8.0.dev20250618
-torchvision==0.23.0.dev20250618
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250618-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250618-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.8.0.dev20250618-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
 
diff --git a/requirements/xpu.txt b/requirements/xpu.txt
index 0d95dc571..b8a140d00 100644
--- a/requirements/xpu.txt
+++ b/requirements/xpu.txt
@@ -11,14 +11,6 @@ jinja2>=3.1.6
 datasets # for benchmark scripts
 numba == 0.60.0 # v0.61 doesn't support Python 3.9. Required for N-gram speculative decoding
 
-torch==2.7.0+xpu
-torchaudio
-torchvision
-pytorch-triton-xpu
---extra-index-url=https://download.pytorch.org/whl/xpu
 
-# Please refer xpu doc, we need manually install intel-extension-for-pytorch 2.6.10+xpu due to there are some conflict dependencies with torch 2.6.0+xpu
 # FIXME: This will be fix in ipex 2.7. just leave this here for awareness.
-intel-extension-for-pytorch==2.7.10+xpu
 oneccl_bind_pt==2.7.0+xpu
---extra-index-url=https://pytorch-extension.intel.com/release-whl/stable/xpu/us/
diff --git a/setup.py b/setup.py
index ea7cd0169..2aa56dac4 100644
--- a/setup.py
+++ b/setup.py
@@ -686,7 +686,6 @@ setup(
     # static metadata should rather go in pyproject.toml
     version=get_vllm_version(),
     ext_modules=ext_modules,
-    install_requires=get_requirements(),
     extras_require={
         "bench": ["pandas", "datasets"],
         "tensorizer": ["tensorizer>=2.9.0"],
@@ -698,3 +697,11 @@ setup(
     cmdclass=cmdclass,
     package_data=package_data,
 )
+
+vllm_musa_dir = "vllm_musa"
+subprocess.run(
+    ["bash", "install.sh"],
+    cwd=vllm_musa_dir,
+    check=True,
+    text=True,
+)
diff --git a/vllm/compilation/collective_fusion.py b/vllm/compilation/collective_fusion.py
index f754fc238..78ffcfcd0 100644
--- a/vllm/compilation/collective_fusion.py
+++ b/vllm/compilation/collective_fusion.py
@@ -6,7 +6,7 @@ import torch
 import torch._inductor.pattern_matcher as pm
 import torch.fx as fx
 from torch._inductor.pattern_matcher import PatternMatcherPass
-from torch.distributed._symmetric_memory import enable_symm_mem_for_group
+# from torch.distributed._symmetric_memory import enable_symm_mem_for_group
 
 from vllm.config import VllmConfig
 from vllm.distributed import get_tp_group
@@ -104,7 +104,7 @@ class AsyncTPPass(VllmInductorPass):
         super().__init__(config)
 
         # Enable symmetric memory for the TP process group
-        enable_symm_mem_for_group(get_tp_group().device_group.group_name)
+        # enable_symm_mem_for_group(get_tp_group().device_group.group_name)
         self.patterns: PatternMatcherPass = PatternMatcherPass(
             pass_name="async_tp_pass")
         GEMMReduceScatterPattern(self.model_dtype,
diff --git a/vllm/compilation/fusion.py b/vllm/compilation/fusion.py
index 951a2861e..69df0dfc5 100644
--- a/vllm/compilation/fusion.py
+++ b/vllm/compilation/fusion.py
@@ -29,8 +29,8 @@ def empty_fp32(*args, **kwargs):
     return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda")
 
 
-RMS_OP = torch.ops._C.rms_norm.default
-RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
+RMS_OP = None
+RMS_ADD_OP = None
 
 
 # Use proxy as NamedTuple direct subclasses cannot have static members
@@ -87,12 +87,12 @@ kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, False, GroupShape.PER_TENSOR, True)
 kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, False, GroupShape.PER_TOKEN, True)
 
 QUANT_OPS: dict[QuantKey, OpOverload] = {
-    kFp8StaticTensorSym:
-    torch.ops._C.static_scaled_fp8_quant.default,  # noqa: E501
-    kFp8DynamicTensorSym:
-    torch.ops._C.dynamic_scaled_fp8_quant.default,  # noqa: E501
-    kFp8DynamicTokenSym:
-    torch.ops._C.dynamic_per_token_scaled_fp8_quant.default,  # noqa: E501
+    # kFp8StaticTensorSym:
+    # torch.ops._C.static_scaled_fp8_quant,  # noqa: E501
+    # kFp8DynamicTensorSym:
+    # torch.ops._C.dynamic_scaled_fp8_quant,  # noqa: E501
+    # kFp8DynamicTokenSym:
+    # torch.ops._C.dynamic_per_token_scaled_fp8_quant,  # noqa: E501
 }
 
 
@@ -111,14 +111,14 @@ class FusedRMSQuantKey(NamedTuple):
 
 
 FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {
-    FusedRMSQuantKey(kFp8StaticTensorSym, False):
-    torch.ops._C.rms_norm_static_fp8_quant.default,  # noqa: E501
-    FusedRMSQuantKey(kFp8StaticTensorSym, True):
-    torch.ops._C.fused_add_rms_norm_static_fp8_quant.default,  # noqa: E501
-    FusedRMSQuantKey(kFp8DynamicTokenSym, False):
-    torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa: E501
-    FusedRMSQuantKey(kFp8DynamicTokenSym, True):
-    torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa: E501
+#     FusedRMSQuantKey(kFp8StaticTensorSym, False):
+#     torch.ops._C.rms_norm_static_fp8_quant,  # noqa: E501
+#     FusedRMSQuantKey(kFp8StaticTensorSym, True):
+#     torch.ops._C.fused_add_rms_norm_static_fp8_quant,  # noqa: E501
+#     FusedRMSQuantKey(kFp8DynamicTokenSym, False):
+#     torch.ops._C.rms_norm_dynamic_per_token_quant,  # noqa: E501
+#     FusedRMSQuantKey(kFp8DynamicTokenSym, True):
+#     torch.ops._C.rms_norm_dynamic_per_token_quant,  # noqa: E501
 }
 
 
diff --git a/vllm/config.py b/vllm/config.py
index 724f69a38..44fbb621f 100644
--- a/vllm/config.py
+++ b/vllm/config.py
@@ -1131,7 +1131,7 @@ class ModelConfig:
         if not hasattr(self.hf_text_config, "model_type"):
             return False
         elif self.hf_text_config.model_type in \
-            ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp'):
+            ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp', 'kimi_k2'):
             return self.hf_text_config.kv_lora_rank is not None
         elif self.hf_text_config.model_type == 'eagle':
             # if the model is an EAGLE module, check for the
@@ -4081,6 +4081,9 @@ class CompilationConfig:
     splitting certain operations such as attention into subgraphs. Thus this
     flag cannot be used together with splitting_ops. This may provide
     performance benefits for smaller models."""
+    simple_cuda_graph: bool = False
+    """Whether to use simple CUDA graph, which uses full CUDA graphs without
+    torch.compile. This can speed up the startup time of vLLM."""
 
     pass_config: PassConfig = field(default_factory=PassConfig)
     """Custom inductor passes, see PassConfig for more details"""
@@ -4205,6 +4208,13 @@ class CompilationConfig:
         if isinstance(self.pass_config, dict):
             self.pass_config = PassConfig(**self.pass_config)
 
+        if self.simple_cuda_graph:
+            # When using simple CUDA graph, we skip torch.compile.
+            self.level = CompilationLevel.NO_COMPILATION
+            # Simple CUDA graph does not support piecewise CUDA graphs.
+            self.full_cuda_graph = True
+            self.use_cudagraph = True
+
     def init_backend(self, vllm_config: "VllmConfig") -> Union[str, Callable]:
         if self.level == CompilationLevel.NO_COMPILATION:
             raise ValueError("No compilation level is set.")
@@ -4559,9 +4569,14 @@ class VllmConfig:
             not self.model_config.enforce_eager:
             # By default, V1 uses piecewise CUDA graphs. If full_cuda_graph
             # is set to True, full CUDA graphs will be used.
+            self.use_cudagraph = True
             self.compilation_config.cudagraph_num_of_warmups = 1
-            self.compilation_config.level = CompilationLevel.PIECEWISE
-            self.compilation_config.set_splitting_ops_for_v1()
+            if self.compilation_config.simple_cuda_graph:
+                self.compilation_config.full_cuda_graph = True
+                self.compilation_config.level = CompilationLevel.NO_COMPILATION
+            else:
+                self.compilation_config.level = CompilationLevel.PIECEWISE
+                self.compilation_config.set_splitting_ops_for_v1()
 
         self._set_cudagraph_sizes()
 
diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py
index 67f71643d..a49e17a09 100644
--- a/vllm/distributed/utils.py
+++ b/vllm/distributed/utils.py
@@ -20,8 +20,7 @@ from typing import Any, Optional
 import torch
 from torch.distributed import ProcessGroup, TCPStore
 from torch.distributed.distributed_c10d import (Backend, PrefixStore,
-                                                _get_default_timeout,
-                                                _unregister_process_group)
+                                                _get_default_timeout)
 from torch.distributed.rendezvous import rendezvous
 
 import vllm.envs as envs
@@ -533,4 +532,3 @@ def stateless_destroy_torch_distributed_process_group(
         from torch.distributed.distributed_c10d import _shutdown_backend
         _shutdown_backend(pg)
 
-    _unregister_process_group(pg.group_name)
diff --git a/vllm/engine/llm_engine.py b/vllm/engine/llm_engine.py
index 25fa1c305..282afe172 100644
--- a/vllm/engine/llm_engine.py
+++ b/vllm/engine/llm_engine.py
@@ -227,6 +227,8 @@ class LLMEngine:
         self.observability_config = vllm_config.observability_config or ObservabilityConfig(  # noqa
         )
 
+        vllm_config.cache_config.block_size=64
+
         logger.info(
             "Initializing a V0 LLM engine (v%s) with config: %s, "
             "use_cached_outputs=%s, ",
diff --git a/vllm/env_override.py b/vllm/env_override.py
index ef425d433..653956df4 100644
--- a/vllm/env_override.py
+++ b/vllm/env_override.py
@@ -38,4 +38,4 @@ os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1'
 # see https://github.com/vllm-project/vllm/issues/10480
 os.environ['TORCHINDUCTOR_COMPILE_THREADS'] = '1'
 # see https://github.com/vllm-project/vllm/issues/10619
-torch._inductor.config.compile_threads = 1
+import torch_musa
diff --git a/vllm/model_executor/layers/quantization/gptq.py b/vllm/model_executor/layers/quantization/gptq.py
index d3ab1be3b..5f5ee0f21 100644
--- a/vllm/model_executor/layers/quantization/gptq.py
+++ b/vllm/model_executor/layers/quantization/gptq.py
@@ -86,7 +86,7 @@ class GPTQConfig(QuantizationConfig):
 
     @classmethod
     def get_supported_act_dtypes(cls) -> list[torch.dtype]:
-        return [torch.half]
+        return [torch.half, torch.bfloat16]
 
     @classmethod
     # Need to figure it out
diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py
index 2fa1294b7..af9b78ff9 100644
--- a/vllm/model_executor/models/deepseek_v2.py
+++ b/vllm/model_executor/models/deepseek_v2.py
@@ -739,14 +739,20 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
         self.num_expert_groups = config.n_group
 
         self.moe_layers: list[FusedMoE] = []
+        example_moe = None
         for layer in self.model.layers:
+            if isinstance(layer, PPMissingLayer):
+                continue
+
             assert isinstance(layer, DeepseekV2DecoderLayer)
             if isinstance(layer.mlp, DeepseekV2MoE):
+                # Pick last one layer since the first ones may be dense layers.
+                example_moe = layer.mlp
                 self.moe_layers.append(layer.mlp.experts)
 
-        # Pick last one layer since the first ones may be dense layers.
-        example_moe = typing.cast(
-            DeepseekV2MoE, self.model.layers[config.num_hidden_layers - 1].mlp)
+        if example_moe is None:
+            raise RuntimeError("No DeepseekV2MoE layer found in model.layers.")
+
         self.num_logical_experts = example_moe.n_logical_experts
         self.num_physical_experts = example_moe.n_physical_experts
         self.num_local_physical_experts = example_moe.n_local_physical_experts
@@ -882,15 +888,14 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
                     # available replicas.
                     weight_loader = typing.cast(Callable[..., bool],
                                                 param.weight_loader)
-                    success = weight_loader(param,
-                                            loaded_weight,
-                                            name_mapped,
-                                            shard_id=shard_id,
-                                            expert_id=expert_id,
-                                            return_success=True)
-                    if success:
-                        name = name_mapped
-                        break
+                    weight_loader(
+                        param,
+                        loaded_weight,
+                        name_mapped,
+                        shard_id=shard_id,
+                        expert_id=expert_id)
+                    name = name_mapped
+                    break
                 else:
                     if is_expert_weight:
                         # We've checked that this is an expert weight
diff --git a/vllm/utils/__init__.py b/vllm/utils/__init__.py
index 9550b056f..25a88b591 100644
--- a/vllm/utils/__init__.py
+++ b/vllm/utils/__init__.py
@@ -2256,6 +2256,7 @@ def direct_register_custom_op(
         dispatch_key: str = "CUDA",
         tags: tuple[torch.Tag, ...] = (),
 ):
+    return
     """
     `torch.library.custom_op` can have significant overhead because it
     needs to consider complicated dispatching logic. This function
diff --git a/vllm/worker/model_runner.py b/vllm/worker/model_runner.py
index 82db6617b..9b2a5d071 100644
--- a/vllm/worker/model_runner.py
+++ b/vllm/worker/model_runner.py
@@ -77,6 +77,23 @@ TModelInputForGPU = TypeVar('TModelInputForGPU', bound="ModelInputForGPU")
 torch._dynamo.config.cache_size_limit = 128
 torch._dynamo.config.accumulated_cache_size_limit = 128
 
+# --- FLAGSCALE MODIFICATION BEG ---
+# Know more about FlagGems: https://github.com/FlagOpen/FlagGems
+import os
+if os.getenv("USE_FLAGGEMS", "false").lower() in ("1", "true", "yes"):
+   try:
+       print("Try to using FLAGGEMS...")
+       import flag_gems
+       flag_gems.enable(record=True, path="/tmp/gems_oplist.log.txt")
+       flag_gems.apply_gems_patches_to_vllm(verbose=True)         # [Experimental] Optional, not required now. Will be officially released later.
+       logger.info("Successfully enabled flag_gems as default ops implementation.")
+   except ImportError as e:
+       # Throw an exception directly if failure occurs
+       raise ImportError("Failed to import 'flag_gems'. Please install flag_gems or set USE_FLAGGEMS=false to disable it.") from e
+   except Exception as e:
+       # Throw an exception directly if failure occurs
+       raise RuntimeError(f"Failed to enable 'flag_gems': {e}. Please check your flag_gems installation or set USE_FLAGGEMS=false to disable it.") from e
+# --- FLAGSCALE MODIFICATION END ---
 
 @dataclass(frozen=True)
 class ModelInputForGPU(ModelRunnerInputBase):
diff --git a/vllm_musa/csrc/activation_kernels.mu b/vllm_musa/csrc/activation_kernels.mu
new file mode 100644
index 000000000..c3ef94a87
--- /dev/null
+++ b/vllm_musa/csrc/activation_kernels.mu
@@ -0,0 +1,221 @@
+#include <torch/all.h>
+#include "musa.h"
+
+#include <cmath>
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+
+namespace vllm_musa {
+
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
+          bool act_first>
+__device__ __forceinline__ scalar_t compute(const scalar_t& x,
+                                            const scalar_t& y) {
+  return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
+}
+// Activation and gating kernel template.
+
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
+          bool act_first>
+__global__ void act_and_mul_kernel(
+    scalar_t* __restrict__ out,          // [..., d]
+    const scalar_t* __restrict__ input,  // [..., 2, d]
+    const int d) {
+  const int64_t token_idx = blockIdx.x;
+  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+    const scalar_t x = VLLM_LDG(&input[token_idx * 2 * d + idx]);
+    const scalar_t y = VLLM_LDG(&input[token_idx * 2 * d + d + idx]);
+    out[token_idx * d + idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
+  }
+}
+
+template <typename T>
+__device__ __forceinline__ T silu_kernel(const T& x) {
+  // x * sigmoid(x)
+  return (T)(((float)x) / (1.0f + expf((float)-x)));
+}
+
+template <typename T>
+__device__ __forceinline__ T gelu_kernel(const T& x) {
+  // Equivalent to PyTorch GELU with 'none' approximation.
+  // Refer to:
+  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
+  const float f = (float)x;
+  constexpr float ALPHA = M_SQRT1_2;
+  return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA)));
+}
+
+template <typename T>
+__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
+  // Equivalent to PyTorch GELU with 'tanh' approximation.
+  // Refer to:
+  // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
+  const float f = (float)x;
+  constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
+  constexpr float KAPPA = 0.044715;
+  float x_cube = f * f * f;
+  float inner = BETA * (f + KAPPA * x_cube);
+  return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
+}
+
+}  // namespace vllm_musa
+
+// Launch activation and gating kernel.
+// Use ACT_FIRST (bool) indicating whether to apply the activation function
+// first.
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST)                 \
+  int d = input.size(-1) / 2;                                            \
+  int64_t num_tokens = input.numel() / input.size(-1);                   \
+  dim3 grid(num_tokens);                                                 \
+  dim3 block(std::min(d, 1024));                                         \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));      \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();          \
+  VLLM_DISPATCH_FLOATING_TYPES(                                          \
+      input.scalar_type(), "act_and_mul_kernel", [&] {                   \
+        vllm_musa::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST>  \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),       \
+                                         input.data_ptr<scalar_t>(), d); \
+      });
+
+void silu_and_mul(torch::Tensor& out,    // [..., d]
+                  torch::Tensor& input)  // [..., 2 * d]
+{
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm_musa::silu_kernel, true);
+}
+
+void mul_and_silu(torch::Tensor& out,    // [..., d]
+                  torch::Tensor& input)  // [..., 2 * d]
+{
+  // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
+  // applies the silu to the latter half of the input.
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm_musa::silu_kernel, false);
+}
+
+void gelu_and_mul(torch::Tensor& out,    // [..., d]
+                  torch::Tensor& input)  // [..., 2 * d]
+{
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm_musa::gelu_kernel, true);
+}
+
+void gelu_tanh_and_mul(torch::Tensor& out,    // [..., d]
+                       torch::Tensor& input)  // [..., 2 * d]
+{
+  LAUNCH_ACTIVATION_GATE_KERNEL(vllm_musa::gelu_tanh_kernel, true);
+}
+
+namespace vllm_musa {
+
+template <typename T>
+__device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) {
+  const float f = (float)x;
+  return (T)(f > threshold ? f : 0.0f);
+}
+
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float)>
+__global__ void act_and_mul_kernel_with_param(
+    scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
+    const float param) {
+  const int64_t token_idx = blockIdx.x;
+  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+    const scalar_t x = VLLM_LDG(&input[token_idx * 2 * d + idx]);
+    const scalar_t y = VLLM_LDG(&input[token_idx * 2 * d + d + idx]);
+    out[token_idx * d + idx] = ACT_FN(x, param) * y;
+  }
+}
+
+}  // namespace vllm_musa
+
+#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PARAM)         \
+  int d = input.size(-1) / 2;                                           \
+  int64_t num_tokens = input.numel() / input.size(-1);                  \
+  dim3 grid(num_tokens);                                                \
+  dim3 block(std::min(d, 1024));                                        \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));     \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();         \
+  VLLM_DISPATCH_FLOATING_TYPES(                                         \
+      input.scalar_type(), "act_and_mul_kernel_with_param", [&] {       \
+        vllm_musa::act_and_mul_kernel_with_param<scalar_t, KERNEL<scalar_t>> \
+            <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),      \
+                                         input.data_ptr<scalar_t>(), d, \
+                                         PARAM);                        \
+      });
+
+void fatrelu_and_mul(torch::Tensor& out,    // [..., d],
+                     torch::Tensor& input,  // [..., 2 * d]
+                     double threshold) {
+  LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(vllm_musa::fatrelu_kernel, threshold);
+}
+namespace vllm_musa {
+
+// Element-wise activation kernel template.
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&)>
+__global__ void activation_kernel(
+    scalar_t* __restrict__ out,          // [..., d]
+    const scalar_t* __restrict__ input,  // [..., d]
+    const int d) {
+  const int64_t token_idx = blockIdx.x;
+  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+    const scalar_t x = VLLM_LDG(&input[token_idx * d + idx]);
+    out[token_idx * d + idx] = ACT_FN(x);
+  }
+}
+
+}  // namespace vllm_musa
+
+// Launch element-wise activation kernel.
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL)                                       \
+  int d = input.size(-1);                                                      \
+  int64_t num_tokens = input.numel() / d;                                      \
+  dim3 grid(num_tokens);                                                       \
+  dim3 block(std::min(d, 1024));                                               \
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));            \
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();                \
+  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \
+    vllm_musa::activation_kernel<scalar_t, KERNEL<scalar_t>>                        \
+        <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),                 \
+                                     input.data_ptr<scalar_t>(), d);           \
+  });
+
+namespace vllm_musa {
+
+template <typename T>
+__device__ __forceinline__ T gelu_new_kernel(const T& x) {
+  const float x3 = (float)(x * x * x);
+  const T t = (T)tanhf((T)(0.79788456f * (float)(x + (T)(0.044715f * x3))));
+  return ((T)0.5) * x * (((T)1.0) + t);
+}
+
+template <typename T>
+__device__ __forceinline__ T gelu_fast_kernel(const T& x) {
+  const float f = (float)x;
+  const T t =
+      (T)tanhf(((T)(f * 0.79788456f)) * (((T)1.0) + (T)(0.044715f * f) * x));
+  return ((T)0.5) * x * (((T)1.0) + t);
+}
+
+template <typename T>
+__device__ __forceinline__ T gelu_quick_kernel(const T& x) {
+  // x * sigmoid(1.702 * x)
+  return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x)));
+}
+
+}  // namespace vllm_musa
+
+void gelu_new(torch::Tensor& out,    // [..., d]
+              torch::Tensor& input)  // [..., d]
+{
+  LAUNCH_ACTIVATION_KERNEL(vllm_musa::gelu_new_kernel);
+}
+
+void gelu_fast(torch::Tensor& out,    // [..., d]
+               torch::Tensor& input)  // [..., d]
+{
+  LAUNCH_ACTIVATION_KERNEL(vllm_musa::gelu_fast_kernel);
+}
+
+void gelu_quick(torch::Tensor& out,    // [..., d]
+                torch::Tensor& input)  // [..., d]
+{
+  LAUNCH_ACTIVATION_KERNEL(vllm_musa::gelu_quick_kernel);
+}
diff --git a/vllm_musa/csrc/attention/attention_dtypes.h b/vllm_musa/csrc/attention/attention_dtypes.h
new file mode 100644
index 000000000..02cb8f15f
--- /dev/null
+++ b/vllm_musa/csrc/attention/attention_dtypes.h
@@ -0,0 +1,7 @@
+#pragma once
+
+#include "attention_generic.muh"
+#include "dtype_float16.muh"
+#include "dtype_float32.muh"
+#include "dtype_bfloat16.muh"
+#include "dtype_fp8.muh"
diff --git a/vllm_musa/csrc/attention/attention_generic.muh b/vllm_musa/csrc/attention/attention_generic.muh
new file mode 100644
index 000000000..5d34a6681
--- /dev/null
+++ b/vllm_musa/csrc/attention/attention_generic.muh
@@ -0,0 +1,65 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <stdint.h>
+
+namespace vllm_musa {
+
+// A vector type to store Q, K, V elements.
+template <typename T, int VEC_SIZE>
+struct Vec {};
+
+// A vector type to store FP32 accumulators.
+template <typename T>
+struct FloatVec {};
+
+// Template vector operations.
+template <typename Acc, typename A, typename B>
+inline __device__ Acc mul(A a, B b);
+
+template <typename T>
+inline __device__ float sum(T v);
+
+template <typename T>
+inline __device__ float dot(T a, T b) {
+  return sum(mul<T, T, T>(a, b));
+}
+
+template <typename A, typename T>
+inline __device__ float dot(T a, T b) {
+  return sum(mul<A, T, T>(a, b));
+}
+
+template <typename T>
+inline __device__ void zero(T& dst) {
+  constexpr int WORDS = sizeof(T) / 4;
+  union {
+    T raw;
+    uint32_t words[WORDS];
+  } tmp;
+
+#pragma unroll
+  for (int ii = 0; ii < WORDS; ++ii) {
+    tmp.words[ii] = 0u;
+  }
+  dst = tmp.raw;
+}
+
+}  // namespace vllm_musa
diff --git a/vllm_musa/csrc/attention/attention_kernels.muh b/vllm_musa/csrc/attention/attention_kernels.muh
new file mode 100644
index 000000000..d32f81e4a
--- /dev/null
+++ b/vllm_musa/csrc/attention/attention_kernels.muh
@@ -0,0 +1,647 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <torch/all.h>
+#include <algorithm>
+
+#include "../musa.h"
+#include "attention_dtypes.h"
+#include "attention_utils.muh"
+
+#define WARP_SIZE 32
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
+
+namespace vllm_musa {
+
+// Utility function for attention softmax.
+template <int NUM_WARPS>
+inline __device__ float block_sum(float* red_smem, float sum) {
+  // Decompose the thread index into warp / lane.
+  int warp = threadIdx.x / WARP_SIZE;
+  int lane = threadIdx.x % WARP_SIZE;
+
+  // Compute the sum per warp.
+#pragma unroll
+  for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
+    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
+  }
+
+  // Warp leaders store the data to shared memory.
+  if (lane == 0) {
+    red_smem[warp] = sum;
+  }
+
+  // Make sure the data is in shared memory.
+  __syncthreads();
+
+  // The warps compute the final sums.
+  if (lane < NUM_WARPS) {
+    sum = red_smem[lane];
+  }
+
+  // Parallel reduction inside the warp.
+#pragma unroll
+  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
+    sum += VLLM_SHFL_XOR_SYNC(sum, mask);
+  }
+
+  // Broadcast to other threads.
+  return VLLM_SHFL_SYNC(sum, 0);
+}
+
+// TODO(woosuk): Merge the last two dimensions of the grid.
+// Grid: (num_heads, num_seqs, max_num_partitions).
+template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
+          int NUM_THREADS, vllm_musa::Fp8KVCacheDataType KV_DTYPE,
+          bool IS_BLOCK_SPARSE,
+          int PARTITION_SIZE = 0>  // Zero means no partitioning.
+__device__ void paged_attention_kernel(
+    float* __restrict__ exp_sums,  // [num_seqs, num_heads, max_num_partitions]
+    float* __restrict__ max_logits,  // [num_seqs, num_heads,
+                                     // max_num_partitions]
+    scalar_t* __restrict__ out,  // [num_seqs, num_heads, max_num_partitions,
+                                 // head_size]
+    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
+    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
+                                          // head_size/x, block_size, x]
+    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
+                                          // head_size, block_size]
+    const int num_kv_heads,               // [num_heads]
+    const float scale,
+    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    const int* __restrict__ seq_lens,      // [num_seqs]
+    const int max_num_blocks_per_seq,
+    const float* __restrict__ alibi_slopes,  // [num_heads]
+    const int q_stride, const int kv_block_stride, const int kv_head_stride,
+    const float* k_scale, const float* v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
+  const int seq_idx = blockIdx.y;
+  const int partition_idx = blockIdx.z;
+  const int max_num_partitions = gridDim.z;
+  constexpr bool USE_PARTITIONING = PARTITION_SIZE > 0;
+  const int seq_len = seq_lens[seq_idx];
+  if (USE_PARTITIONING && partition_idx * PARTITION_SIZE >= seq_len) {
+    // No work to do. Terminate the thread block.
+    return;
+  }
+
+  const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
+  const int num_blocks_per_partition =
+      USE_PARTITIONING ? PARTITION_SIZE / BLOCK_SIZE : num_seq_blocks;
+
+  // [start_block_idx, end_block_idx) is the range of blocks to process.
+  const int start_block_idx =
+      USE_PARTITIONING ? partition_idx * num_blocks_per_partition : 0;
+  const int end_block_idx =
+      MIN(start_block_idx + num_blocks_per_partition, num_seq_blocks);
+  const int num_blocks = end_block_idx - start_block_idx;
+
+  // [start_token_idx, end_token_idx) is the range of tokens to process.
+  const int start_token_idx = start_block_idx * BLOCK_SIZE;
+  const int end_token_idx =
+      MIN(start_token_idx + num_blocks * BLOCK_SIZE, seq_len);
+  const int num_tokens = end_token_idx - start_token_idx;
+
+  constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
+  constexpr int NUM_THREAD_GROUPS =
+      NUM_THREADS / THREAD_GROUP_SIZE;  // Note: This assumes THREAD_GROUP_SIZE
+                                        // divides NUM_THREADS
+  assert(NUM_THREADS % THREAD_GROUP_SIZE == 0);
+  constexpr int NUM_TOKENS_PER_THREAD_GROUP =
+      DIVIDE_ROUND_UP(BLOCK_SIZE, WARP_SIZE);
+  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
+  const int thread_idx = threadIdx.x;
+  const int warp_idx = thread_idx / WARP_SIZE;
+  const int lane = thread_idx % WARP_SIZE;
+
+  const int head_idx = blockIdx.x;
+  const int num_heads = gridDim.x;
+  const int num_queries_per_kv = num_heads / num_kv_heads;
+  const int kv_head_idx = head_idx / num_queries_per_kv;
+  const float alibi_slope =
+      alibi_slopes == nullptr ? 0.f : alibi_slopes[head_idx];
+
+  // A vector type to store a part of a key or a query.
+  // The vector size is configured in such a way that the threads in a thread
+  // group fetch or compute 16 bytes at a time. For example, if the size of a
+  // thread group is 4 and the data type is half, then the vector size is 16 /
+  // (4 * sizeof(half)) == 2.
+  constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1);
+  using K_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
+  using Q_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
+  using Quant_vec = typename Vec<cache_t, VEC_SIZE>::Type;
+
+  constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE;
+  constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE;
+
+  const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE;
+  const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE;
+
+  // Load the query to registers.
+  // Each thread in a thread group has a different part of the query.
+  // For example, if the the thread group size is 4, then the first thread in
+  // the group has 0, 4, 8, ... th vectors of the query, and the second thread
+  // has 1, 5, 9, ... th vectors of the query, and so on. NOTE(woosuk): Because
+  // q is split from a qkv tensor, it may not be contiguous.
+  const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
+  __shared__ Q_vec q_vecs[THREAD_GROUP_SIZE][NUM_VECS_PER_THREAD];
+#pragma unroll
+  for (int i = thread_group_idx; i < NUM_VECS_PER_THREAD;
+       i += NUM_THREAD_GROUPS) {
+    const int vec_idx = thread_group_offset + i * THREAD_GROUP_SIZE;
+    q_vecs[thread_group_offset][i] =
+        *reinterpret_cast<const Q_vec*>(q_ptr + vec_idx * VEC_SIZE);
+  }
+  __syncthreads();  // TODO(naed90): possible speedup if this is replaced with a
+                    // memory wall right before we use q_vecs
+
+  // Memory planning.
+  extern __shared__ char shared_mem[];
+  // NOTE(woosuk): We use FP32 for the softmax logits for better accuracy.
+  float* logits = reinterpret_cast<float*>(shared_mem);
+  // Workspace for reduction.
+  __shared__ float red_smem[2 * NUM_WARPS];
+
+  // x == THREAD_GROUP_SIZE * VEC_SIZE
+  // Each thread group fetches x elements from the key at a time.
+  constexpr int x = 16 / sizeof(cache_t);
+  float qk_max = -FLT_MAX;
+
+  // Iterate over the key blocks.
+  // Each warp fetches a block of keys for each iteration.
+  // Each thread group in a warp fetches a key from the block, and computes
+  // dot product with the query.
+  const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;
+
+  // blocksparse specific vars
+  int bs_block_offset;
+  int q_bs_block_id;
+  if constexpr (IS_BLOCK_SPARSE) {
+    // const int num_blocksparse_blocks = DIVIDE_ROUND_UP(seq_len,
+    // blocksparse_block_size);
+    q_bs_block_id = (seq_len - 1) / blocksparse_block_size;
+    if (blocksparse_head_sliding_step >= 0)
+      // sliding on q heads
+      bs_block_offset =
+          (tp_rank * num_heads + head_idx) * blocksparse_head_sliding_step + 1;
+    else
+      // sliding on kv heads
+      bs_block_offset = (tp_rank * num_kv_heads + kv_head_idx) *
+                            (-blocksparse_head_sliding_step) +
+                        1;
+  }
+
+  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
+       block_idx += NUM_WARPS) {
+    // NOTE(woosuk): The block number is stored in int32. However, we cast it to
+    // int64 because int32 can lead to overflow when this variable is multiplied
+    // by large numbers (e.g., kv_block_stride).
+    // For blocksparse attention: skip computation on blocks that are not
+    // attended
+    if constexpr (IS_BLOCK_SPARSE) {
+      const int k_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
+      const bool is_remote =
+          ((k_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0);
+      const bool is_local =
+          (k_bs_block_id > q_bs_block_id - blocksparse_local_blocks);
+      if (!is_remote && !is_local) {
+        for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
+          const int physical_block_offset =
+              (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
+          const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
+
+          if (thread_group_offset == 0) {
+            // NOTE(linxihui): assign very large number to skipped tokens to
+            // avoid contribution to the sumexp softmax normalizer. This will
+            // not be used at computing sum(softmax*v) as the blocks will be
+            // skipped.
+            logits[token_idx - start_token_idx] = -FLT_MAX;
+          }
+        }
+        continue;
+      }
+    }
+    const int64_t physical_block_number =
+        static_cast<int64_t>(block_table[block_idx]);
+
+    // Load a key to registers.
+    // Each thread in a thread group has a different part of the key.
+    // For example, if the the thread group size is 4, then the first thread in
+    // the group has 0, 4, 8, ... th vectors of the key, and the second thread
+    // has 1, 5, 9, ... th vectors of the key, and so on.
+    for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
+      const int physical_block_offset =
+          (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
+      const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
+      K_vec k_vecs[NUM_VECS_PER_THREAD];
+
+#pragma unroll
+      for (int j = 0; j < NUM_VECS_PER_THREAD; j++) {
+        const cache_t* k_ptr =
+            k_cache + physical_block_number * kv_block_stride +
+            kv_head_idx * kv_head_stride + physical_block_offset * x;
+        const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE;
+        const int offset1 = (vec_idx * VEC_SIZE) / x;
+        const int offset2 = (vec_idx * VEC_SIZE) % x;
+
+        k_vecs[j] = *reinterpret_cast<const K_vec*>(
+            k_ptr + offset1 * BLOCK_SIZE * x + offset2);
+      }
+
+      // Compute dot product.
+      // This includes a reduction across the threads in the same thread group.
+      float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(
+                             q_vecs[thread_group_offset], k_vecs);
+      // Add the ALiBi bias if slopes are given.
+      qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0;
+
+      if (thread_group_offset == 0) {
+        // Store the partial reductions to shared memory.
+        // NOTE(woosuk): It is required to zero out the masked logits.
+        const bool mask = token_idx >= seq_len;
+        logits[token_idx - start_token_idx] = mask ? 0.f : qk;
+        // Update the max value.
+        qk_max = mask ? qk_max : fmaxf(qk_max, qk);
+      }
+    }
+  }
+
+  // Perform reduction across the threads in the same warp to get the
+  // max qk value for each "warp" (not across the thread block yet).
+  // The 0-th thread of each thread group already has its max qk value.
+#pragma unroll
+  for (int mask = WARP_SIZE / 2; mask >= THREAD_GROUP_SIZE; mask /= 2) {
+    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
+  }
+  if (lane == 0) {
+    red_smem[warp_idx] = qk_max;
+  }
+  __syncthreads();
+
+  // TODO(woosuk): Refactor this part.
+  // Get the max qk value for the sequence.
+  qk_max = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
+#pragma unroll
+  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
+    qk_max = fmaxf(qk_max, VLLM_SHFL_XOR_SYNC(qk_max, mask));
+  }
+  // Broadcast the max qk value to all threads.
+  qk_max = VLLM_SHFL_SYNC(qk_max, 0);
+
+  // Get the sum of the exp values.
+  float exp_sum = 0.f;
+  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
+    float val = __expf(logits[i] - qk_max);
+    logits[i] = val;
+    exp_sum += val;
+  }
+  exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], exp_sum);
+
+  // Compute softmax.
+  const float inv_sum = __fdividef(1.f, exp_sum + 1e-6f);
+  for (int i = thread_idx; i < num_tokens; i += NUM_THREADS) {
+    logits[i] *= inv_sum;
+  }
+  __syncthreads();
+
+  // If partitioning is enabled, store the max logit and exp_sum.
+  if (USE_PARTITIONING && thread_idx == 0) {
+    float* max_logits_ptr = max_logits +
+                            seq_idx * num_heads * max_num_partitions +
+                            head_idx * max_num_partitions + partition_idx;
+    *max_logits_ptr = qk_max;
+    float* exp_sums_ptr = exp_sums + seq_idx * num_heads * max_num_partitions +
+                          head_idx * max_num_partitions + partition_idx;
+    *exp_sums_ptr = exp_sum;
+  }
+
+  // Each thread will fetch 16 bytes from the value cache at a time.
+  constexpr int V_VEC_SIZE = MIN(16 / sizeof(scalar_t), BLOCK_SIZE);
+  using V_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
+  using L_vec = typename Vec<scalar_t, V_VEC_SIZE>::Type;
+  using V_quant_vec = typename Vec<cache_t, V_VEC_SIZE>::Type;
+  using Float_L_vec = typename FloatVec<L_vec>::Type;
+
+  constexpr int NUM_V_VECS_PER_ROW = BLOCK_SIZE / V_VEC_SIZE;
+  constexpr int NUM_ROWS_PER_ITER = WARP_SIZE / NUM_V_VECS_PER_ROW;
+  constexpr int NUM_ROWS_PER_THREAD =
+      DIVIDE_ROUND_UP(HEAD_SIZE, NUM_ROWS_PER_ITER);
+
+  // NOTE(woosuk): We use FP32 for the accumulator for better accuracy.
+  float accs[NUM_ROWS_PER_THREAD];
+#pragma unroll
+  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+    accs[i] = 0.f;
+  }
+
+  scalar_t zero_value;
+  zero(zero_value);
+  for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx;
+       block_idx += NUM_WARPS) {
+    // NOTE(woosuk): The block number is stored in int32. However, we cast it to
+    // int64 because int32 can lead to overflow when this variable is multiplied
+    // by large numbers (e.g., kv_block_stride).
+    // For blocksparse attention: skip computation on blocks that are not
+    // attended
+    if constexpr (IS_BLOCK_SPARSE) {
+      int v_bs_block_id = block_idx * BLOCK_SIZE / blocksparse_block_size;
+      if (!((v_bs_block_id + bs_block_offset) % blocksparse_vert_stride == 0) &&
+          !((v_bs_block_id > q_bs_block_id - blocksparse_local_blocks))) {
+        continue;
+      }
+    }
+    const int64_t physical_block_number =
+        static_cast<int64_t>(block_table[block_idx]);
+    const int physical_block_offset = (lane % NUM_V_VECS_PER_ROW) * V_VEC_SIZE;
+    const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
+    L_vec logits_vec;
+    from_float(logits_vec, *reinterpret_cast<Float_L_vec*>(logits + token_idx -
+                                                           start_token_idx));
+
+    const cache_t* v_ptr = v_cache + physical_block_number * kv_block_stride +
+                           kv_head_idx * kv_head_stride;
+#pragma unroll
+    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
+      if (row_idx < HEAD_SIZE) {
+        const int offset = row_idx * BLOCK_SIZE + physical_block_offset;
+        V_vec v_vec;
+
+        v_vec = *reinterpret_cast<const V_vec*>(v_ptr + offset);
+        if (block_idx == num_seq_blocks - 1) {
+          // NOTE(woosuk): When v_vec contains the tokens that are out of the
+          // context, we should explicitly zero out the values since they may
+          // contain NaNs. See
+          // https://github.com/vllm-project/vllm/issues/641#issuecomment-1682544472
+          scalar_t* v_vec_ptr = reinterpret_cast<scalar_t*>(&v_vec);
+#pragma unroll
+          for (int j = 0; j < V_VEC_SIZE; j++) {
+            v_vec_ptr[j] = token_idx + j < seq_len ? v_vec_ptr[j] : zero_value;
+          }
+        }
+        accs[i] += dot(logits_vec, v_vec);
+      }
+    }
+  }
+
+  // Perform reduction within each warp.
+#pragma unroll
+  for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+    float acc = accs[i];
+#pragma unroll
+    for (int mask = NUM_V_VECS_PER_ROW / 2; mask >= 1; mask /= 2) {
+      acc += VLLM_SHFL_XOR_SYNC(acc, mask);
+    }
+    accs[i] = acc;
+  }
+
+  // NOTE(woosuk): A barrier is required because the shared memory space for
+  // logits is reused for the output.
+  __syncthreads();
+
+  // Perform reduction across warps.
+  float* out_smem = reinterpret_cast<float*>(shared_mem);
+#pragma unroll
+  for (int i = NUM_WARPS; i > 1; i /= 2) {
+    int mid = i / 2;
+    // Upper warps write to shared memory.
+    if (warp_idx >= mid && warp_idx < i) {
+      float* dst = &out_smem[(warp_idx - mid) * HEAD_SIZE];
+#pragma unroll
+      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+        const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
+        if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
+          dst[row_idx] = accs[i];
+        }
+      }
+    }
+    __syncthreads();
+
+    // Lower warps update the output.
+    if (warp_idx < mid) {
+      const float* src = &out_smem[warp_idx * HEAD_SIZE];
+#pragma unroll
+      for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+        const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
+        if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
+          accs[i] += src[row_idx];
+        }
+      }
+    }
+    __syncthreads();
+  }
+
+  // Write the final output.
+  if (warp_idx == 0) {
+    scalar_t* out_ptr =
+        out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
+        head_idx * max_num_partitions * HEAD_SIZE + partition_idx * HEAD_SIZE;
+#pragma unroll
+    for (int i = 0; i < NUM_ROWS_PER_THREAD; i++) {
+      const int row_idx = lane / NUM_V_VECS_PER_ROW + i * NUM_ROWS_PER_ITER;
+      if (row_idx < HEAD_SIZE && lane % NUM_V_VECS_PER_ROW == 0) {
+        from_float(*(out_ptr + row_idx), accs[i]);
+      }
+    }
+  }
+}
+
+// Grid: (num_heads, num_seqs, 1).
+template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
+          int NUM_THREADS, vllm_musa::Fp8KVCacheDataType KV_DTYPE,
+          bool IS_BLOCK_SPARSE>
+__global__ void paged_attention_v1_kernel(
+    scalar_t* __restrict__ out,           // [num_seqs, num_heads, head_size]
+    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
+    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
+                                          // head_size/x, block_size, x]
+    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
+                                          // head_size, block_size]
+    const int num_kv_heads,               // [num_heads]
+    const float scale,
+    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    const int* __restrict__ seq_lens,      // [num_seqs]
+    const int max_num_blocks_per_seq,
+    const float* __restrict__ alibi_slopes,  // [num_heads]
+    const int q_stride, const int kv_block_stride, const int kv_head_stride,
+    const float* k_scale, const float* v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
+  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
+                         KV_DTYPE, IS_BLOCK_SPARSE>(
+      /* exp_sums */ nullptr, /* max_logits */ nullptr, out, q, k_cache,
+      v_cache, num_kv_heads, scale, block_tables, seq_lens,
+      max_num_blocks_per_seq, alibi_slopes, q_stride, kv_block_stride,
+      kv_head_stride, k_scale, v_scale, tp_rank, blocksparse_local_blocks,
+      blocksparse_vert_stride, blocksparse_block_size,
+      blocksparse_head_sliding_step);
+}
+
+// Grid: (num_heads, num_seqs, max_num_partitions).
+template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
+          int NUM_THREADS, vllm_musa::Fp8KVCacheDataType KV_DTYPE,
+          bool IS_BLOCK_SPARSE,
+          int PARTITION_SIZE>
+__global__ void paged_attention_v2_kernel(
+    float* __restrict__ exp_sums,  // [num_seqs, num_heads, max_num_partitions]
+    float* __restrict__ max_logits,       // [num_seqs, num_heads,
+                                          // max_num_partitions]
+    scalar_t* __restrict__ tmp_out,       // [num_seqs, num_heads,
+                                          // max_num_partitions, head_size]
+    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
+    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads,
+                                          // head_size/x, block_size, x]
+    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads,
+                                          // head_size, block_size]
+    const int num_kv_heads,               // [num_heads]
+    const float scale,
+    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    const int* __restrict__ seq_lens,      // [num_seqs]
+    const int max_num_blocks_per_seq,
+    const float* __restrict__ alibi_slopes,  // [num_heads]
+    const int q_stride, const int kv_block_stride, const int kv_head_stride,
+    const float* k_scale, const float* v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
+  paged_attention_kernel<scalar_t, cache_t, HEAD_SIZE, BLOCK_SIZE, NUM_THREADS,
+                         KV_DTYPE, IS_BLOCK_SPARSE, PARTITION_SIZE>(
+      exp_sums, max_logits, tmp_out, q, k_cache, v_cache, num_kv_heads, scale,
+      block_tables, seq_lens, max_num_blocks_per_seq, alibi_slopes, q_stride,
+      kv_block_stride, kv_head_stride, k_scale, v_scale, tp_rank,
+      blocksparse_local_blocks, blocksparse_vert_stride, blocksparse_block_size,
+      blocksparse_head_sliding_step);
+}
+
+// Grid: (num_heads, num_seqs).
+template <typename scalar_t, int HEAD_SIZE, int NUM_THREADS,
+          int PARTITION_SIZE>
+__global__ void paged_attention_v2_reduce_kernel(
+    scalar_t* __restrict__ out,            // [num_seqs, num_heads, head_size]
+    const float* __restrict__ exp_sums,    // [num_seqs, num_heads,
+                                           // max_num_partitions]
+    const float* __restrict__ max_logits,  // [num_seqs, num_heads,
+                                           // max_num_partitions]
+    const scalar_t* __restrict__ tmp_out,  // [num_seqs, num_heads,
+                                           // max_num_partitions, head_size]
+    const int* __restrict__ seq_lens,      // [num_seqs]
+    const int max_num_partitions) {
+  const int num_heads = gridDim.x;
+  const int head_idx = blockIdx.x;
+  const int seq_idx = blockIdx.y;
+  const int seq_len = seq_lens[seq_idx];
+  const int num_partitions = DIVIDE_ROUND_UP(seq_len, PARTITION_SIZE);
+  if (num_partitions == 1) {
+    // No need to reduce. Only copy tmp_out to out.
+    scalar_t* out_ptr =
+        out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
+    const scalar_t* tmp_out_ptr =
+        tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
+        head_idx * max_num_partitions * HEAD_SIZE;
+    for (int i = threadIdx.x; i < HEAD_SIZE; i += blockDim.x) {
+      out_ptr[i] = tmp_out_ptr[i];
+    }
+    // Terminate the thread block.
+    return;
+  }
+
+  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
+  const int warp_idx = threadIdx.x / WARP_SIZE;
+  const int lane = threadIdx.x % WARP_SIZE;
+
+  // Size: 2 * num_partitions.
+  extern __shared__ char shared_mem[];
+  // Workspace for reduction.
+  __shared__ float red_smem[2 * NUM_WARPS];
+
+  // Load max logits to shared memory.
+  float* shared_max_logits = reinterpret_cast<float*>(shared_mem);
+  const float* max_logits_ptr = max_logits +
+                                seq_idx * num_heads * max_num_partitions +
+                                head_idx * max_num_partitions;
+  float max_logit = -FLT_MAX;
+  for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
+    const float l = max_logits_ptr[i];
+    shared_max_logits[i] = l;
+    max_logit = fmaxf(max_logit, l);
+  }
+  __syncthreads();
+
+  // Get the global max logit.
+  // Reduce within the warp.
+#pragma unroll
+  for (int mask = WARP_SIZE / 2; mask >= 1; mask /= 2) {
+    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
+  }
+  if (lane == 0) {
+    red_smem[warp_idx] = max_logit;
+  }
+  __syncthreads();
+  // Reduce across warps.
+  max_logit = lane < NUM_WARPS ? red_smem[lane] : -FLT_MAX;
+#pragma unroll
+  for (int mask = NUM_WARPS / 2; mask >= 1; mask /= 2) {
+    max_logit = fmaxf(max_logit, VLLM_SHFL_XOR_SYNC(max_logit, mask));
+  }
+  // Broadcast the max value to all threads.
+  max_logit = VLLM_SHFL_SYNC(max_logit, 0);
+
+  // Load rescaled exp sums to shared memory.
+  float* shared_exp_sums =
+      reinterpret_cast<float*>(shared_mem + sizeof(float) * num_partitions);
+  const float* exp_sums_ptr = exp_sums +
+                              seq_idx * num_heads * max_num_partitions +
+                              head_idx * max_num_partitions;
+  float global_exp_sum = 0.0f;
+  for (int i = threadIdx.x; i < num_partitions; i += blockDim.x) {
+    float l = shared_max_logits[i];
+    float rescaled_exp_sum = exp_sums_ptr[i] * expf(l - max_logit);
+    global_exp_sum += rescaled_exp_sum;
+    shared_exp_sums[i] = rescaled_exp_sum;
+  }
+  __syncthreads();
+  global_exp_sum = block_sum<NUM_WARPS>(&red_smem[NUM_WARPS], global_exp_sum);
+  const float inv_global_exp_sum = __fdividef(1.0f, global_exp_sum + 1e-6f);
+
+  // Aggregate tmp_out to out.
+  const scalar_t* tmp_out_ptr =
+      tmp_out + seq_idx * num_heads * max_num_partitions * HEAD_SIZE +
+      head_idx * max_num_partitions * HEAD_SIZE;
+  scalar_t* out_ptr =
+      out + seq_idx * num_heads * HEAD_SIZE + head_idx * HEAD_SIZE;
+#pragma unroll
+  for (int i = threadIdx.x; i < HEAD_SIZE; i += NUM_THREADS) {
+    float acc = 0.0f;
+    for (int j = 0; j < num_partitions; ++j) {
+      acc += to_float(tmp_out_ptr[j * HEAD_SIZE + i]) * shared_exp_sums[j] *
+             inv_global_exp_sum;
+    }
+    from_float(out_ptr[i], acc);
+  }
+}
+
+}  // namespace vllm_musa
+
+#undef WARP_SIZE
+#undef MAX
+#undef MIN
+#undef DIVIDE_ROUND_UP
diff --git a/vllm_musa/csrc/attention/attention_utils.muh b/vllm_musa/csrc/attention/attention_utils.muh
new file mode 100644
index 000000000..e4bd89b0b
--- /dev/null
+++ b/vllm_musa/csrc/attention/attention_utils.muh
@@ -0,0 +1,57 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include "../cuda_compat.h"
+#include "attention_dtypes.h"
+
+#include <float.h>
+#include <type_traits>
+
+namespace vllm_musa {
+
+// Q*K^T operation.
+template <int THREAD_GROUP_SIZE, typename Vec, int N>
+inline __device__ float qk_dot_(const Vec (&q)[N], const Vec (&k)[N]) {
+  using A_vec = typename FloatVec<Vec>::Type;
+  // Compute the parallel products for Q*K^T (treat vector lanes separately).
+  A_vec qk_vec = mul<A_vec, Vec, Vec>(q[0], k[0]);
+#pragma unroll
+  for (int ii = 1; ii < N; ++ii) {
+    qk_vec = vllm_musa::fma(q[ii], k[ii], qk_vec);
+  }
+
+  // Finalize the reduction across lanes.
+  float qk = sum(qk_vec);
+#pragma unroll
+  for (int mask = THREAD_GROUP_SIZE / 2; mask >= 1; mask /= 2) {
+    qk += VLLM_SHFL_XOR_SYNC(qk, mask);
+  }
+  return qk;
+}
+
+template <typename T, int THREAD_GROUP_SIZE>
+struct Qk_dot {
+  template <typename Vec, int N>
+  static inline __device__ float dot(const Vec (&q)[N], const Vec (&k)[N]) {
+    return qk_dot_<THREAD_GROUP_SIZE>(q, k);
+  }
+};
+
+}  // namespace vllm_musa
diff --git a/vllm_musa/csrc/attention/dtype_bfloat16.muh b/vllm_musa/csrc/attention/dtype_bfloat16.muh
new file mode 100644
index 000000000..ee76cf188
--- /dev/null
+++ b/vllm_musa/csrc/attention/dtype_bfloat16.muh
@@ -0,0 +1,449 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
+ * and
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include "attention_generic.muh"
+#include "dtype_float32.muh"
+#include "../musa.h"
+
+#include <stdint.h>
+
+namespace vllm_musa {
+
+// Define custom BF16 vector data types.
+struct bf16_4_t {
+  __nv_bfloat162 x;
+  __nv_bfloat162 y;
+};
+
+struct bf16_8_t {
+  __nv_bfloat162 x;
+  __nv_bfloat162 y;
+  __nv_bfloat162 z;
+  __nv_bfloat162 w;
+};
+
+// BF16 vector types for Q, K, V.
+template <>
+struct Vec<__nv_bfloat16, 1> {
+  using Type = __nv_bfloat16;
+};
+template <>
+struct Vec<__nv_bfloat16, 2> {
+  using Type = __nv_bfloat162;
+};
+template <>
+struct Vec<__nv_bfloat16, 4> {
+  using Type = bf16_4_t;
+};
+template <>
+struct Vec<__nv_bfloat16, 8> {
+  using Type = bf16_8_t;
+};
+
+// FP32 accumulator vector types corresponding to Vec.
+template <>
+struct FloatVec<__nv_bfloat16> {
+  using Type = float;
+};
+template <>
+struct FloatVec<__nv_bfloat162> {
+  using Type = float2;
+};
+template <>
+struct FloatVec<bf16_4_t> {
+  using Type = Float4_;
+};
+template <>
+struct FloatVec<bf16_8_t> {
+  using Type = Float8_;
+};
+
+// Utility functions for type conversions.
+inline __device__ float2 bf1622float2(const __nv_bfloat162 val) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  return __bfloat1622float2(val);
+#endif
+  __builtin_unreachable();  // Suppress missing return statement warning
+}
+
+inline __device__ __nv_bfloat162 bf162bf162(const __nv_bfloat16 val) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  return __bfloat162bfloat162(val);
+#endif
+  __builtin_unreachable();  // Suppress missing return statement warning
+}
+
+// Vector addition.
+inline __device__ __nv_bfloat16 add(__nv_bfloat16 a, __nv_bfloat16 b) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  return a + b;
+#endif
+  __builtin_unreachable();  // Suppress missing return statement warning
+}
+
+inline __device__ __nv_bfloat162 add(__nv_bfloat162 a, __nv_bfloat162 b) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  return __hadd2(a, b);
+#endif
+  __builtin_unreachable();  // Suppress missing return statement warning
+}
+
+inline __device__ bf16_4_t add(bf16_4_t a, bf16_4_t b) {
+  bf16_4_t c;
+  c.x = add(a.x, b.x);
+  c.y = add(a.y, b.y);
+  return c;
+}
+
+inline __device__ bf16_8_t add(bf16_8_t a, bf16_8_t b) {
+  bf16_8_t c;
+  c.x = add(a.x, b.x);
+  c.y = add(a.y, b.y);
+  c.z = add(a.z, b.z);
+  c.w = add(a.w, b.w);
+  return c;
+}
+
+inline __device__ float2 add(__nv_bfloat162 a, float2 fb) {
+  float2 fa = bf1622float2(a);
+  return add(fa, fb);
+}
+
+inline __device__ Float4_ add(bf16_4_t a, Float4_ fb) {
+  Float4_ fc;
+  fc.x = add(a.x, fb.x);
+  fc.y = add(a.y, fb.y);
+  return fc;
+}
+
+inline __device__ Float8_ add(bf16_8_t a, Float8_ fb) {
+  Float8_ fc;
+  fc.x = add(a.x, fb.x);
+  fc.y = add(a.y, fb.y);
+  fc.z = add(a.z, fb.z);
+  fc.w = add(a.w, fb.w);
+  return fc;
+}
+
+// Vector multiplication.
+template <>
+inline __device__ __nv_bfloat16 mul(__nv_bfloat16 a, __nv_bfloat16 b) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  return __hmul(a, b);
+#endif
+  __builtin_unreachable();  // Suppress missing return statement warning
+}
+
+template <>
+inline __device__ __nv_bfloat162 mul(__nv_bfloat162 a, __nv_bfloat162 b) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  return __hmul2(a, b);
+#endif
+  __builtin_unreachable();  // Suppress missing return statement warning
+}
+
+template <>
+inline __device__ __nv_bfloat162 mul(__nv_bfloat16 a, __nv_bfloat162 b) {
+  return mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(bf162bf162(a), b);
+}
+
+template <>
+inline __device__ bf16_4_t mul(bf16_4_t a, bf16_4_t b) {
+  bf16_4_t c;
+  c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x);
+  c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y);
+  return c;
+}
+
+template <>
+inline __device__ bf16_4_t mul(__nv_bfloat16 a, bf16_4_t b) {
+  __nv_bfloat162 s = bf162bf162(a);
+  bf16_4_t c;
+  c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x);
+  c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y);
+  return c;
+}
+
+template <>
+inline __device__ bf16_8_t mul(bf16_8_t a, bf16_8_t b) {
+  bf16_8_t c;
+  c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.x, b.x);
+  c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.y, b.y);
+  c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.z, b.z);
+  c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(a.w, b.w);
+  return c;
+}
+
+template <>
+inline __device__ bf16_8_t mul(__nv_bfloat16 a, bf16_8_t b) {
+  __nv_bfloat162 s = bf162bf162(a);
+  bf16_8_t c;
+  c.x = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.x);
+  c.y = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.y);
+  c.z = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.z);
+  c.w = mul<__nv_bfloat162, __nv_bfloat162, __nv_bfloat162>(s, b.w);
+  return c;
+}
+
+template <>
+inline __device__ float mul(__nv_bfloat16 a, __nv_bfloat16 b) {
+  float fa = __bfloat162float(a);
+  float fb = __bfloat162float(b);
+  return fa * fb;
+}
+
+template <>
+inline __device__ float2 mul(__nv_bfloat162 a, __nv_bfloat162 b) {
+  float2 fa = bf1622float2(a);
+  float2 fb = bf1622float2(b);
+  return mul<float2, float2, float2>(fa, fb);
+}
+
+template <>
+inline __device__ float2 mul(__nv_bfloat16 a, __nv_bfloat162 b) {
+  return mul<float2, __nv_bfloat162, __nv_bfloat162>(bf162bf162(a), b);
+}
+
+template <>
+inline __device__ Float4_ mul(bf16_4_t a, bf16_4_t b) {
+  Float4_ fc;
+  fc.x = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.x, b.x);
+  fc.y = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.y, b.y);
+  return fc;
+}
+
+template <>
+inline __device__ Float4_ mul(__nv_bfloat16 a, bf16_4_t b) {
+  __nv_bfloat162 s = bf162bf162(a);
+  Float4_ fc;
+  fc.x = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.x);
+  fc.y = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.y);
+  return fc;
+}
+
+template <>
+inline __device__ Float8_ mul(bf16_8_t a, bf16_8_t b) {
+  Float8_ fc;
+  fc.x = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.x, b.x);
+  fc.y = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.y, b.y);
+  fc.z = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.z, b.z);
+  fc.w = mul<float2, __nv_bfloat162, __nv_bfloat162>(a.w, b.w);
+  return fc;
+}
+
+template <>
+inline __device__ Float8_ mul(__nv_bfloat16 a, bf16_8_t b) {
+  __nv_bfloat162 s = bf162bf162(a);
+  Float8_ fc;
+  fc.x = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.x);
+  fc.y = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.y);
+  fc.z = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.z);
+  fc.w = mul<float2, __nv_bfloat162, __nv_bfloat162>(s, b.w);
+  return fc;
+}
+
+// Vector fused multiply-add.
+inline __device__ __nv_bfloat162 fma(__nv_bfloat162 a, __nv_bfloat162 b,
+                                     __nv_bfloat162 c) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  return __hfma2(a, b, c);
+#endif
+  __builtin_unreachable();  // Suppress missing return statement warning
+}
+
+inline __device__ __nv_bfloat162 fma(__nv_bfloat16 a, __nv_bfloat162 b,
+                                     __nv_bfloat162 c) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  return __hfma2(bf162bf162(a), b, c);
+#endif
+  __builtin_unreachable();  // Suppress missing return statement warning
+}
+
+inline __device__ bf16_4_t fma(bf16_4_t a, bf16_4_t b, bf16_4_t c) {
+  bf16_4_t d;
+  d.x = fma(a.x, b.x, c.x);
+  d.y = fma(a.y, b.y, c.y);
+  return d;
+}
+
+inline __device__ bf16_4_t fma(__nv_bfloat16 a, bf16_4_t b, bf16_4_t c) {
+  __nv_bfloat162 s = bf162bf162(a);
+  bf16_4_t d;
+  d.x = fma(s, b.x, c.x);
+  d.y = fma(s, b.y, c.y);
+  return d;
+}
+
+inline __device__ bf16_8_t fma(bf16_8_t a, bf16_8_t b, bf16_8_t c) {
+  bf16_8_t d;
+  d.x = fma(a.x, b.x, c.x);
+  d.y = fma(a.y, b.y, c.y);
+  d.z = fma(a.z, b.z, c.z);
+  d.w = fma(a.w, b.w, c.w);
+  return d;
+}
+
+inline __device__ bf16_8_t fma(__nv_bfloat16 a, bf16_8_t b, bf16_8_t c) {
+  __nv_bfloat162 s = bf162bf162(a);
+  bf16_8_t d;
+  d.x = fma(s, b.x, c.x);
+  d.y = fma(s, b.y, c.y);
+  d.z = fma(s, b.z, c.z);
+  d.w = fma(s, b.w, c.w);
+  return d;
+}
+
+inline __device__ float fma(__nv_bfloat16 a, __nv_bfloat16 b, float fc) {
+  return __bfloat162float(a) * __bfloat162float(b) + fc;
+}
+
+inline __device__ float2 fma(__nv_bfloat162 a, __nv_bfloat162 b, float2 fc) {
+  float2 fa = bf1622float2(a);
+  float2 fb = bf1622float2(b);
+  return fma(fa, fb, fc);
+}
+
+inline __device__ float2 fma(__nv_bfloat16 a, __nv_bfloat162 b, float2 fc) {
+  return fma(bf162bf162(a), b, fc);
+}
+
+inline __device__ Float4_ fma(bf16_4_t a, bf16_4_t b, Float4_ fc) {
+  Float4_ fd;
+  fd.x = fma(a.x, b.x, fc.x);
+  fd.y = fma(a.y, b.y, fc.y);
+  return fd;
+}
+
+inline __device__ Float4_ fma(__nv_bfloat16 a, bf16_4_t b, Float4_ fc) {
+  __nv_bfloat162 s = bf162bf162(a);
+  Float4_ fd;
+  fd.x = fma(s, b.x, fc.x);
+  fd.y = fma(s, b.y, fc.y);
+  return fd;
+}
+
+inline __device__ Float8_ fma(bf16_8_t a, bf16_8_t b, Float8_ fc) {
+  Float8_ fd;
+  fd.x = fma(a.x, b.x, fc.x);
+  fd.y = fma(a.y, b.y, fc.y);
+  fd.z = fma(a.z, b.z, fc.z);
+  fd.w = fma(a.w, b.w, fc.w);
+  return fd;
+}
+
+inline __device__ Float8_ fma(__nv_bfloat16 a, bf16_8_t b, Float8_ fc) {
+  __nv_bfloat162 s = bf162bf162(a);
+  Float8_ fd;
+  fd.x = fma(s, b.x, fc.x);
+  fd.y = fma(s, b.y, fc.y);
+  fd.z = fma(s, b.z, fc.z);
+  fd.w = fma(s, b.w, fc.w);
+  return fd;
+}
+
+// Vector sum.
+template <>
+inline __device__ float sum(__nv_bfloat16 v) {
+  return __bfloat162float(v);
+}
+
+template <>
+inline __device__ float sum(__nv_bfloat162 v) {
+  float2 vf = bf1622float2(v);
+  return vf.x + vf.y;
+}
+
+template <>
+inline __device__ float sum(bf16_4_t v) {
+  return sum(v.x) + sum(v.y);
+}
+
+template <>
+inline __device__ float sum(bf16_8_t v) {
+  return sum(v.x) + sum(v.y) + sum(v.z) + sum(v.w);
+}
+
+// From float32 to bfloat16.
+inline __device__ void from_float(__nv_bfloat16& dst, float src) {
+  dst = __float2bfloat16(src);
+}
+
+inline __device__ void from_float(__nv_bfloat162& dst, float2 src) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  dst = __float22bfloat162_rn(src);
+#endif
+}
+
+inline __device__ void from_float(bf16_4_t& dst, Float4_ src) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  dst.x = __float22bfloat162_rn(src.x);
+  dst.y = __float22bfloat162_rn(src.y);
+#endif
+}
+
+inline __device__ void from_float(bf16_8_t& dst, Float8_ src) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  dst.x = __float22bfloat162_rn(src.x);
+  dst.y = __float22bfloat162_rn(src.y);
+  dst.z = __float22bfloat162_rn(src.z);
+  dst.w = __float22bfloat162_rn(src.w);
+#endif
+}
+
+// From bfloat16 to float32.
+inline __device__ float to_float(__nv_bfloat16 u) {
+  return __bfloat162float(u);
+}
+
+// Zero-out a variable.
+inline __device__ void zero(__nv_bfloat16& dst) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+  assert(false);
+#else
+  // Same as CUDART_ZERO_BF16 introduced in CUDA 12.2.
+  dst = __ushort_as_bfloat16((unsigned short)0x0000U);
+#endif
+}
+
+}  // namespace vllm_musa
diff --git a/vllm_musa/csrc/attention/dtype_float16.muh b/vllm_musa/csrc/attention/dtype_float16.muh
new file mode 100644
index 000000000..259b6d35f
--- /dev/null
+++ b/vllm_musa/csrc/attention/dtype_float16.muh
@@ -0,0 +1,502 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
+ * and
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include "attention_generic.muh"
+#include "dtype_float32.muh"
+
+#include "../musa.h"
+
+#include <stdint.h>
+
+namespace vllm_musa {
+
+// FP16 vector types for Q, K, V.
+template <>
+struct Vec<uint16_t, 1> {
+  using Type = uint16_t;
+};
+template <>
+struct Vec<uint16_t, 2> {
+  using Type = uint32_t;
+};
+template <>
+struct Vec<uint16_t, 4> {
+  using Type = uint2;
+};
+template <>
+struct Vec<uint16_t, 8> {
+  using Type = uint4;
+};
+
+// FP32 accumulator vector types corresponding to Vec.
+template <>
+struct FloatVec<uint16_t> {
+  using Type = float;
+};
+template <>
+struct FloatVec<uint32_t> {
+  using Type = float2;
+};
+template <>
+struct FloatVec<uint2> {
+  using Type = Float4_;
+};
+template <>
+struct FloatVec<uint4> {
+  using Type = Float8_;
+};
+
+// Utility functions for type conversions.
+inline __device__ uint32_t h0_h0(uint16_t a) {
+#ifndef USE_ROCM
+  uint32_t b;
+  asm volatile("mov.b32 %0, {%1, %1};" : "=r"(b) : "h"(a));
+  return b;
+#else
+  union {
+    uint32_t u32;
+    uint16_t u16[2];
+  } tmp;
+  tmp.u16[0] = a;
+  tmp.u16[1] = a;
+  return tmp.u32;
+#endif
+}
+
+inline __device__ float half_to_float(uint16_t h) {
+  float f;
+#ifndef USE_ROCM
+  asm volatile("cvt.f32.f16 %0, %1;\n" : "=f"(f) : "h"(h));
+#else
+  asm volatile("v_cvt_f32_f16 %0, %1;" : "=v"(f) : "v"(h));
+#endif
+  return f;
+}
+
+inline __device__ float2 half2_to_float2(uint32_t v) {
+#ifndef USE_ROCM
+  uint16_t lo, hi;
+  asm volatile("mov.b32 {%0, %1}, %2;\n" : "=h"(lo), "=h"(hi) : "r"(v));
+  return make_float2(half_to_float(lo), half_to_float(hi));
+#else
+  union {
+    uint32_t u32;
+    uint16_t u16[2];
+  } tmp;
+  tmp.u32 = v;
+  float2 ret;
+  ret.x = half_to_float(tmp.u16[0]);
+  ret.y = half_to_float(tmp.u16[1]);
+  return ret;
+#endif
+}
+
+inline __device__ uint16_t float_to_half(float f) {
+  union {
+    uint32_t u32;
+    uint16_t u16[2];
+  } tmp;
+#ifndef USE_ROCM
+  asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f));
+#else
+  asm volatile("v_cvt_f16_f32 %0, %1;\n" : "=v"(tmp.u32) : "v"(f));
+#endif
+  return tmp.u16[0];
+}
+
+inline __device__ uint32_t float2_to_half2(float2 f) {
+  union {
+    uint32_t u32;
+    uint16_t u16[2];
+  } tmp;
+#ifndef USE_ROCM
+  #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+  asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n"
+               : "=r"(tmp.u32)
+               : "f"(f.y), "f"(f.x));
+  #else
+  asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f.x));
+  asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[1]) : "f"(f.y));
+  #endif
+#else
+  tmp.u16[0] = float_to_half(f.x);
+  tmp.u16[1] = float_to_half(f.y);
+#endif
+  return tmp.u32;
+}
+
+// Vector addition.
+inline __device__ uint16_t add(uint16_t a, uint16_t b) {
+  uint16_t c;
+#ifndef USE_ROCM
+  asm volatile("add.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b));
+#else
+  asm volatile("v_add_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b));
+#endif
+  return c;
+}
+
+inline __device__ uint32_t add(uint32_t a, uint32_t b) {
+  uint32_t c;
+#ifndef USE_ROCM
+  asm volatile("add.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b));
+#else
+  asm volatile("v_pk_add_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b));
+#endif
+  return c;
+}
+
+inline __device__ uint2 add(uint2 a, uint2 b) {
+  uint2 c;
+  c.x = add(a.x, b.x);
+  c.y = add(a.y, b.y);
+  return c;
+}
+
+inline __device__ uint4 add(uint4 a, uint4 b) {
+  uint4 c;
+  c.x = add(a.x, b.x);
+  c.y = add(a.y, b.y);
+  c.z = add(a.z, b.z);
+  c.w = add(a.w, b.w);
+  return c;
+}
+
+inline __device__ float2 add(uint32_t a, float2 fb) {
+  float2 fa = half2_to_float2(a);
+  return add(fa, fb);
+}
+
+inline __device__ Float4_ add(uint2 a, Float4_ fb) {
+  Float4_ fc;
+  fc.x = add(a.x, fb.x);
+  fc.y = add(a.y, fb.y);
+  return fc;
+}
+
+inline __device__ Float8_ add(uint4 a, Float8_ fb) {
+  Float8_ fc;
+  fc.x = add(a.x, fb.x);
+  fc.y = add(a.y, fb.y);
+  fc.z = add(a.z, fb.z);
+  fc.w = add(a.w, fb.w);
+  return fc;
+}
+
+// Vector multiplication.
+template <>
+inline __device__ uint16_t mul(uint16_t a, uint16_t b) {
+  uint16_t c;
+#ifndef USE_ROCM
+  asm volatile("mul.f16 %0, %1, %2;\n" : "=h"(c) : "h"(a), "h"(b));
+#else
+  asm volatile("v_mul_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b));
+#endif
+  return c;
+}
+
+template <>
+inline __device__ uint32_t mul(uint32_t a, uint32_t b) {
+  uint32_t c;
+#ifndef USE_ROCM
+  asm volatile("mul.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b));
+#else
+  asm volatile("v_pk_mul_f16 %0, %1, %2;\n" : "=v"(c) : "v"(a), "v"(b));
+#endif
+  return c;
+}
+
+template <>
+inline __device__ uint32_t mul(uint16_t a, uint32_t b) {
+  return mul<uint32_t, uint32_t, uint32_t>(h0_h0(a), b);
+}
+
+template <>
+inline __device__ uint2 mul(uint2 a, uint2 b) {
+  uint2 c;
+  c.x = mul<uint32_t, uint32_t, uint32_t>(a.x, b.x);
+  c.y = mul<uint32_t, uint32_t, uint32_t>(a.y, b.y);
+  return c;
+}
+
+template <>
+inline __device__ uint2 mul(uint16_t a, uint2 b) {
+  uint32_t s = h0_h0(a);
+  uint2 c;
+  c.x = mul<uint32_t, uint32_t, uint32_t>(s, b.x);
+  c.y = mul<uint32_t, uint32_t, uint32_t>(s, b.y);
+  return c;
+}
+
+template <>
+inline __device__ uint4 mul(uint4 a, uint4 b) {
+  uint4 c;
+  c.x = mul<uint32_t, uint32_t, uint32_t>(a.x, b.x);
+  c.y = mul<uint32_t, uint32_t, uint32_t>(a.y, b.y);
+  c.z = mul<uint32_t, uint32_t, uint32_t>(a.z, b.z);
+  c.w = mul<uint32_t, uint32_t, uint32_t>(a.w, b.w);
+  return c;
+}
+
+template <>
+inline __device__ uint4 mul(uint16_t a, uint4 b) {
+  uint32_t s = h0_h0(a);
+  uint4 c;
+  c.x = mul<uint32_t, uint32_t, uint32_t>(s, b.x);
+  c.y = mul<uint32_t, uint32_t, uint32_t>(s, b.y);
+  c.z = mul<uint32_t, uint32_t, uint32_t>(s, b.z);
+  c.w = mul<uint32_t, uint32_t, uint32_t>(s, b.w);
+  return c;
+}
+
+template <>
+inline __device__ float mul(uint16_t a, uint16_t b) {
+  float fa = half_to_float(a);
+  float fb = half_to_float(b);
+  return fa * fb;
+}
+
+template <>
+inline __device__ float2 mul(uint32_t a, uint32_t b) {
+  float2 fa = half2_to_float2(a);
+  float2 fb = half2_to_float2(b);
+  return mul<float2, float2, float2>(fa, fb);
+}
+
+template <>
+inline __device__ float2 mul(uint16_t a, uint32_t b) {
+  return mul<float2, uint32_t, uint32_t>(h0_h0(a), b);
+}
+
+template <>
+inline __device__ Float4_ mul(uint2 a, uint2 b) {
+  Float4_ fc;
+  fc.x = mul<float2, uint32_t, uint32_t>(a.x, b.x);
+  fc.y = mul<float2, uint32_t, uint32_t>(a.y, b.y);
+  return fc;
+}
+
+template <>
+inline __device__ Float4_ mul(uint16_t a, uint2 b) {
+  uint32_t s = h0_h0(a);
+  Float4_ fc;
+  fc.x = mul<float2, uint32_t, uint32_t>(s, b.x);
+  fc.y = mul<float2, uint32_t, uint32_t>(s, b.y);
+  return fc;
+}
+
+template <>
+inline __device__ Float8_ mul(uint4 a, uint4 b) {
+  Float8_ fc;
+  fc.x = mul<float2, uint32_t, uint32_t>(a.x, b.x);
+  fc.y = mul<float2, uint32_t, uint32_t>(a.y, b.y);
+  fc.z = mul<float2, uint32_t, uint32_t>(a.z, b.z);
+  fc.w = mul<float2, uint32_t, uint32_t>(a.w, b.w);
+  return fc;
+}
+
+template <>
+inline __device__ Float8_ mul(uint16_t a, uint4 b) {
+  uint32_t s = h0_h0(a);
+  Float8_ fc;
+  fc.x = mul<float2, uint32_t, uint32_t>(s, b.x);
+  fc.y = mul<float2, uint32_t, uint32_t>(s, b.y);
+  fc.z = mul<float2, uint32_t, uint32_t>(s, b.z);
+  fc.w = mul<float2, uint32_t, uint32_t>(s, b.w);
+  return fc;
+}
+
+// Vector fused multiply-add.
+inline __device__ uint32_t fma(uint32_t a, uint32_t b, uint32_t c) {
+  uint32_t d;
+#ifndef USE_ROCM
+  asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n"
+               : "=r"(d)
+               : "r"(a), "r"(b), "r"(c));
+#else
+  asm volatile("v_pk_fma_f16 %0, %1, %2, %3;\n"
+               : "=v"(d)
+               : "v"(a), "v"(b), "v"(c));
+#endif
+  return d;
+}
+
+inline __device__ uint32_t fma(uint16_t a, uint32_t b, uint32_t c) {
+  return fma(h0_h0(a), b, c);
+}
+
+inline __device__ uint2 fma(uint2 a, uint2 b, uint2 c) {
+  uint2 d;
+  d.x = fma(a.x, b.x, c.x);
+  d.y = fma(a.y, b.y, c.y);
+  return d;
+}
+
+inline __device__ uint2 fma(uint16_t a, uint2 b, uint2 c) {
+  uint32_t s = h0_h0(a);
+  uint2 d;
+  d.x = fma(s, b.x, c.x);
+  d.y = fma(s, b.y, c.y);
+  return d;
+}
+
+inline __device__ uint4 fma(uint4 a, uint4 b, uint4 c) {
+  uint4 d;
+  d.x = fma(a.x, b.x, c.x);
+  d.y = fma(a.y, b.y, c.y);
+  d.z = fma(a.z, b.z, c.z);
+  d.w = fma(a.w, b.w, c.w);
+  return d;
+}
+
+inline __device__ uint4 fma(uint16_t a, uint4 b, uint4 c) {
+  uint32_t s = h0_h0(a);
+  uint4 d;
+  d.x = fma(s, b.x, c.x);
+  d.y = fma(s, b.y, c.y);
+  d.z = fma(s, b.z, c.z);
+  d.w = fma(s, b.w, c.w);
+  return d;
+}
+
+inline __device__ float fma(uint16_t a, uint16_t b, float fc) {
+  float fa = half_to_float(a);
+  float fb = half_to_float(b);
+  return fa * fb + fc;
+}
+
+inline __device__ float2 fma(uint32_t a, uint32_t b, float2 fc) {
+  float2 fa = half2_to_float2(a);
+  float2 fb = half2_to_float2(b);
+  return fma(fa, fb, fc);
+}
+
+inline __device__ float2 fma(uint16_t a, uint32_t b, float2 fc) {
+  return fma(h0_h0(a), b, fc);
+}
+
+inline __device__ Float4_ fma(uint2 a, uint2 b, Float4_ fc) {
+  Float4_ fd;
+  fd.x = fma(a.x, b.x, fc.x);
+  fd.y = fma(a.y, b.y, fc.y);
+  return fd;
+}
+
+inline __device__ Float4_ fma(uint16_t a, uint2 b, Float4_ fc) {
+  uint32_t s = h0_h0(a);
+  Float4_ fd;
+  fd.x = fma(s, b.x, fc.x);
+  fd.y = fma(s, b.y, fc.y);
+  return fd;
+}
+
+inline __device__ Float8_ fma(uint4 a, uint4 b, Float8_ fc) {
+  Float8_ fd;
+  fd.x = fma(a.x, b.x, fc.x);
+  fd.y = fma(a.y, b.y, fc.y);
+  fd.z = fma(a.z, b.z, fc.z);
+  fd.w = fma(a.w, b.w, fc.w);
+  return fd;
+}
+
+inline __device__ Float8_ fma(uint16_t a, uint4 b, Float8_ fc) {
+  uint32_t s = h0_h0(a);
+  Float8_ fd;
+  fd.x = fma(s, b.x, fc.x);
+  fd.y = fma(s, b.y, fc.y);
+  fd.z = fma(s, b.z, fc.z);
+  fd.w = fma(s, b.w, fc.w);
+  return fd;
+}
+
+// Vector sum.
+template <>
+inline __device__ float sum(uint16_t v) {
+  return half_to_float(v);
+}
+
+template <>
+inline __device__ float sum(uint32_t v) {
+  float2 tmp = half2_to_float2(v);
+  return tmp.x + tmp.y;
+}
+
+template <>
+inline __device__ float sum(uint2 v) {
+  uint32_t c = add(v.x, v.y);
+  return sum(c);
+}
+
+template <>
+inline __device__ float sum(uint4 v) {
+  uint32_t c = add(v.x, v.y);
+  c = add(c, v.z);
+  c = add(c, v.w);
+  return sum(c);
+}
+
+// From float32 to float16.
+inline __device__ void from_float(uint16_t& dst, float src) {
+  dst = float_to_half(src);
+}
+
+inline __device__ void from_float(uint32_t& dst, float2 src) {
+  dst = float2_to_half2(src);
+}
+
+inline __device__ void from_float(uint2& dst, Float4_ src) {
+  dst.x = float2_to_half2(src.x);
+  dst.y = float2_to_half2(src.y);
+}
+
+inline __device__ void from_float(uint4& dst, Float8_ src) {
+  dst.x = float2_to_half2(src.x);
+  dst.y = float2_to_half2(src.y);
+  dst.z = float2_to_half2(src.z);
+  dst.w = float2_to_half2(src.w);
+}
+
+// From float16 to float32.
+inline __device__ float to_float(uint16_t u) { return half_to_float(u); }
+
+inline __device__ float2 to_float(uint32_t u) { return half2_to_float2(u); }
+
+inline __device__ Float4_ to_float(uint2 u) {
+  Float4_ tmp;
+  tmp.x = half2_to_float2(u.x);
+  tmp.y = half2_to_float2(u.y);
+  return tmp;
+}
+
+inline __device__ Float8_ to_float(uint4 u) {
+  Float8_ tmp;
+  tmp.x = half2_to_float2(u.x);
+  tmp.y = half2_to_float2(u.y);
+  tmp.z = half2_to_float2(u.z);
+  tmp.w = half2_to_float2(u.w);
+  return tmp;
+}
+
+// Zero-out a variable.
+inline __device__ void zero(uint16_t& dst) { dst = uint16_t(0); }
+
+}  // namespace vllm_musa
diff --git a/vllm_musa/csrc/attention/dtype_float32.muh b/vllm_musa/csrc/attention/dtype_float32.muh
new file mode 100644
index 000000000..55da0b932
--- /dev/null
+++ b/vllm_musa/csrc/attention/dtype_float32.muh
@@ -0,0 +1,251 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
+ * and
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention_utils.h
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include "attention_generic.muh"
+
+#include <stdint.h>
+
+namespace vllm_musa {
+
+// Define custom FP32 vector data types.
+struct Float4_ {
+  float2 x;
+  float2 y;
+};
+
+struct Float8_ {
+  float2 x;
+  float2 y;
+  float2 z;
+  float2 w;
+};
+
+// FP32 vector types for Q, K, V.
+template <>
+struct Vec<float, 1> {
+  using Type = float;
+};
+template <>
+struct Vec<float, 2> {
+  using Type = float2;
+};
+template <>
+struct Vec<float, 4> {
+  using Type = float4;
+};
+
+// FP32 accumulator vector types corresponding to Vec.
+template <>
+struct FloatVec<float> {
+  using Type = float;
+};
+template <>
+struct FloatVec<float2> {
+  using Type = float2;
+};
+template <>
+struct FloatVec<float4> {
+  using Type = float4;
+};
+
+// Vector addition.
+inline __device__ float add(float a, float b) { return a + b; }
+
+inline __device__ float2 add(float2 a, float2 b) {
+  float2 c;
+  c.x = add(a.x, b.x);
+  c.y = add(a.y, b.y);
+  return c;
+}
+
+inline __device__ float4 add(float4 a, float4 b) {
+  float4 c;
+  c.x = add(a.x, b.x);
+  c.y = add(a.y, b.y);
+  c.z = add(a.z, b.z);
+  c.w = add(a.w, b.w);
+  return c;
+}
+
+// Vector multiplication.
+template <>
+inline __device__ float mul<float, float>(float a, float b) {
+  return a * b;
+}
+
+template <>
+inline __device__ float2 mul(float2 a, float2 b) {
+  float2 c;
+  c.x = a.x * b.x;
+  c.y = a.y * b.y;
+  return c;
+}
+
+template <>
+inline __device__ float2 mul(float a, float2 b) {
+  float2 c;
+  c.x = a * b.x;
+  c.y = a * b.y;
+  return c;
+}
+
+template <>
+inline __device__ float4 mul(float4 a, float4 b) {
+  float4 c;
+  c.x = a.x * b.x;
+  c.y = a.y * b.y;
+  c.z = a.z * b.z;
+  c.w = a.w * b.w;
+  return c;
+}
+
+template <>
+inline __device__ float4 mul(float a, float4 b) {
+  float4 c;
+  c.x = a * b.x;
+  c.y = a * b.y;
+  c.z = a * b.z;
+  c.w = a * b.w;
+  return c;
+}
+
+// Vector fused multiply-add.
+inline __device__ float fma(float a, float b, float c) { return a * b + c; }
+
+inline __device__ float2 fma(float2 a, float2 b, float2 c) {
+  float2 d;
+  d.x = fma(a.x, b.x, c.x);
+  d.y = fma(a.y, b.y, c.y);
+  return d;
+}
+
+inline __device__ float2 fma(float a, float2 b, float2 c) {
+  float2 d;
+  d.x = fma(a, b.x, c.x);
+  d.y = fma(a, b.y, c.y);
+  return d;
+}
+
+inline __device__ float4 fma(float4 a, float4 b, float4 c) {
+  float4 d;
+  d.x = fma(a.x, b.x, c.x);
+  d.y = fma(a.y, b.y, c.y);
+  d.z = fma(a.z, b.z, c.z);
+  d.w = fma(a.w, b.w, c.w);
+  return d;
+}
+
+inline __device__ float4 fma(float a, float4 b, float4 c) {
+  float4 d;
+  d.x = fma(a, b.x, c.x);
+  d.y = fma(a, b.y, c.y);
+  d.z = fma(a, b.z, c.z);
+  d.w = fma(a, b.w, c.w);
+  return d;
+}
+
+inline __device__ Float4_ fma(float a, Float4_ b, Float4_ c) {
+  Float4_ d;
+  d.x = fma(a, b.x, c.x);
+  d.y = fma(a, b.y, c.y);
+  return d;
+}
+
+inline __device__ Float8_ fma(float a, Float8_ b, Float8_ c) {
+  Float8_ d;
+  d.x = fma(a, b.x, c.x);
+  d.y = fma(a, b.y, c.y);
+  d.z = fma(a, b.z, c.z);
+  d.w = fma(a, b.w, c.w);
+  return d;
+}
+
+// Vector sum.
+template <>
+inline __device__ float sum(float v) {
+  return v;
+}
+
+template <>
+inline __device__ float sum(float2 v) {
+  return v.x + v.y;
+}
+
+template <>
+inline __device__ float sum(float4 v) {
+  return v.x + v.y + v.z + v.w;
+}
+
+template <>
+inline __device__ float sum(Float4_ v) {
+  return v.x.x + v.x.y + v.y.x + v.y.y;
+}
+
+template <>
+inline __device__ float sum(Float8_ v) {
+  return v.x.x + v.x.y + v.y.x + v.y.y + v.z.x + v.z.y + v.w.x + v.w.y;
+}
+
+// Vector dot product.
+inline __device__ float dot(float a, float b) { return a * b; }
+
+inline __device__ float dot(float2 a, float2 b) {
+  float2 c = mul<float2, float2, float2>(a, b);
+  return c.x + c.y;
+}
+
+inline __device__ float dot(Float4_ a, Float4_ b) {
+  float2 acc = mul<float2, float2, float2>(a.x, b.x);
+  acc = fma(a.y, b.y, acc);
+  return acc.x + acc.y;
+}
+
+inline __device__ float dot(Float8_ a, Float8_ b) {
+  float2 acc = mul<float2, float2, float2>(a.x, b.x);
+  acc = fma(a.y, b.y, acc);
+  acc = fma(a.z, b.z, acc);
+  acc = fma(a.w, b.w, acc);
+  return acc.x + acc.y;
+}
+
+// From float to float.
+inline __device__ void from_float(float& dst, float src) { dst = src; }
+
+inline __device__ void from_float(float2& dst, float2 src) { dst = src; }
+
+inline __device__ void from_float(float4& dst, float4 src) { dst = src; }
+
+// From float to float.
+inline __device__ float to_float(float u) { return u; }
+
+inline __device__ float2 to_float(float2 u) { return u; }
+
+inline __device__ float4 to_float(float4 u) { return u; }
+
+inline __device__ Float4_ to_float(Float4_ u) { return u; }
+
+inline __device__ Float8_ to_float(Float8_ u) { return u; }
+
+// Zero-out a variable.
+inline __device__ void zero(float& dst) { dst = 0.f; }
+
+}  // namespace vllm_musa
diff --git a/vllm_musa/csrc/attention/dtype_fp8.muh b/vllm_musa/csrc/attention/dtype_fp8.muh
new file mode 100644
index 000000000..e47fa73e1
--- /dev/null
+++ b/vllm_musa/csrc/attention/dtype_fp8.muh
@@ -0,0 +1,34 @@
+#pragma once
+
+#include "attention_generic.muh"
+
+namespace vllm_musa {
+
+enum class Fp8KVCacheDataType {
+  kAuto = 0,
+  kFp8E4M3 = 1,
+  kFp8E5M2 = 2,
+};
+
+// fp8 vector types for quantization of kv cache
+template <>
+struct Vec<uint8_t, 1> {
+  using Type = uint8_t;
+};
+
+template <>
+struct Vec<uint8_t, 2> {
+  using Type = uint16_t;
+};
+
+template <>
+struct Vec<uint8_t, 4> {
+  using Type = uint32_t;
+};
+
+template <>
+struct Vec<uint8_t, 8> {
+  using Type = uint2;
+};
+
+}  // namespace vllm_musa
diff --git a/vllm_musa/csrc/attention/merge_attn_states.mu b/vllm_musa/csrc/attention/merge_attn_states.mu
new file mode 100644
index 000000000..63d20169e
--- /dev/null
+++ b/vllm_musa/csrc/attention/merge_attn_states.mu
@@ -0,0 +1,175 @@
+#include <optional>
+#include <torch/all.h>
+#include <algorithm>
+
+#include "attention_dtypes.h"
+#include "attention_utils.muh"
+
+namespace vllm_musa {
+
+// Implements section 2.2 of https://www.arxiv.org/pdf/2501.01005
+// can be used to combine partial attention results (in the split-KV case)
+template <typename scalar_t, const uint NUM_THREADS>
+__global__ void merge_attn_states_kernel(
+    scalar_t* output, float* output_lse, const scalar_t* prefix_output,
+    const float* prefix_lse, const scalar_t* suffix_output,
+    const float* suffix_lse, const uint num_tokens, const uint num_heads,
+    const uint head_size) {
+  using pack_128b_t = uint4;
+  const uint pack_size = 16 / sizeof(scalar_t);
+  const uint threads_per_head = head_size / pack_size;
+
+  const uint global_idx = blockIdx.x * NUM_THREADS + threadIdx.x;
+  const uint token_head_threads = num_tokens * num_heads * threads_per_head;
+
+  if (global_idx >= token_head_threads) return;
+
+  // global_idx -> token_idx + head_idx + pack_idx
+  const uint token_head_idx = global_idx / threads_per_head;
+  const uint pack_idx = global_idx % threads_per_head;
+
+  const uint token_idx = token_head_idx / num_heads;
+  const uint head_idx = token_head_idx % num_heads;
+
+  const uint pack_offset = pack_idx * pack_size;  // (0~15)*8, etc.
+  const uint head_offset =
+      token_idx * num_heads * head_size + head_idx * head_size;
+  const scalar_t* prefix_head_ptr = prefix_output + head_offset;
+  const scalar_t* suffix_head_ptr = suffix_output + head_offset;
+  scalar_t* output_head_ptr = output + head_offset;
+
+  float p_lse = prefix_lse[head_idx * num_tokens + token_idx];
+  float s_lse = suffix_lse[head_idx * num_tokens + token_idx];
+  p_lse = std::isinf(p_lse) ? -std::numeric_limits<float>::infinity() : p_lse;
+  s_lse = std::isinf(s_lse) ? -std::numeric_limits<float>::infinity() : s_lse;
+
+  const float max_lse = fmaxf(p_lse, s_lse);
+  p_lse = p_lse - max_lse;
+  s_lse = s_lse - max_lse;
+  const float p_se = expf(p_lse);
+  const float s_se = expf(s_lse);
+  const float out_se = p_se + s_se;
+  const float p_scale = p_se / out_se;
+  const float s_scale = s_se / out_se;
+
+  if (pack_offset < head_size) {
+    // Pack 128b load
+    pack_128b_t p_out_pack = reinterpret_cast<const pack_128b_t*>(
+        prefix_head_ptr)[pack_offset / pack_size];
+    pack_128b_t s_out_pack = reinterpret_cast<const pack_128b_t*>(
+        suffix_head_ptr)[pack_offset / pack_size];
+    pack_128b_t o_out_pack;
+
+#pragma unroll
+    for (uint i = 0; i < pack_size; ++i) {
+      // Always use float for FMA to keep high precision.
+      // half(uint16_t), bfloat16, float -> float.
+      const float p_out_f =
+          to_float(reinterpret_cast<const scalar_t*>(&p_out_pack)[i]);
+      const float s_out_f =
+          to_float(reinterpret_cast<const scalar_t*>(&s_out_pack)[i]);
+      // fma: a * b + c = p_out_f * p_scale + (s_out_f * s_scale)
+      const float o_out_f = p_out_f * p_scale + (s_out_f * s_scale);
+      // float -> half(uint16_t), bfloat16, float.
+      from_float(reinterpret_cast<scalar_t*>(&o_out_pack)[i], o_out_f);
+    }
+
+    // Pack 128b storage
+    reinterpret_cast<pack_128b_t*>(output_head_ptr)[pack_offset / pack_size] =
+        o_out_pack;
+  }
+  // We only need to write to output_lse once per head.
+  if (output_lse != nullptr && pack_idx == 0) {
+    float out_lse = logf(out_se) + max_lse;
+    output_lse[head_idx * num_tokens + token_idx] = out_lse;
+  }
+}
+
+}  // namespace vllm_musa
+
+// The following macro is used to dispatch the conversion function based on
+// the output data type. The FN is a macro that calls a function with
+// template<typename scalar_t>.
+#define DISPATCH_BY_SCALAR_DTYPE(scalar_dtype, fn)                      \
+  {                                                                     \
+    if (scalar_dtype == at::ScalarType::Float) {                        \
+      fn(float);                                                        \
+    } else if (scalar_dtype == at::ScalarType::Half) {                  \
+      fn(uint16_t);                                                     \
+    } else if (scalar_dtype == at::ScalarType::BFloat16) {              \
+      fn(__nv_bfloat16);                                                \
+    } else {                                                            \
+      TORCH_CHECK(false, "Unsupported data type of O: ", scalar_dtype); \
+    }                                                                   \
+  }
+
+#define LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS)                     \
+  {                                                                         \
+    vllm_musa::merge_attn_states_kernel<scalar_t, NUM_THREADS>                   \
+        <<<grid, block, 0, stream>>>(                                       \
+            reinterpret_cast<scalar_t*>(output.data_ptr()), output_lse_ptr, \
+            reinterpret_cast<scalar_t*>(prefix_output.data_ptr()),          \
+            reinterpret_cast<float*>(prefix_lse.data_ptr()),                \
+            reinterpret_cast<scalar_t*>(suffix_output.data_ptr()),          \
+            reinterpret_cast<float*>(suffix_lse.data_ptr()), num_tokens,    \
+            num_heads, head_size);                                          \
+  }
+
+/*@brief Merges the attention states from prefix and suffix
+ * into the output tensor. NUM_TOKENS: n, NUM_HEADS: h, HEAD_SIZE: d
+ *
+ * @param output [n,h,d] The output tensor to store the merged attention states.
+ * @param output_lse [h,d] Optional tensor to store the log-sum-exp values.
+ * @param prefix_output [n,h,d] The prefix attention states.
+ * @param prefix_lse [h,n] The log-sum-exp values for the prefix attention
+ * states.
+ * @param suffix_output [n,h,d] The suffix attention states.
+ * @param suffix_lse [h,n] The log-sum-exp values for the suffix attention
+ * states.
+ */
+template <typename scalar_t>
+void merge_attn_states_launcher(torch::Tensor& output,
+                                std::optional<torch::Tensor> output_lse,
+                                const torch::Tensor& prefix_output,
+                                const torch::Tensor& prefix_lse,
+                                const torch::Tensor& suffix_output,
+                                const torch::Tensor& suffix_lse) {
+  constexpr uint NUM_THREADS = 128;
+  const uint num_tokens = output.size(0);
+  const uint num_heads = output.size(1);
+  const uint head_size = output.size(2);
+  const uint pack_size = 16 / sizeof(scalar_t);
+  TORCH_CHECK(head_size % pack_size == 0,
+              "headsize must be multiple of pack_size:", pack_size);
+  float* output_lse_ptr = nullptr;
+  if (output_lse.has_value()) {
+    output_lse_ptr = output_lse.value().data_ptr<float>();
+  }
+  // Process one pack elements per thread. for float, the
+  // pack_size is 4 for half/bf16, the pack_size is 8.
+  const uint threads_per_head = head_size / pack_size;
+  const uint total_threads = num_tokens * num_heads * threads_per_head;
+
+  dim3 block(NUM_THREADS);
+  dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS);
+  const at::cuda::OptionalCUDAGuard device_guard(prefix_output.device());
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  LAUNCH_MERGE_ATTN_STATES(scalar_t, NUM_THREADS);
+}
+
+#define CALL_MERGE_ATTN_STATES_LAUNCHER(scalar_t)                           \
+  {                                                                         \
+    merge_attn_states_launcher<scalar_t>(output, output_lse, prefix_output, \
+                                         prefix_lse, suffix_output,         \
+                                         suffix_lse);                       \
+  }
+
+void merge_attn_states(torch::Tensor& output,
+                       std::optional<torch::Tensor> output_lse,
+                       const torch::Tensor& prefix_output,
+                       const torch::Tensor& prefix_lse,
+                       const torch::Tensor& suffix_output,
+                       const torch::Tensor& suffix_lse) {
+  DISPATCH_BY_SCALAR_DTYPE(output.dtype(), CALL_MERGE_ATTN_STATES_LAUNCHER);
+}
diff --git a/vllm_musa/csrc/attention/paged_attention_v1.mu b/vllm_musa/csrc/attention/paged_attention_v1.mu
new file mode 100644
index 000000000..c69857c30
--- /dev/null
+++ b/vllm_musa/csrc/attention/paged_attention_v1.mu
@@ -0,0 +1,193 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "attention_kernels.muh"
+#include "../quantization/quant_utils.muh"
+
+#define WARP_SIZE 32
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
+
+#define LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE)                                \
+  VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(                     \
+      ((void*)vllm_musa::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE,        \
+                                              BLOCK_SIZE, NUM_THREADS,      \
+                                              KV_DTYPE, IS_BLOCK_SPARSE>),  \
+      shared_mem_size);                                                     \
+  vllm_musa::paged_attention_v1_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,        \
+                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE>   \
+      <<<grid, block, shared_mem_size, stream>>>(                           \
+          out_ptr, query_ptr, key_cache_ptr, value_cache_ptr, num_kv_heads, \
+          scale, block_tables_ptr, seq_lens_ptr, max_num_blocks_per_seq,    \
+          alibi_slopes_ptr, q_stride, kv_block_stride, kv_head_stride,      \
+          k_scale_ptr, v_scale_ptr, tp_rank, blocksparse_local_blocks,      \
+          blocksparse_vert_stride, blocksparse_block_size,                  \
+          blocksparse_head_sliding_step);
+
+// TODO(woosuk): Tune NUM_THREADS.
+template <typename T, typename CACHE_T, int BLOCK_SIZE,
+          vllm_musa::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
+          int NUM_THREADS = 128>
+void paged_attention_v1_launcher(
+    torch::Tensor& out, torch::Tensor& query, torch::Tensor& key_cache,
+    torch::Tensor& value_cache, int num_kv_heads, float scale,
+    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
+    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
+  int num_seqs = query.size(0);
+  int num_heads = query.size(1);
+  int head_size = query.size(2);
+  int max_num_blocks_per_seq = block_tables.size(1);
+  int q_stride = query.stride(0);
+  int kv_block_stride = key_cache.stride(0);
+  int kv_head_stride = key_cache.stride(1);
+
+  [[maybe_unused]] int thread_group_size = MAX(WARP_SIZE / BLOCK_SIZE, 1);
+  assert(head_size % thread_group_size == 0);
+
+  // NOTE: alibi_slopes is optional.
+  const float* alibi_slopes_ptr =
+      alibi_slopes
+          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
+          : nullptr;
+
+  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
+  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
+  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
+  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
+  int* block_tables_ptr = block_tables.data_ptr<int>();
+  int* seq_lens_ptr = seq_lens.data_ptr<int>();
+  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
+  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
+
+  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
+  int padded_max_seq_len =
+      DIVIDE_ROUND_UP(max_seq_len, BLOCK_SIZE) * BLOCK_SIZE;
+  int logits_size = padded_max_seq_len * sizeof(float);
+  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
+  // Python-side check in vllm.worker.worker._check_if_can_support_max_seq_len
+  // Keep that in sync with the logic here!
+  int shared_mem_size = std::max(logits_size, outputs_size);
+
+  dim3 grid(num_heads, num_seqs, 1);
+  dim3 block(NUM_THREADS);
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  switch (head_size) {
+    // NOTE(woosuk): To reduce the compilation time, we only compile for the
+    // head sizes that we use in the model. However, we can easily extend this
+    // to support any head size which is a multiple of 16.
+    case 32:
+      LAUNCH_PAGED_ATTENTION_V1(32);
+      break;
+    case 64:
+      LAUNCH_PAGED_ATTENTION_V1(64);
+      break;
+    case 80:
+      LAUNCH_PAGED_ATTENTION_V1(80);
+      break;
+    case 96:
+      LAUNCH_PAGED_ATTENTION_V1(96);
+      break;
+    case 112:
+      LAUNCH_PAGED_ATTENTION_V1(112);
+      break;
+    case 120:
+      LAUNCH_PAGED_ATTENTION_V1(120);
+      break;
+    case 128:
+      LAUNCH_PAGED_ATTENTION_V1(128);
+      break;
+    case 192:
+      LAUNCH_PAGED_ATTENTION_V1(192);
+      break;
+    case 256:
+      LAUNCH_PAGED_ATTENTION_V1(256);
+      break;
+    default:
+      TORCH_CHECK(false, "Unsupported head size: ", head_size);
+      break;
+  }
+}
+
+#define CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)  \
+  paged_attention_v1_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,              \
+                              IS_BLOCK_SPARSE>(                              \
+      out, query, key_cache, value_cache, num_kv_heads, scale, block_tables, \
+      seq_lens, max_seq_len, alibi_slopes, k_scale, v_scale, tp_rank,        \
+      blocksparse_local_blocks, blocksparse_vert_stride,                     \
+      blocksparse_block_size, blocksparse_head_sliding_step);
+
+#define CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
+  if (is_block_sparse) {                                                   \
+    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
+  } else {                                                                 \
+    CALL_V1_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
+  }
+
+// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
+// 1, 2, 4, 64, 128, 256.
+#define CALL_V1_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
+  switch (block_size) {                                           \
+    case 8:                                                       \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
+      break;                                                      \
+    case 16:                                                      \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
+      break;                                                      \
+    case 32:                                                      \
+      CALL_V1_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
+      break;                                                      \
+    default:                                                      \
+      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                      \
+  }
+
+void paged_attention_v1(
+    torch::Tensor& out,    // [num_seqs, num_heads, head_size]
+    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::Tensor&
+        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
+    torch::Tensor&
+        value_cache,       // [num_blocks, num_heads, head_size, block_size]
+    int64_t num_kv_heads,  // [num_heads]
+    double scale,
+    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::Tensor& seq_lens,      // [num_seqs]
+    int64_t block_size, int64_t max_seq_len,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step) {
+  const bool is_block_sparse = (blocksparse_vert_stride > 1);
+
+  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
+                             CALL_V1_LAUNCHER_BLOCK_SIZE)
+}
+
+#undef WARP_SIZE
+#undef MAX
+#undef MIN
+#undef DIVIDE_ROUND_UP
\ No newline at end of file
diff --git a/vllm_musa/csrc/attention/paged_attention_v2.mu b/vllm_musa/csrc/attention/paged_attention_v2.mu
new file mode 100644
index 000000000..790975958
--- /dev/null
+++ b/vllm_musa/csrc/attention/paged_attention_v2.mu
@@ -0,0 +1,206 @@
+/*
+ * Adapted from
+ * https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
+ * Copyright (c) 2023, The vLLM team.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "attention_kernels.muh"
+#include "../quantization/quant_utils.muh"
+#include "../musa.h"
+#define WARP_SIZE 128
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))
+
+#define LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE)                                   \
+  vllm_musa::paged_attention_v2_kernel<T, CACHE_T, HEAD_SIZE, BLOCK_SIZE,           \
+                                  NUM_THREADS, KV_DTYPE, IS_BLOCK_SPARSE,      \
+                                  PARTITION_SIZE>                              \
+      <<<grid, block, shared_mem_size, stream>>>(                              \
+          exp_sums_ptr, max_logits_ptr, tmp_out_ptr, query_ptr, key_cache_ptr, \
+          value_cache_ptr, num_kv_heads, scale, block_tables_ptr,              \
+          seq_lens_ptr, max_num_blocks_per_seq, alibi_slopes_ptr, q_stride,    \
+          kv_block_stride, kv_head_stride, k_scale_ptr, v_scale_ptr, tp_rank,  \
+          blocksparse_local_blocks, blocksparse_vert_stride,                   \
+          blocksparse_block_size, blocksparse_head_sliding_step);              \
+  vllm_musa::paged_attention_v2_reduce_kernel<T, HEAD_SIZE, NUM_THREADS,            \
+                                         PARTITION_SIZE>                       \
+      <<<reduce_grid, block, reduce_shared_mem_size, stream>>>(                \
+          out_ptr, exp_sums_ptr, max_logits_ptr, tmp_out_ptr, seq_lens_ptr,    \
+          max_num_partitions);
+
+template <typename T, typename CACHE_T, int BLOCK_SIZE,
+          vllm_musa::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
+          int NUM_THREADS = 128, int PARTITION_SIZE = 512>
+void paged_attention_v2_launcher(
+    torch::Tensor& out, torch::Tensor& exp_sums, torch::Tensor& max_logits,
+    torch::Tensor& tmp_out, torch::Tensor& query, torch::Tensor& key_cache,
+    torch::Tensor& value_cache, int num_kv_heads, float scale,
+    torch::Tensor& block_tables, torch::Tensor& seq_lens, int max_seq_len,
+    const std::optional<torch::Tensor>& alibi_slopes, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int tp_rank,
+    const int blocksparse_local_blocks, const int blocksparse_vert_stride,
+    const int blocksparse_block_size, const int blocksparse_head_sliding_step) {
+  int num_seqs = query.size(0);
+  int num_heads = query.size(1);
+  int head_size = query.size(2);
+  int max_num_blocks_per_seq = block_tables.size(1);
+  int q_stride = query.stride(0);
+  int kv_block_stride = key_cache.stride(0);
+  int kv_head_stride = key_cache.stride(1);
+
+  [[maybe_unused]] int thread_group_size = MAX(WARP_SIZE / BLOCK_SIZE, 1);
+  assert(head_size % thread_group_size == 0);
+
+  // NOTE: alibi_slopes is optional.
+  const float* alibi_slopes_ptr =
+      alibi_slopes
+          ? reinterpret_cast<const float*>(alibi_slopes.value().data_ptr())
+          : nullptr;
+
+  T* out_ptr = reinterpret_cast<T*>(out.data_ptr());
+  float* exp_sums_ptr = reinterpret_cast<float*>(exp_sums.data_ptr());
+  float* max_logits_ptr = reinterpret_cast<float*>(max_logits.data_ptr());
+  T* tmp_out_ptr = reinterpret_cast<T*>(tmp_out.data_ptr());
+  T* query_ptr = reinterpret_cast<T*>(query.data_ptr());
+  CACHE_T* key_cache_ptr = reinterpret_cast<CACHE_T*>(key_cache.data_ptr());
+  CACHE_T* value_cache_ptr = reinterpret_cast<CACHE_T*>(value_cache.data_ptr());
+  int* block_tables_ptr = block_tables.data_ptr<int>();
+  int* seq_lens_ptr = seq_lens.data_ptr<int>();
+  const float* k_scale_ptr = reinterpret_cast<const float*>(k_scale.data_ptr());
+  const float* v_scale_ptr = reinterpret_cast<const float*>(v_scale.data_ptr());
+
+  constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
+  int max_num_partitions = DIVIDE_ROUND_UP(max_seq_len, PARTITION_SIZE);
+  int logits_size = PARTITION_SIZE * sizeof(float);
+  int outputs_size = (NUM_WARPS / 2) * head_size * sizeof(float);
+
+  // For paged attention v2 kernel.
+  dim3 grid(num_heads, num_seqs, max_num_partitions);
+  int shared_mem_size = std::max(logits_size, outputs_size);
+  // For paged attention v2 reduce kernel.
+  dim3 reduce_grid(num_heads, num_seqs);
+  int reduce_shared_mem_size = 2 * max_num_partitions * sizeof(float);
+
+  dim3 block(NUM_THREADS);
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  switch (head_size) {
+    // NOTE(woosuk): To reduce the compilation time, we only compile for the
+    // head sizes that we use in the model. However, we can easily extend this
+    // to support any head size which is a multiple of 16.
+    case 32:
+      LAUNCH_PAGED_ATTENTION_V2(32);
+      break;
+    case 64:
+      LAUNCH_PAGED_ATTENTION_V2(64);
+      break;
+    case 80:
+      LAUNCH_PAGED_ATTENTION_V2(80);
+      break;
+    case 96:
+      LAUNCH_PAGED_ATTENTION_V2(96);
+      break;
+    case 112:
+      LAUNCH_PAGED_ATTENTION_V2(112);
+      break;
+    case 120:
+      LAUNCH_PAGED_ATTENTION_V2(120);
+      break;
+    case 128:
+      LAUNCH_PAGED_ATTENTION_V2(128);
+      break;
+    case 192:
+      LAUNCH_PAGED_ATTENTION_V2(192);
+      break;
+    case 256:
+      LAUNCH_PAGED_ATTENTION_V2(256);
+      break;
+    case 512:
+      LAUNCH_PAGED_ATTENTION_V2(512);
+      break;
+    default:
+      TORCH_CHECK(false, "Unsupported head size: ", head_size);
+      break;
+  }
+}
+
+#define CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, KV_DTYPE, IS_BLOCK_SPARSE)   \
+  paged_attention_v2_launcher<T, CACHE_T, BLOCK_SIZE, KV_DTYPE,               \
+                              IS_BLOCK_SPARSE>(                               \
+      out, exp_sums, max_logits, tmp_out, query, key_cache, value_cache,      \
+      num_kv_heads, scale, block_tables, seq_lens, max_seq_len, alibi_slopes, \
+      k_scale, v_scale, tp_rank, blocksparse_local_blocks,                    \
+      blocksparse_vert_stride, blocksparse_block_size,                        \
+      blocksparse_head_sliding_step);
+
+#define CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE) \
+  if (is_block_sparse) {                                                   \
+    CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, true);       \
+  } else {                                                                 \
+    CALL_V2_LAUNCHER(T, CACHE_T, BLOCK_SIZE, IS_FP8_KV_CACHE, false);      \
+  }
+
+// NOTE(woosuk): To reduce the compilation time, we omitted block sizes
+// 1, 2, 4, 64, 128, 256.
+#define CALL_V2_LAUNCHER_BLOCK_SIZE(T, CACHE_T, KV_DTYPE)         \
+  switch (block_size) {                                           \
+    case 8:                                                       \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 8, KV_DTYPE);         \
+      break;                                                      \
+    case 16:                                                      \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 16, KV_DTYPE);        \
+      break;                                                      \
+    case 32:                                                      \
+      CALL_V2_LAUNCHER_SPARSITY(T, CACHE_T, 32, KV_DTYPE);        \
+      break;                                                      \
+    default:                                                      \
+      TORCH_CHECK(false, "Unsupported block size: ", block_size); \
+      break;                                                      \
+  }
+
+void paged_attention_v2(
+    torch::Tensor& out,         // [num_seqs, num_heads, head_size]
+    torch::Tensor& exp_sums,    // [num_seqs, num_heads, max_num_partitions]
+    torch::Tensor& max_logits,  // [num_seqs, num_heads, max_num_partitions]
+    torch::Tensor&
+        tmp_out,  // [num_seqs, num_heads, max_num_partitions, head_size]
+    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
+    torch::Tensor&
+        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
+    torch::Tensor&
+        value_cache,       // [num_blocks, num_heads, head_size, block_size]
+    int64_t num_kv_heads,  // [num_heads]
+    double scale,
+    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
+    torch::Tensor& seq_lens,      // [num_seqs]
+    int64_t block_size, int64_t max_seq_len,
+    const std::optional<torch::Tensor>& alibi_slopes,
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step) {
+  const bool is_block_sparse = (blocksparse_vert_stride > 1);
+  DISPATCH_BY_KV_CACHE_DTYPE(query.dtype(), kv_cache_dtype,
+                             CALL_V2_LAUNCHER_BLOCK_SIZE)
+}
+
+#undef WARP_SIZE
+#undef MAX
+#undef MIN
+#undef DIVIDE_ROUND_UP
\ No newline at end of file
diff --git a/vllm_musa/csrc/cache_kernels.mu b/vllm_musa/csrc/cache_kernels.mu
new file mode 100644
index 000000000..380b0b814
--- /dev/null
+++ b/vllm_musa/csrc/cache_kernels.mu
@@ -0,0 +1,634 @@
+#include <torch/all.h>
+#include "musa.h"
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+#include "attention/dtype_fp8.muh"
+#include "quantization/quant_utils.muh"
+#include "linear/common.muh"
+
+#include <algorithm>
+#include <cassert>
+#include <map>
+#include <vector>
+
+void swap_blocks(torch::Tensor& src, torch::Tensor& dst,
+                 const torch::Tensor& block_mapping) {
+  torch::Device src_device = src.device();
+  torch::Device dst_device = dst.device();
+  cudaMemcpyKind memcpy_type;
+
+  if (src_device.is_privateuseone() && dst_device.is_privateuseone()) {
+    TORCH_CHECK(src_device.index() == dst_device.index(),
+                "src and dst must be on the same GPU");
+    memcpy_type = cudaMemcpyDeviceToDevice;
+  } else if (src_device.is_privateuseone() && dst_device.is_cpu()) {
+    memcpy_type = cudaMemcpyDeviceToHost;
+  } else if (src_device.is_cpu() && dst_device.is_privateuseone()) {
+    memcpy_type = cudaMemcpyHostToDevice;
+  } else {
+    TORCH_CHECK(false, "Invalid device combination");
+  }
+
+  // NOTE(youkaichao): keep in mind that `block_mapping` should be
+  // a cpu tensor, otherwise every `item` call will require a gpu-cpu
+  // synchronization.
+  TORCH_CHECK(block_mapping.device().is_cpu(), "block_mapping must be on CPU");
+
+  char* src_ptr = static_cast<char*>(src.data_ptr());
+  char* dst_ptr = static_cast<char*>(dst.data_ptr());
+
+  // We use the stride instead of numel in case the cache is padded for memory
+  // alignment reasons, we assume the blocks data (inclusive of any padding)
+  // is contiguous in memory
+  const int64_t block_size_in_bytes = src.element_size() * src.stride(0);
+  const at::cuda::OptionalCUDAGuard device_guard(
+      src_device.is_privateuseone() ? src_device : dst_device);
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  // NOTE(woosuk): This can be slow if the number of blocks is large.
+  const int64_t num_blocks = block_mapping.size(0);
+  for (size_t i = 0; i < num_blocks; i++) {
+    int64_t src_block_number = block_mapping[i][0].item<int64_t>();
+    int64_t dst_block_number = block_mapping[i][1].item<int64_t>();
+    int64_t src_offset = src_block_number * block_size_in_bytes;
+    int64_t dst_offset = dst_block_number * block_size_in_bytes;
+    cudaMemcpyAsync(dst_ptr + dst_offset, src_ptr + src_offset,
+                    block_size_in_bytes, memcpy_type, stream);
+  }
+}
+
+namespace vllm_musa {
+
+// Grid: (num_layers, num_pairs)
+template <typename scalar_t>
+__global__ void copy_blocks_kernel(int64_t* key_cache_ptrs,
+                                   int64_t* value_cache_ptrs,
+                                   const int64_t* __restrict__ block_mapping,
+                                   const int numel_per_block) {
+  const int layer_idx = blockIdx.x;
+  const int pair_idx = blockIdx.y;
+
+  scalar_t* key_cache = reinterpret_cast<scalar_t*>(key_cache_ptrs[layer_idx]);
+  scalar_t* value_cache =
+      reinterpret_cast<scalar_t*>(value_cache_ptrs[layer_idx]);
+  int64_t src_block_number = block_mapping[2 * pair_idx];
+  int64_t dst_block_number = block_mapping[2 * pair_idx + 1];
+
+  const int64_t src_block_offset = src_block_number * numel_per_block;
+  const int64_t dst_block_offset = dst_block_number * numel_per_block;
+  for (int i = threadIdx.x; i < numel_per_block; i += blockDim.x) {
+    int64_t src_offset = src_block_offset + i;
+    int64_t dst_offset = dst_block_offset + i;
+    key_cache[dst_offset] = key_cache[src_offset];
+  }
+  for (int i = threadIdx.x; i < numel_per_block; i += blockDim.x) {
+    int64_t src_offset = src_block_offset + i;
+    int64_t dst_offset = dst_block_offset + i;
+    value_cache[dst_offset] = value_cache[src_offset];
+  }
+}
+
+}  // namespace vllm_musa
+
+// Note: the key_caches and value_caches vectors are constant but
+// not the Tensors they contain. The vectors need to be const refs
+// in order to satisfy pytorch's C++ operator registration code.
+void copy_blocks(std::vector<torch::Tensor> const& key_caches,
+                 std::vector<torch::Tensor> const& value_caches,
+                 const torch::Tensor& block_mapping) {
+  int num_layers = key_caches.size();
+  TORCH_CHECK(num_layers == value_caches.size());
+  if (num_layers == 0) {
+    return;
+  }
+  torch::Device cache_device = key_caches[0].device();
+  TORCH_CHECK(cache_device.is_privateuseone());
+
+  // Create data structures for the kernel.
+  // Create an array of pointers to the key and value caches.
+  int64_t key_cache_ptrs[num_layers];
+  int64_t value_cache_ptrs[num_layers];
+  for (int layer_idx = 0; layer_idx < num_layers; ++layer_idx) {
+    key_cache_ptrs[layer_idx] =
+        reinterpret_cast<int64_t>(key_caches[layer_idx].data_ptr());
+    value_cache_ptrs[layer_idx] =
+        reinterpret_cast<int64_t>(value_caches[layer_idx].data_ptr());
+  }
+
+  // block_mapping is a 2D tensor with shape (num_pairs, 2).
+  int num_pairs = block_mapping.size(0);
+
+  // Move the data structures to the GPU.
+  // NOTE: This synchronizes the CPU and GPU.
+  torch::Tensor key_cache_ptrs_tensor =
+      torch::from_blob(key_cache_ptrs, {num_layers}, torch::kInt64)
+          .to(cache_device);
+  torch::Tensor value_cache_ptrs_tensor =
+      torch::from_blob(value_cache_ptrs, {num_layers}, torch::kInt64)
+          .to(cache_device);
+
+  // Launch the kernel.
+  const int numel_per_block = key_caches[0][0].numel();
+  dim3 grid(num_layers, num_pairs);
+  dim3 block(std::min(1024, numel_per_block));
+  const at::cuda::OptionalCUDAGuard device_guard(cache_device);
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(
+      key_caches[0].scalar_type(), "copy_blocks_kernel", ([&] {
+        vllm_musa::copy_blocks_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            key_cache_ptrs_tensor.data_ptr<int64_t>(),
+            value_cache_ptrs_tensor.data_ptr<int64_t>(),
+            block_mapping.data_ptr<int64_t>(), numel_per_block);
+      }));
+}
+
+namespace vllm_musa {
+
+template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
+__global__ void reshape_and_cache_kernel(
+    const scalar_t* __restrict__ key,    // [num_tokens, num_heads, head_size]
+    const scalar_t* __restrict__ value,  // [num_tokens, num_heads, head_size]
+    cache_t* __restrict__ key_cache,     // [num_blocks, num_heads, head_size/x,
+                                         // block_size, x]
+    cache_t* __restrict__ value_cache,   // [num_blocks, num_heads, head_size,
+                                         // block_size]
+    const int64_t* __restrict__ slot_mapping,  // [num_tokens]
+    const int key_stride, const int value_stride, const int num_heads,
+    const int head_size, const int block_size, const int x,
+    const float* k_scale, const float* v_scale) {
+  const int64_t token_idx = blockIdx.x;
+  const int64_t slot_idx = slot_mapping[token_idx];
+  if (slot_idx < 0) {
+    // Padding token that should be ignored.
+    return;
+  }
+
+  const int64_t block_idx = slot_idx / block_size;
+  const int64_t block_offset = slot_idx % block_size;
+
+  const int n = num_heads * head_size;
+  for (int i = threadIdx.x; i < n; i += blockDim.x) {
+    const int64_t src_key_idx = token_idx * key_stride + i;
+    const int64_t src_value_idx = token_idx * value_stride + i;
+
+    const int head_idx = i / head_size;
+    const int head_offset = i % head_size;
+    const int x_idx = head_offset / x;
+    const int x_offset = head_offset % x;
+
+    const int64_t 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;
+    const int64_t tgt_value_idx =
+        block_idx * num_heads * head_size * block_size +
+        head_idx * head_size * block_size + head_offset * block_size +
+        block_offset;
+    scalar_t tgt_key = key[src_key_idx];
+    scalar_t tgt_value = value[src_value_idx];
+    key_cache[tgt_key_idx] = tgt_key;
+    value_cache[tgt_value_idx] = tgt_value;
+  }
+}
+
+}  // namespace vllm_musa
+
+// KV_T is the stored data type of kv-cache.
+// CACHE_T is the data type of key and value tensors.
+// KV_DTYPE is the real data type of kv-cache.
+#define CALL_RESHAPE_AND_CACHE(KV_T, CACHE_T, KV_DTYPE)               \
+  vllm_musa::reshape_and_cache_kernel<KV_T, CACHE_T, KV_DTYPE>             \
+      <<<grid, block, 0, stream>>>(                                   \
+          reinterpret_cast<KV_T*>(key.data_ptr()),                    \
+          reinterpret_cast<KV_T*>(value.data_ptr()),                  \
+          reinterpret_cast<CACHE_T*>(key_cache.data_ptr()),           \
+          reinterpret_cast<CACHE_T*>(value_cache.data_ptr()),         \
+          slot_mapping.data_ptr<int64_t>(), key_stride, value_stride, \
+          num_heads, head_size, block_size, x,                        \
+          reinterpret_cast<const float*>(k_scale.data_ptr()),         \
+          reinterpret_cast<const float*>(v_scale.data_ptr()));
+
+void reshape_and_cache(
+    torch::Tensor& key,    // [num_tokens, num_heads, head_size]
+    torch::Tensor& value,  // [num_tokens, num_heads, head_size]
+    torch::Tensor&
+        key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
+    torch::Tensor&
+        value_cache,  // [num_blocks, num_heads, head_size, block_size]
+    torch::Tensor& slot_mapping,  // [num_tokens]
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale) {
+  int num_tokens = key.size(0);
+  int num_heads = key.size(1);
+  int head_size = key.size(2);
+  int block_size = key_cache.size(3);
+  int x = key_cache.size(4);
+
+  int key_stride = key.stride(0);
+  int value_stride = value.stride(0);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min(num_heads * head_size, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
+                             CALL_RESHAPE_AND_CACHE)
+}
+
+namespace vllm_musa {
+  template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
+  __global__ void concat_and_cache_mla_kernel(
+      const scalar_t* __restrict__ kv_c,  // [num_tokens, kv_lora_rank]
+      const scalar_t* __restrict__ k_pe,  // [num_tokens, pe_dim]
+      cache_t* __restrict__ kv_cache,  // [num_blocks, block_size, (kv_lora_rank
+                                       // + pe_dim)]
+      const int64_t* __restrict__ slot_mapping,  // [num_tokens]
+      const int block_stride,                    //
+      const int entry_stride,                    //
+      const int kv_c_stride,                     //
+      const int k_pe_stride,                     //
+      const int kv_lora_rank,                    //
+      const int pe_dim,                          //
+      const int block_size,                      //
+      const float* scale                         //
+  ) {
+    const int64_t token_idx = blockIdx.x;
+    const int64_t slot_idx = slot_mapping[token_idx];
+    // NOTE: slot_idx can be -1 if the token is padded
+    if (slot_idx < 0) {
+      return;
+    }
+    const int64_t block_idx = slot_idx / block_size;
+    const int64_t block_offset = slot_idx % block_size;
+
+    auto copy = [&](const scalar_t* __restrict__ src, cache_t* __restrict__ dst,
+                    int src_stride, int dst_stride, int size, int offset) {
+      for (int i = threadIdx.x; i < size; i += blockDim.x) {
+        const int64_t src_idx = token_idx * src_stride + i;
+        const int64_t dst_idx =
+            block_idx * block_stride + block_offset * entry_stride + i + offset;
+        dst[dst_idx] = src[src_idx];
+      }
+    };
+
+    copy(kv_c, kv_cache, kv_c_stride, block_stride, kv_lora_rank, 0);
+    copy(k_pe, kv_cache, k_pe_stride, block_stride, pe_dim, kv_lora_rank);
+  }
+
+}  // namespace vllm_musa
+
+// KV_T is the stored data type of kv-cache.
+// CACHE_T is the data type of key and value tensors.
+// KV_DTYPE is the real data type of kv-cache.
+#define CALL_CONCAT_AND_CACHE_MLA(KV_T, CACHE_T, KV_DTYPE)              \
+  vllm_musa::concat_and_cache_mla_kernel<KV_T, CACHE_T, KV_DTYPE>            \
+      <<<grid, block, 0, stream>>>(                                     \
+          reinterpret_cast<KV_T*>(kv_c.data_ptr()),                     \
+          reinterpret_cast<KV_T*>(k_pe.data_ptr()),                     \
+          reinterpret_cast<CACHE_T*>(kv_cache.data_ptr()),              \
+          slot_mapping.data_ptr<int64_t>(), block_stride, entry_stride, \
+          kv_c_stride, k_pe_stride, kv_lora_rank, pe_dim, block_size,   \
+          reinterpret_cast<const float*>(scale.data_ptr()));
+
+void concat_and_cache_mla(
+    torch::Tensor& kv_c,          // [num_tokens, kv_lora_rank]
+    torch::Tensor& k_pe,          // [num_tokens, pe_dim]
+    torch::Tensor& kv_cache,      // [num_blocks, block_size, (kv_lora_rank +
+                                  // pe_dim)]
+    torch::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
+    const std::string& kv_cache_dtype, torch::Tensor& scale) {
+  // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
+  // slot_mapping.size(0) because of padding for CUDA graphs.
+  // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
+  // both include padding.
+  // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
+  // since key includes padding for CUDA graphs, while slot_mapping does not.
+  // In this case, slot_mapping.size(0) represents the actual number of tokens
+  // before padding.
+  // For compatibility with both cases, we use slot_mapping.size(0) as the
+  // number of tokens.
+  int num_tokens = slot_mapping.size(0);
+  int kv_lora_rank = kv_c.size(1);
+  int pe_dim = k_pe.size(1);
+  int block_size = kv_cache.size(1);
+
+  TORCH_CHECK(kv_cache.size(2) == kv_lora_rank + pe_dim);
+
+  int kv_c_stride = kv_c.stride(0);
+  int k_pe_stride = k_pe.stride(0);
+  int block_stride = kv_cache.stride(0);
+  int entry_stride = kv_cache.stride(1);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min(kv_lora_rank, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(kv_c));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  DISPATCH_BY_KV_CACHE_DTYPE(kv_c.dtype(), kv_cache_dtype,
+                             CALL_CONCAT_AND_CACHE_MLA);
+}
+
+namespace vllm_musa {
+  // Kernel for MLA, which works on a single joint kv_cache
+  // Grid: (num_layers, num_pairs)
+  template <typename scalar_t>
+  __global__ void copy_blocks_mla_kernel(
+      int64_t* cache_ptrs, const int64_t* __restrict__ block_mapping,
+      const int mem_footprint_per_block) {
+    const int layer_idx = blockIdx.x;
+    const int pair_idx = blockIdx.y;
+    scalar_t* cache = reinterpret_cast<scalar_t*>(cache_ptrs[layer_idx]);
+    int64_t src_block = block_mapping[2 * pair_idx];
+    int64_t dst_block = block_mapping[2 * pair_idx + 1];
+    int64_t src_offset = src_block * mem_footprint_per_block;
+    int64_t dst_offset = dst_block * mem_footprint_per_block;
+    for (int i = threadIdx.x; i < mem_footprint_per_block; i += blockDim.x) {
+      cache[dst_offset + i] = cache[src_offset + i];
+    }
+  }
+}  // namespace vllm_musa
+
+// copy blocks kernel for MLA (assumes a joint KV-cache)
+void copy_blocks_mla(std::vector<torch::Tensor> const& kv_caches,
+  const torch::Tensor& block_mapping) {
+  int num_layers = kv_caches.size();
+  if (num_layers == 0) {
+    return;
+  }
+  torch::Device cache_device = kv_caches[0].device();
+  TORCH_CHECK(cache_device.is_privateuseone(), "kv_cache must be on MUSA");
+
+  std::vector<int64_t> cache_ptrs(num_layers);
+  for (int layer_idx = 0; layer_idx < num_layers; ++layer_idx) {
+    cache_ptrs[layer_idx] =
+    reinterpret_cast<int64_t>(kv_caches[layer_idx].data_ptr());
+  }
+  torch::Tensor cache_ptrs_tensor =
+  torch::from_blob(cache_ptrs.data(), {num_layers}, torch::kInt64)
+  .to(cache_device);
+
+  int num_pairs = block_mapping.size(0);
+  // We use the stride instead of numel in case the cache is padded for memory
+  // alignment reasons, we assume the blocks data (inclusive of any padding)
+  // is contiguous in memory
+  int mem_footprint_per_block = kv_caches[0].stride(0);
+  dim3 grid(num_layers, num_pairs);
+  dim3 block(std::min(1024, mem_footprint_per_block));
+  const at::cuda::OptionalCUDAGuard device_guard(cache_device);
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(
+    kv_caches[0].scalar_type(), "copy_blocks_mla_kernel", ([&] {
+      vllm_musa::copy_blocks_mla_kernel<scalar_t><<<grid, block, 0, stream>>>(
+    cache_ptrs_tensor.data_ptr<int64_t>(),
+    block_mapping.data_ptr<int64_t>(), mem_footprint_per_block);
+  }));
+}
+
+namespace vllm_musa {
+
+// grid is launched with dimensions (batch, num_splits)
+template <typename scalar_t>
+__global__ void gather_cache(
+    const scalar_t* __restrict__ src_cache,   // [NUM_BLOCKS, BLOCK_SIZE,
+                                              // ENTRIES...]
+    scalar_t* __restrict__ dst,               // [TOT_TOKENS, ENTRIES...]
+    const int32_t* __restrict__ block_table,  // [BATCH, BLOCK_INDICES]
+    const int32_t* __restrict__ cu_seq_lens,  // [BATCH+1]
+    const int32_t block_size, const int32_t entry_size,
+    const int64_t block_table_stride, const int64_t cache_block_stride,
+    const int64_t cache_entry_stride, const int64_t dst_entry_stride,
+    const int32_t* __restrict__ seq_starts) {  // Optional: starting offsets per
+                                               // batch
+
+  const int64_t bid = blockIdx.x;  // Batch ID
+  const int32_t num_splits = gridDim.y;
+  const int32_t split = blockIdx.y;
+  const int32_t seq_start = cu_seq_lens[bid];
+  const int32_t seq_end = cu_seq_lens[bid + 1];
+  const int32_t seq_len = seq_end - seq_start;
+  const int32_t tot_blocks = ceil_div(seq_len, block_size);
+  const int32_t split_blocks = ceil_div(tot_blocks, num_splits);
+
+  const int32_t split_start = split * split_blocks;
+  const int32_t split_end = min((split + 1) * split_blocks, tot_blocks);
+
+  const bool is_active_split = (split_start < tot_blocks);
+  const bool is_last_split = (split_end == tot_blocks);
+
+  if (!is_active_split) return;
+
+  int32_t full_blocks_end = split_end;
+  int32_t partial_block_size = 0;
+
+  // Adjust the pointer for the block_table for this batch.
+  // If seq_starts is provided, compute an offset based on (seq_starts[bid] /
+  // page_size)
+  const int32_t batch_offset = bid * block_table_stride;
+  int32_t offset = 0;
+  if (seq_starts != nullptr) {
+    offset = seq_starts[bid] / block_size;
+  }
+  const int32_t* batch_block_table = block_table + batch_offset + offset;
+
+  // Adjust dst pointer based on the cumulative sequence lengths.
+  dst += seq_start * dst_entry_stride;
+
+  if (is_last_split) {
+    partial_block_size = seq_len % block_size;
+    if (partial_block_size) full_blocks_end -= 1;
+  }
+
+  auto copy_entry = [&](const scalar_t* __restrict__ _src,
+                        scalar_t* __restrict__ _dst) {
+    for (int i = threadIdx.x; i < entry_size; i += blockDim.x)
+      _dst[i] = _src[i];
+  };
+
+  for (int pid = split_start; pid < full_blocks_end; ++pid) {
+    auto block_id = batch_block_table[pid];
+    auto block_start_ptr = src_cache + block_id * cache_block_stride;
+    auto block_dst_ptr = dst + pid * block_size * dst_entry_stride;
+    for (int eid = 0; eid < block_size; ++eid) {
+      copy_entry(block_start_ptr + eid * cache_entry_stride,
+                 block_dst_ptr + eid * dst_entry_stride);
+    }
+  }
+
+  if (partial_block_size) {
+    auto block_id = batch_block_table[full_blocks_end];
+    auto block_start_ptr = src_cache + block_id * cache_block_stride;
+    auto block_dst_ptr = dst + full_blocks_end * block_size * dst_entry_stride;
+    for (int eid = 0; eid < partial_block_size; ++eid) {
+      copy_entry(block_start_ptr + eid * cache_entry_stride,
+                 block_dst_ptr + eid * dst_entry_stride);
+    }
+  }
+}
+
+}  // namespace vllm_musa
+
+// Macro to dispatch the kernel based on the data type.
+#define CALL_GATHER_CACHE(CPY_DTYPE)                                    \
+  vllm_musa::gather_cache<CPY_DTYPE><<<grid, block, 0, stream>>>(            \
+      reinterpret_cast<CPY_DTYPE*>(src_cache.data_ptr()),               \
+      reinterpret_cast<CPY_DTYPE*>(dst.data_ptr()),                     \
+      block_table.data_ptr<int32_t>(), cu_seq_lens.data_ptr<int32_t>(), \
+      block_size, entry_size, block_table_stride, cache_block_stride,   \
+      cache_entry_stride, dst_entry_stride, seq_starts_ptr);
+
+// Gather sequences from the cache into the destination tensor.
+//  - cu_seq_lens contains the cumulative sequence lengths for each batch
+//  - block_table contains the cache block indices for each sequence
+//  - Optionally, seq_starts (if provided) offsets the starting block index by
+//  (seq_starts[bid] / page_size)
+void gather_cache(
+    torch::Tensor const& src_cache,    // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
+    torch::Tensor const& dst,          // [TOT_TOKENS, ENTRIES...]
+    torch::Tensor const& block_table,  // [BATCH, BLOCK_INDICES]
+    torch::Tensor const& cu_seq_lens,  // [BATCH+1]
+    int64_t batch_size,
+    std::optional<torch::Tensor> seq_starts = std::nullopt) {
+  at::cuda::OptionalCUDAGuard device_guard(src_cache.device());
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  int32_t block_size = src_cache.size(1);
+  int32_t entry_size = src_cache.flatten(2, -1).size(2);
+
+  TORCH_CHECK(block_table.dtype() == torch::kInt32,
+              "block_table must be int32");
+  TORCH_CHECK(cu_seq_lens.dtype() == torch::kInt32,
+              "cu_seq_lens must be int32");
+  if (seq_starts.has_value()) {
+    TORCH_CHECK(seq_starts.value().dtype() == torch::kInt32,
+                "seq_starts must be int32");
+  }
+
+  TORCH_CHECK(src_cache.device() == dst.device(),
+              "src_cache and dst must be on the same device");
+  TORCH_CHECK(src_cache.device() == block_table.device(),
+              "src_cache and block_table must be on the same device");
+  TORCH_CHECK(src_cache.device() == cu_seq_lens.device(),
+              "src_cache and cu_seq_lens must be on the same device");
+  if (seq_starts.has_value()) {
+    TORCH_CHECK(src_cache.device() == seq_starts.value().device(),
+                "src_cache and seq_starts must be on the same device");
+  }
+
+  int64_t block_table_stride = block_table.stride(0);
+  int64_t cache_block_stride = src_cache.stride(0);
+  int64_t cache_entry_stride = src_cache.stride(1);
+  int64_t dst_entry_stride = dst.stride(0);
+
+  // Decide on the number of splits based on the batch size.
+  int num_splits = batch_size > 128 ? 2 : batch_size > 64 ? 4 : 16;
+  dim3 grid(batch_size, num_splits);
+  dim3 block(1024);
+
+  TORCH_CHECK(src_cache.dtype() == dst.dtype(),
+              "src_cache and dst must have the same dtype");
+
+  const int dtype_bits = src_cache.element_size() * 8;
+  const int32_t* seq_starts_ptr =
+      seq_starts.has_value() ? seq_starts.value().data_ptr<int32_t>() : nullptr;
+
+  if (dtype_bits == 32) {
+    CALL_GATHER_CACHE(uint32_t);
+  } else if (dtype_bits == 16) {
+    CALL_GATHER_CACHE(uint16_t);
+  } else if (dtype_bits == 8) {
+    CALL_GATHER_CACHE(uint8_t);
+  } else {
+    TORCH_CHECK(false, "Unsupported data type width: ", dtype_bits);
+  }
+}
+
+namespace vllm_musa {
+  template <typename scalar_t, typename cache_t, Fp8KVCacheDataType kv_dt>
+  __global__ void reshape_and_cache_flash_kernel(
+      const scalar_t* __restrict__ key,    // [num_tokens, num_heads, head_size]
+      const scalar_t* __restrict__ value,  // [num_tokens, num_heads, head_size]
+      cache_t* __restrict__ key_cache,     // [num_blocks, block_size, num_heads,
+                                          // head_size]
+      cache_t* __restrict__ value_cache,   // [num_blocks, block_size, num_heads,
+                                          // head_size]
+      const int64_t* __restrict__ slot_mapping,  // [num_tokens]
+      const int block_stride, const int key_stride, const int value_stride,
+      const int num_heads, const int head_size, const int block_size,
+      const float* k_scale, const float* v_scale) {
+    const int64_t token_idx = blockIdx.x;
+    const int64_t slot_idx = slot_mapping[token_idx];
+    // NOTE: slot_idx can be -1 if the token is padded
+    if (slot_idx < 0) {
+      return;
+    }
+    const int64_t block_idx = slot_idx / block_size;
+    const int64_t block_offset = slot_idx % block_size;
+    const int n = num_heads * head_size;
+    for (int i = threadIdx.x; i < n; i += blockDim.x) {
+      const int64_t src_key_idx = token_idx * key_stride + i;
+      const int64_t src_value_idx = token_idx * value_stride + i;
+      const int head_idx = i / head_size;
+      const int head_offset = i % head_size;
+      const int64_t tgt_key_value_idx = block_idx * block_stride +
+                                        block_offset * num_heads * head_size +
+                                        head_idx * head_size + head_offset;
+      scalar_t tgt_key = key[src_key_idx];
+      scalar_t tgt_value = value[src_value_idx];
+      key_cache[tgt_key_value_idx] = tgt_key;
+      value_cache[tgt_key_value_idx] = tgt_value;
+    }
+  }
+}
+
+// KV_T is the stored data type of kv-cache.
+// CACHE_T is the data type of key and value tensors.
+// KV_DTYPE is the real data type of kv-cache.
+#define CALL_RESHAPE_AND_CACHE_FLASH(KV_T, CACHE_T, KV_DTYPE)         \
+  vllm_musa::reshape_and_cache_flash_kernel<KV_T, CACHE_T, KV_DTYPE>  \
+      <<<grid, block, 0, stream>>>(                                   \
+          reinterpret_cast<KV_T*>(key.data_ptr()),                    \
+          reinterpret_cast<KV_T*>(value.data_ptr()),                  \
+          reinterpret_cast<CACHE_T*>(key_cache.data_ptr()),           \
+          reinterpret_cast<CACHE_T*>(value_cache.data_ptr()),         \
+          slot_mapping.data_ptr<int64_t>(), block_stride, key_stride, \
+          value_stride, num_heads, head_size, block_size,             \
+          reinterpret_cast<const float*>(k_scale.data_ptr()),         \
+          reinterpret_cast<const float*>(v_scale.data_ptr()));
+
+void reshape_and_cache_flash(
+    torch::Tensor& key,        // [num_tokens, num_heads, head_size]
+    torch::Tensor& value,      // [num_tokens, num_heads, head_size]
+    torch::Tensor& key_cache,  // [num_blocks, block_size, num_heads, head_size]
+    torch::Tensor&
+        value_cache,  // [num_blocks, block_size, num_heads, head_size]
+    torch::Tensor& slot_mapping,  // [num_tokens] or [num_actual_tokens]
+    const std::string& kv_cache_dtype, torch::Tensor& k_scale,
+    torch::Tensor& v_scale) {
+  // NOTE(woosuk): In vLLM V1, key.size(0) can be different from
+  // slot_mapping.size(0) because of padding for CUDA graphs.
+  // In vLLM V0, key.size(0) is always equal to slot_mapping.size(0) because
+  // both include padding.
+  // In vLLM V1, however, key.size(0) can be larger than slot_mapping.size(0)
+  // since key includes padding for CUDA graphs, while slot_mapping does not.
+  // In this case, slot_mapping.size(0) represents the actual number of tokens
+  // before padding.
+  // For compatibility with both cases, we use slot_mapping.size(0) as the
+  // number of tokens.
+  int num_tokens = slot_mapping.size(0);
+  int num_heads = key.size(1);
+  int head_size = key.size(2);
+  int block_size = key_cache.size(1);
+
+  int key_stride = key.stride(0);
+  int value_stride = value.stride(0);
+  int block_stride = key_cache.stride(0);
+  TORCH_CHECK(key_cache.stride(0) == value_cache.stride(0));
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min(num_heads * head_size, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(key));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  DISPATCH_BY_KV_CACHE_DTYPE(key.dtype(), kv_cache_dtype,
+                             CALL_RESHAPE_AND_CACHE_FLASH);
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/cuda_compat.h b/vllm_musa/csrc/cuda_compat.h
new file mode 100644
index 000000000..10cf6ebeb
--- /dev/null
+++ b/vllm_musa/csrc/cuda_compat.h
@@ -0,0 +1,18 @@
+#pragma once
+
+#define WARP_SIZE 32
+
+#define VLLM_LDG(arg) __ldg(arg)
+
+#define VLLM_SHFL_XOR_SYNC(var, lane_mask) \
+    __shfl_xor_sync(uint32_t(-1), var, lane_mask)
+#define VLLM_SHFL_XOR_SYNC_WIDTH(var, lane_mask, width) \
+    __shfl_xor_sync(uint32_t(-1), var, lane_mask, width)
+
+#define VLLM_SHFL_SYNC(var, src_lane) __shfl_sync(uint32_t(-1), var, src_lane)
+
+#define VLLM_SHFL_DOWN_SYNC(var, lane_delta) \
+    __shfl_down_sync(uint32_t(-1), var, lane_delta)
+
+#define VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(FUNC, VAL) \
+    cudaFuncSetAttribute(FUNC, cudaFuncAttributeMaxDynamicSharedMemorySize, VAL)
diff --git a/vllm_musa/csrc/custom_all_reduce.mu b/vllm_musa/csrc/custom_all_reduce.mu
new file mode 100644
index 000000000..0951cea9b
--- /dev/null
+++ b/vllm_musa/csrc/custom_all_reduce.mu
@@ -0,0 +1,193 @@
+#include "musa.h"
+#include <torch/all.h>
+
+#include "custom_all_reduce.muh"
+
+// Fake pointer type, must match fptr_t type in ops.h.
+// We use this type alias to indicate when pointers are passed in as int64_t.
+using fptr_t = int64_t;
+static_assert(sizeof(void*) == sizeof(fptr_t));
+
+fptr_t init_custom_ar(const std::vector<fptr_t>& fake_ipc_ptrs,
+                      torch::Tensor& rank_data, int64_t rank,
+                      bool full_nvlink) {
+  // printf("init_custom_ar\n");
+  int world_size = fake_ipc_ptrs.size();
+  if (world_size > 8)
+    throw std::invalid_argument("world size > 8 is not supported");
+  if (world_size % 2 != 0)
+    throw std::invalid_argument("Odd num gpus is not supported for now");
+  if (rank < 0 || rank >= world_size)
+    throw std::invalid_argument("invalid rank passed in");
+
+  vllm_musa::Signal* ipc_ptrs[8];
+  for (int i = 0; i < world_size; i++) {
+    ipc_ptrs[i] = reinterpret_cast<vllm_musa::Signal*>(fake_ipc_ptrs[i]);
+    //printf("===== ipc_ptrs[%d]: %p, world_size: %d\n", i, ipc_ptrs[i], world_size);
+  }
+  return (fptr_t) new vllm_musa::CustomAllreduce(ipc_ptrs, rank_data.data_ptr(),
+                                            rank_data.numel(), rank, world_size,
+                                            full_nvlink);
+}
+
+/**
+ * Make sure tensor t's data lies completely within ((char)t.data_ptr()) +
+ * t.numel() * t.element_size(). This is slightly weaker than t.is_contiguous()
+ * because it allows transpose of contiguous slice (i.e. slicing the first
+ * dimension). Currently, we require this because stride information is not
+ * passed into the kernels and we treat input tensors as flat.
+ *
+ * Examples
+ * A = torch.zeros(3, 3, 3)
+ * 1. A: OK
+ * 2. A[1:]: OK
+ * 3. A.permute(2, 0, 1): OK
+ * 4. A[1:].permute(2, 0, 1): OK
+ * 5. A[None].expand(2, -1, -1, -1): Not OK
+ * 6. A[:, 1:, 1:]: Not OK
+ */
+bool _is_weak_contiguous(torch::Tensor& t) {
+  return t.is_contiguous() ||
+         (t.storage().nbytes() - t.storage_offset() * t.element_size() ==
+          t.numel() * t.element_size());
+}
+
+/**
+ * Performs an out-of-place allreduce and stores result in out.
+ *
+ * If _reg_buffer is null, assumes inp.data_ptr() is already IPC-registered.
+ * Otherwise, _reg_buffer is assumed to be IPC-registered and inp is first
+ * copied into _reg_buffer.
+ */
+void all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
+                fptr_t _reg_buffer, int64_t reg_buffer_sz_bytes) {
+  // printf("all_reduce from mu\n");
+  auto fa = reinterpret_cast<vllm_musa::CustomAllreduce*>(_fa);
+  const at::musa::OptionalMUSAGuard device_guard(device_of(inp));
+  auto stream = c10::musa::getCurrentMUSAStream().stream();
+
+  TORCH_CHECK_EQ(inp.scalar_type(), out.scalar_type());
+  TORCH_CHECK_EQ(inp.numel(), out.numel());
+  TORCH_CHECK(_is_weak_contiguous(out));
+  TORCH_CHECK(_is_weak_contiguous(inp));
+  auto input_size = inp.numel() * inp.element_size();
+  auto reg_buffer = reinterpret_cast<void*>(_reg_buffer);
+
+  int threads = 128;
+  int block_limit = 72;
+  if(at::musa::getMUSAArch() == 220)
+  {
+    threads = 512;
+  }
+  else if (at::musa::getMUSAArch() == 310)
+  {
+    threads = 64;
+  }
+
+  if (reg_buffer) {
+    TORCH_CHECK_LE(input_size, reg_buffer_sz_bytes);
+
+    //AT_MUSA_CHECK(musaMemcpyAsync(reg_buffer, inp.data_ptr(), input_size,
+    //                              musaMemcpyDeviceToDevice, stream));
+    ::musa::dnn::Tensor tensor_src;
+    ::musa::dnn::Tensor tensor_dst;
+
+    if (inp.scalar_type() == at::ScalarType::Float) {
+      tensor_src.SetType(::musa::dnn::Tensor::Type::FLOAT);
+      tensor_dst.SetType(::musa::dnn::Tensor::Type::FLOAT);
+    } else if (inp.scalar_type() == at::ScalarType::Half) {
+      tensor_src.SetType(::musa::dnn::Tensor::Type::HALF);
+      tensor_dst.SetType(::musa::dnn::Tensor::Type::HALF);
+    } else if (inp.scalar_type() == at::ScalarType::BFloat16) {
+      tensor_src.SetType(::musa::dnn::Tensor::Type::BFLOAT16);
+      tensor_dst.SetType(::musa::dnn::Tensor::Type::BFLOAT16);
+    } else {
+      throw std::runtime_error(
+          "custom allreduce only supports float32, float16 and bfloat16");
+    }
+
+    tensor_src.SetNdInfo({inp.numel()});
+    tensor_dst.SetNdInfo({inp.numel()});
+
+    tensor_src.SetAddr(inp.data_ptr());
+    tensor_dst.SetAddr(reg_buffer);
+
+    ::musa::dnn::Unary op;
+    op.SetMode(::musa::dnn::Unary::Mode::IDENTITY);
+    op.SetAlpha(0.0f);
+    op.SetBeta(0.0f);
+    at::musa::muHandle& h = at::GetMudnnHandle();
+    auto ret_code = op.Run(h, tensor_dst, tensor_src);
+
+    if (ret_code != ::musa::dnn::Status::SUCCESS) {
+      printf("Error when mudnn d2d copy, err code:%d\n", static_cast<int>(ret_code));
+      TORCH_CHECK(false);
+    }
+  } else {
+    reg_buffer = inp.data_ptr();
+  }
+  // printf("all_reduce-> out.numel: %d\n", out.numel());
+  switch (out.scalar_type()) {
+    case at::ScalarType::Float: {
+      fa->allreduce<float>(stream, reinterpret_cast<float*>(reg_buffer),
+                           reinterpret_cast<float*>(out.data_ptr()),
+                           out.numel(), threads, block_limit);
+      break;
+    }
+    case at::ScalarType::Half: {
+      fa->allreduce<half>(stream, reinterpret_cast<half*>(reg_buffer),
+                          reinterpret_cast<half*>(out.data_ptr()),
+                          out.numel(), threads, block_limit);
+      break;
+    }
+    case at::ScalarType::BFloat16: {
+      fa->allreduce<__mt_bfloat16>(stream, reinterpret_cast<__mt_bfloat16*>(reg_buffer),
+                                  reinterpret_cast<__mt_bfloat16*>(out.data_ptr()),
+                                  out.numel(), threads, block_limit);
+      break;
+    }
+    default:
+      throw std::runtime_error(
+          "custom allreduce only supports float32, float16 and bfloat16");
+  }
+}
+
+void dispose(fptr_t _fa) {
+  delete reinterpret_cast<vllm_musa::CustomAllreduce*>(_fa);
+}
+
+int64_t meta_size() { return sizeof(vllm_musa::Signal); }
+
+void register_buffer(fptr_t _fa, const std::vector<fptr_t>& fake_ipc_ptrs) {
+  // printf("register_buffer\n");
+  auto fa = reinterpret_cast<vllm_musa::CustomAllreduce*>(_fa);
+  TORCH_CHECK(fake_ipc_ptrs.size() == fa->world_size_);
+  void* ipc_ptrs[8];
+  for (int i = 0; i < fake_ipc_ptrs.size(); i++) {
+    ipc_ptrs[i] = reinterpret_cast<void*>(fake_ipc_ptrs[i]);
+  }
+  fa->register_buffer(ipc_ptrs);
+}
+
+// Use vector<int64_t> to represent byte data for python binding compatibility.
+std::tuple<std::vector<int64_t>, std::vector<int64_t>>
+get_graph_buffer_ipc_meta(fptr_t _fa) {
+  auto fa = reinterpret_cast<vllm_musa::CustomAllreduce*>(_fa);
+  auto [handle, offsets] = fa->get_graph_buffer_ipc_meta();
+  std::vector<int64_t> bytes(handle.begin(), handle.end());
+  return std::make_tuple(bytes, offsets);
+}
+
+// Use vector<int64_t> to represent byte data for python binding compatibility.
+void register_graph_buffers(fptr_t _fa,
+                            const std::vector<std::vector<int64_t>>& handles,
+                            const std::vector<std::vector<int64_t>>& offsets) {
+  auto fa = reinterpret_cast<vllm_musa::CustomAllreduce*>(_fa);
+  std::vector<std::string> bytes;
+  bytes.reserve(handles.size());
+  for (int i = 0; i < handles.size(); i++) {
+    bytes.emplace_back(handles[i].begin(), handles[i].end());
+  }
+  bytes.reserve(handles.size());
+  fa->register_graph_buffers(bytes, offsets);
+}
diff --git a/vllm_musa/csrc/custom_all_reduce.muh b/vllm_musa/csrc/custom_all_reduce.muh
new file mode 100644
index 000000000..e953404ea
--- /dev/null
+++ b/vllm_musa/csrc/custom_all_reduce.muh
@@ -0,0 +1,519 @@
+#pragma once
+
+#include "musa.h"
+#include <musa_bf16.h>
+#include <musa_fp16.h>
+#include <musa_runtime.h>
+
+#include <iostream>
+#include <array>
+#include <limits>
+#include <map>
+#include <unordered_map>
+#include <vector>
+
+#define CUDACHECK(cmd)                                              \
+  do {                                                              \
+    musaError_t e = cmd;                                            \
+    if (e != musaSuccess) {                                         \
+      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
+             musaGetErrorString(e));                                \
+      exit(EXIT_FAILURE);                                           \
+    }                                                               \
+  } while (0)
+
+namespace vllm_musa {
+
+constexpr int kMaxBlocks = 72;
+// Counter may overflow, but it's fine since unsigned int overflow is
+// well-defined behavior.
+using FlagType = uint32_t;
+struct Signal {
+  alignas(128) FlagType self_counter[kMaxBlocks][8];
+  // Two sets of peer counters are needed for two syncs. The reason is that
+  // it's possible for peer GPU block to arrive at the second sync point while
+  // the current GPU block haven't passed the first sync point. Thus, peer GPU
+  // may write counter+1 while current GPU is busy waiting for counter. We use
+  // alternating counter array to avoid this possibility.
+  alignas(128) FlagType peer_counter[2][kMaxBlocks][8];
+};
+
+struct __align__(16) RankData {
+  const void* ptrs[8];
+};
+
+struct __align__(16) RankSignals {
+  Signal* signals[8];
+};
+
+// like std::array, but aligned
+template <typename T, int sz>
+struct __align__(alignof(T) * sz) array_t {
+  T data[sz];
+  using type = T;
+  static constexpr int size = sz;
+};
+
+// use packed type to maximize memory efficiency
+// goal: generate ld.128 and st.128 instructions
+template <typename T>
+struct packed_t {
+  // the (P)acked type for load/store
+  using P = array_t<T, 16 / sizeof(T)>;
+  // the (A)ccumulator type for reduction
+  using A = array_t<float, 16 / sizeof(T)>;
+};
+
+#define DINLINE __device__ __forceinline__
+
+// scalar cast functions
+DINLINE float upcast_s(half val) { return __half2float(val); }
+
+template <typename T>
+DINLINE T downcast_s(float val);
+template <>
+DINLINE half downcast_s(float val) {
+  return __float2half(val);
+}
+
+// scalar add functions
+// for some reason when compiling with Pytorch, the + operator for half and
+// bfloat is disabled so we call the intrinsics directly
+DINLINE half& assign_add(half& a, half b) {
+  a = __hadd(a, b);
+  return a;
+}
+DINLINE float& assign_add(float& a, float b) { return a += b; }
+
+DINLINE float upcast_s(mt_bfloat16 val) { return __bfloat162float(val); }
+template <>
+DINLINE mt_bfloat16 downcast_s(float val) {
+  return __float2bfloat16(val);
+}
+DINLINE mt_bfloat16& assign_add(mt_bfloat16& a, mt_bfloat16 b) {
+  a = __hadd(a, b);
+  return a;
+}
+
+template <typename T, int N>
+DINLINE array_t<T, N>& packed_assign_add(array_t<T, N>& a, array_t<T, N> b) {
+#pragma unroll
+  for (int i = 0; i < N; i++) {
+    assign_add(a.data[i], b.data[i]);
+  }
+  return a;
+}
+
+template <typename T, int N>
+DINLINE array_t<float, N> upcast(array_t<T, N> val) {
+  if constexpr (std::is_same<T, float>::value) {
+    return val;
+  } else {
+    array_t<float, N> out;
+#pragma unroll
+    for (int i = 0; i < N; i++) {
+      out.data[i] = upcast_s(val.data[i]);
+    }
+    return out;
+  }
+}
+
+template <typename O>
+DINLINE O downcast(array_t<float, O::size> val) {
+  if constexpr (std::is_same<typename O::type, float>::value) {
+    return val;
+  } else {
+    O out;
+#pragma unroll
+    for (int i = 0; i < O::size; i++) {
+      out.data[i] = downcast_s<typename O::type>(val.data[i]);
+    }
+    return out;
+  }
+}
+
+static DINLINE void st_flag_release(FlagType* flag_addr, FlagType flag) {
+//   asm volatile("st.release.sys.global.u32 [%1], %0;" ::"r"(flag),
+//                "l"(flag_addr));
+    volatile_store((uint32_t)flag, (uint32_t*)flag_addr);
+}
+
+static DINLINE FlagType ld_flag_acquire(FlagType* flag_addr) {
+//   FlagType flag;
+//   asm volatile("ld.acquire.sys.global.u32 %0, [%1];"
+//                : "=r"(flag)
+//                : "l"(flag_addr));
+//   return flag;
+    asm("DMA.CFI_FLUSHINV.SLC.BYPASS");
+    return (uint32_t)volatile_load((uint32_t*)flag_addr);
+}
+
+static DINLINE void st_flag_volatile(FlagType* flag_addr, FlagType flag) {
+//   asm volatile("st.volatile.global.u32 [%1], %0;" ::"r"(flag), "l"(flag_addr));
+    volatile FlagType *volatile_ptr = (volatile FlagType*)flag_addr;
+    *volatile_ptr = flag;
+}
+
+static DINLINE FlagType ld_flag_volatile(FlagType* flag_addr) {
+//   FlagType flag;
+//   asm volatile("ld.volatile.global.u32 %0, [%1];"
+//                : "=r"(flag)
+//                : "l"(flag_addr));
+//   return flag;
+    volatile FlagType *volatile_ptr = (volatile FlagType *)flag_addr;
+    return *volatile_ptr;
+}
+
+// is_start: whether this is the very first synchronization barrier.
+// need_fence: whether a memory fence is needed. If true, a release-acquire
+// semantic is used to enforce memory access order before and after this
+// barrier.
+template <int ngpus, bool is_start, bool need_fence = false>
+DINLINE void multi_gpu_barrier(const RankSignals& sg, Signal* self_sg,
+                               int rank) {
+  if constexpr (!is_start) __syncthreads_lm();
+  static_assert(
+      !(is_start && need_fence));  // Start barrier shouldn't need fence.
+  if (threadIdx.x < ngpus) {
+    // Increment the counter. Technically we only need one counter, but we use
+    // multiple per block to eliminate the need to share the counter via smem.
+    auto val = self_sg->self_counter[blockIdx.x][threadIdx.x] += 1;
+    // Write the expected counter value to peer and wait for correct value from
+    // peer.
+    auto peer_counter_ptr =
+        &sg.signals[threadIdx.x]->peer_counter[val % 2][blockIdx.x][rank];
+    auto self_counter_ptr =
+        &self_sg->peer_counter[val % 2][blockIdx.x][threadIdx.x];
+    if constexpr (need_fence) {
+      st_flag_release(peer_counter_ptr, val);
+      while (ld_flag_acquire(self_counter_ptr) != val);
+    } else {
+      st_flag_volatile(peer_counter_ptr, val);
+      while (ld_flag_volatile(self_counter_ptr) != val);
+    }
+  }
+  if constexpr (is_start || need_fence) __syncthreads_lm();
+}
+
+#define AS1 __attribute__((address_space(1)))
+
+template <typename P, int ngpus, typename A>
+DINLINE P packed_reduce(const P AS1 *ptrs[], int idx) {
+  auto l_ptr = (const P *AS1)(ptrs[0]);
+  A tmp = upcast(l_ptr[idx]);
+#pragma unroll
+  for (int i = 1; i < ngpus; i++) {
+    // asm("CTRL.NOP");
+    auto i_ptr = (const P *AS1)(ptrs[i]);
+    packed_assign_add(tmp, upcast(i_ptr[idx]));
+  }
+  return downcast<P>(tmp);
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(1024, 1)
+    cross_device_reduce_1stage(RankData* _dp, RankSignals sg, Signal* self_sg,
+                               T* result, int rank, int size) {
+  using P = typename packed_t<T>::P;
+  using A = typename packed_t<T>::A;
+  // note: we don't reorder the address so the accumulation order is the same
+  // for all ranks, ensuring bitwise identical results
+  auto dp = *((RankData * AS1)_dp);
+  multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
+  // do the actual reduction
+  for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+       idx += gridDim.x * blockDim.x) {
+    ((P*)result)[idx] = packed_reduce<P, ngpus, A>((const P AS1 **)&dp.ptrs[0], idx);
+  }
+  multi_gpu_barrier<ngpus, false>(sg, self_sg, rank);
+}
+
+template <typename P>
+DINLINE P* get_tmp_buf(Signal* sg) {
+  return (P*)(((Signal*)sg) + 1);
+}
+
+template <typename T, int ngpus>
+__global__ void __launch_bounds__(1024, 1)
+    cross_device_reduce_2stage(RankData* _dp, RankSignals sg, Signal* self_sg,
+                               T* result, int rank, int size) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  using P = typename packed_t<T>::P;
+  using A = typename packed_t<T>::A;
+  int part = size / ngpus;
+  int start = rank * part;
+  int end = rank == ngpus - 1 ? size : start + part;
+  int largest_part = part + size % ngpus;
+  const P AS1* ptrs[ngpus];
+  P* tmps[ngpus];
+#pragma unroll
+  for (int i = 0; i < ngpus; i++) {
+    int target = (rank + i) % ngpus;
+    ptrs[i] = (const P AS1 *)_dp->ptrs[target];
+    tmps[i] = get_tmp_buf<P>(sg.signals[target]);
+  }
+  auto tmp_out = tmps[0];
+  multi_gpu_barrier<ngpus, true>(sg, self_sg, rank);
+  // stage 1: reduce scatter
+  for (int idx = start + tid; idx < end; idx += stride) {
+    tmp_out[idx - start] = packed_reduce<P, ngpus, A>(ptrs, idx);
+  }
+  multi_gpu_barrier<ngpus, false, true>(sg, self_sg, rank);
+
+  // stage 2: allgather. Note: it's important to match the tid between
+  // the two stages, because visibility across devices is only guaranteed
+  // between threads that have the same tid. If thread i computes the sum of
+  // start + i in the first stage, then thread i also gathers start + i from all
+  // ranks.
+  for (int idx = tid; idx < largest_part; idx += stride) {
+#pragma unroll
+    for (int i = 0; i < ngpus; i++) {
+      int gather_from_rank = ((rank + i) % ngpus);
+      if (gather_from_rank == ngpus - 1 || idx < part) {
+        int dst_idx = gather_from_rank * part + idx;
+        ((P*)result)[dst_idx] = tmps[i][idx];
+      }
+    }
+  }
+}
+
+using IPC_KEY = std::array<uint8_t, sizeof(musaIpcMemHandle_t)>;
+static_assert(sizeof(IPC_KEY) == sizeof(musaIpcMemHandle_t));
+static_assert(alignof(IPC_KEY) == alignof(musaIpcMemHandle_t));
+
+class CustomAllreduce {
+ public:
+  int rank_;
+  int world_size_;
+  bool full_nvlink_;
+
+  RankSignals sg_;
+  // Stores an map from a pointer to its peer pointters from all ranks.
+  std::unordered_map<void*, RankData*> buffers_;
+  Signal* self_sg_;
+
+  // Stores rank data from all ranks. This is mainly for musa graph purposes.
+  // For musa graph to work, all kernel arguments must be fixed during graph
+  // capture time. However, the peer pointers are not known during graph capture
+  // time. Therefore, during capture, we increment the rank data pointer and use
+  // that as the argument to the kernel. The kernel arguments are stored in
+  // graph_unreg_buffers_. The actual peer pointers will be filled in at the
+  // memory pointed to by the pointers in graph_unreg_buffers_ when
+  // the IPC handles are exchanged between ranks.
+  //
+  // The overall process looks like this:
+  // 1. Graph capture.
+  // 2. Each rank obtains the IPC handles for each addresses used during musa
+  // graph capture using get_graph_buffer_ipc_meta.
+  // 3. (In Python) all gather the IPC handles.
+  // 4. Obtain the peer pointers by opening the IPC handles, and store them in
+  // the rank data array at corresponding positions.
+  RankData *d_rank_data_base_, *d_rank_data_end_;
+  std::vector<void*> graph_unreg_buffers_;
+  // a map from IPC handles to opened IPC pointers
+  std::map<IPC_KEY, char*> ipc_handles_;
+
+  /**
+   * Signals are an array of ipc-enabled buffers from all ranks.
+   * For each of the buffer, the layout is as follows:
+   * | -- sizeof(Signal) -- | ------ a few MB ----- |
+   * The first section is for allreduce synchronization, and the second section
+   * is for storing the intermediate results required by some allreduce algos.
+   *
+   * Note: this class does not own any device memory. Any required buffers
+   * are passed in from the constructor.
+   */
+  CustomAllreduce(Signal** signals, void* rank_data, size_t rank_data_sz,
+                  int rank, int world_size, bool full_nvlink = true)
+      : rank_(rank),
+        world_size_(world_size),
+        full_nvlink_(full_nvlink),
+        self_sg_(signals[rank]),
+        d_rank_data_base_(reinterpret_cast<RankData*>(rank_data)),
+        d_rank_data_end_(d_rank_data_base_ + rank_data_sz / sizeof(RankData)) {
+    for (int i = 0; i < world_size_; i++) {
+      sg_.signals[i] = signals[i];
+    }
+  }
+
+  char* open_ipc_handle(const void* ipc_handle) {
+    auto [it, new_handle] =
+        ipc_handles_.insert({*((IPC_KEY*)ipc_handle), nullptr});
+    if (new_handle) {
+      char* ipc_ptr;
+      CUDACHECK(musaIpcOpenMemHandle((void**)&ipc_ptr,
+                                     *((const musaIpcMemHandle_t*)ipc_handle),
+                                     musaIpcMemLazyEnablePeerAccess));
+      it->second = ipc_ptr;
+    }
+    return it->second;
+  }
+
+  std::pair<std::string, std::vector<int64_t>> get_graph_buffer_ipc_meta() {
+    auto num_buffers = graph_unreg_buffers_.size();
+    auto handle_sz = sizeof(musaIpcMemHandle_t);
+    std::string handles(handle_sz * num_buffers, static_cast<char>(0));
+    std::vector<int64_t> offsets(num_buffers);
+    for (int i = 0; i < num_buffers; i++) {
+      auto ptr = graph_unreg_buffers_[i];
+      void* base_ptr;
+      // note: must share the base address of each allocation, or we get wrong
+      // address
+      if (muPointerGetAttribute(&base_ptr,
+                                MU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
+                                (MUdeviceptr)ptr) != MUSA_SUCCESS)
+        throw std::runtime_error("failed to get pointer attr");
+      CUDACHECK(musaIpcGetMemHandle(
+          (musaIpcMemHandle_t*)&handles[i * handle_sz], base_ptr));
+      offsets[i] = ((char*)ptr) - ((char*)base_ptr);
+    }
+    return std::make_pair(handles, offsets);
+  }
+
+  void check_rank_data_capacity(size_t num = 1) {
+    if (d_rank_data_base_ + num > d_rank_data_end_)
+      throw std::runtime_error(
+          "Rank data buffer is overflowed by " +
+          std::to_string(d_rank_data_base_ + num - d_rank_data_end_));
+  }
+
+  /**
+   * Register already-shared IPC pointers.
+   */
+  void register_buffer(void** ptrs) {
+    check_rank_data_capacity();
+    RankData data;
+    for (int i = 0; i < world_size_; i++) {
+      data.ptrs[i] = ptrs[i];
+    }
+    auto d_data = d_rank_data_base_++;
+    CUDACHECK(
+        musaMemcpy(d_data, &data, sizeof(RankData), musaMemcpyHostToDevice));
+    buffers_[ptrs[rank_]] = d_data;
+  }
+
+  // Note: when registering graph buffers, we intentionally choose to not
+  // deduplicate the addresses. That means if the allocator reuses some
+  // addresses, they will be registered again. This is to account for the remote
+  // possibility of different allocation patterns between ranks. For example,
+  // rank 1 may get the same input address for the second allreduce, but rank 2
+  // got a different address. IPC handles have internal reference counting
+  // mechanism so overhead should be small.
+  void register_graph_buffers(
+      const std::vector<std::string>& handles,
+      const std::vector<std::vector<int64_t>>& offsets) {
+    auto num_buffers = graph_unreg_buffers_.size();
+    check_rank_data_capacity(num_buffers);
+    std::vector<RankData> rank_data(num_buffers);
+    for (int i = 0; i < num_buffers; i++) {
+      auto self_ptr = graph_unreg_buffers_[i];
+      auto& rd = rank_data[i];
+      for (int j = 0; j < world_size_; j++) {
+        if (j != rank_) {
+          char* handle =
+              open_ipc_handle(&handles[j][i * sizeof(musaIpcMemHandle_t)]);
+          handle += offsets[j][i];
+          rd.ptrs[j] = handle;
+        } else {
+          rd.ptrs[j] = self_ptr;
+        }
+      }
+    }
+    CUDACHECK(musaMemcpy(d_rank_data_base_, rank_data.data(),
+                         sizeof(RankData) * num_buffers,
+                         musaMemcpyHostToDevice));
+    d_rank_data_base_ += num_buffers;
+    graph_unreg_buffers_.clear();
+  }
+
+  /**
+   * Performs allreduce, assuming input has already been registered.
+   *
+   * Block and grid default configs are results after careful grid search. Using
+   * 36 blocks give the best or close to the best runtime on the devices I
+   * tried: A100, A10, A30, T4, V100. You'll notice that NCCL kernels also only
+   * take a small amount of SMs. Not quite sure the underlying reason, but my
+   * guess is that too many SMs will cause contention on NVLink bus.
+   */
+  template <typename T>
+  void allreduce(musaStream_t stream, T* input, T* output, int size,
+                 int threads = 512, int block_limit = 36) {
+    auto d = packed_t<T>::P::size;
+    if (size % d != 0)
+      throw std::runtime_error(
+          "custom allreduce currently requires input length to be multiple "
+          "of " +
+          std::to_string(d));
+    if (block_limit > kMaxBlocks)
+      throw std::runtime_error("max supported block limit is " +
+                               std::to_string(kMaxBlocks) + ". Got " +
+                               std::to_string(block_limit));
+
+    RankData* ptrs;
+    musaStreamCaptureStatus status;
+    CUDACHECK(musaStreamIsCapturing(stream, &status));
+    if (status == musaStreamCaptureStatusActive) {
+      ptrs = d_rank_data_base_ + graph_unreg_buffers_.size();
+      graph_unreg_buffers_.push_back(input);
+    } else {
+      auto it = buffers_.find(input);
+      if (it == buffers_.end())
+      throw std::runtime_error(
+          "buffer address " +
+          std::to_string(reinterpret_cast<uint64_t>(input)) +
+          " is not registered!");
+      ptrs = it->second;
+    }
+
+    size /= d;
+    auto bytes = size * sizeof(typename packed_t<T>::P);
+    int blocks = std::min(block_limit, (size + threads - 1) / threads);
+#define KL(ngpus, name)                                                       \
+  name<T, ngpus><<<blocks, threads, 0, stream>>>(ptrs, sg_, self_sg_, output, \
+                                                 rank_, size);
+    // TODO(hanzhi713): Threshold is different for A100 and H100.
+    // Add per device threshold.
+#define REDUCE_CASE(ngpus)                            \
+  case ngpus: {                                       \
+    if (world_size_ == 2) {                           \
+      KL(ngpus, cross_device_reduce_1stage);          \
+    } else if (full_nvlink_) {                        \
+      if ((world_size_ <= 4 && bytes < 512 * 1024) || \
+          (world_size_ <= 8 && bytes < 256 * 1024)) { \
+        KL(ngpus, cross_device_reduce_1stage);        \
+      } else {                                        \
+        KL(ngpus, cross_device_reduce_2stage);        \
+      }                                               \
+    }                                                 \
+    break;                                            \
+  }
+
+    switch (world_size_) {
+      REDUCE_CASE(2)
+      REDUCE_CASE(4)
+      REDUCE_CASE(6)
+      REDUCE_CASE(8)
+      default:
+        throw std::runtime_error(
+            "custom allreduce only supports num gpus in (2,4,6,8). Actual num "
+            "gpus = " +
+            std::to_string(world_size_));
+    }
+#undef REDUCE_CASE
+#undef KL
+  }
+
+  ~CustomAllreduce() {
+    for (auto [_, ptr] : ipc_handles_) {
+      CUDACHECK(musaIpcCloseMemHandle(ptr));
+    }
+  }
+};
+/**
+ * To inspect PTX/SASS, copy paste this header file to compiler explorer and add
+ a template instantiation:
+ * template void vllm::CustomAllreduce::allreduce<half>(musaStream_t, half *,
+ half *, int, int, int);
+*/
+}  // namespace vllm_musa
\ No newline at end of file
diff --git a/vllm_musa/csrc/custom_all_reduce_test.mu b/vllm_musa/csrc/custom_all_reduce_test.mu
new file mode 100644
index 000000000..93a05d10e
--- /dev/null
+++ b/vllm_musa/csrc/custom_all_reduce_test.mu
@@ -0,0 +1,360 @@
+/**
+ * This is a standalone test for custom allreduce.
+ * To compile, make sure you have MPI and NCCL installed in your system.
+ * export MPI_HOME=XXX
+ * nvcc -O2 -arch=native -std=c++17 custom_all_reduce_test.cu -o
+ * custom_all_reduce_test -lmccl -I${MPI_HOME}/include -lmpi
+ *
+ * Warning: this C++ test is not designed to be very readable and was used
+ * during the rapid prototyping process.
+ *
+ * To run:
+ * mpirun --allow-run-as-root -np 8 ./custom_all_reduce_test
+ */
+ #include <musa.h>
+ #include <murand_kernel.h>
+ #include <stdio.h>
+ #include <stdlib.h>
+
+ #include <limits>
+ #include <vector>
+
+ #include "musa_profiler_api.h"
+ #include "custom_all_reduce.muh"
+ #include "mpi.h"
+#include "mccl.h"
+#include "custom_all_reduce.muh"
+
+#define MUSACHECK(cmd)                                              \
+  do {                                                              \
+    musaError_t e = cmd;                                            \
+    if (e != musaSuccess) {                                         \
+      printf("Failed: Cuda error %s:%d '%s'\n", __FILE__, __LINE__, \
+             musaGetErrorString(e));                                \
+      exit(EXIT_FAILURE);                                           \
+    }                                                               \
+  } while (0)
+
+ #define MPICHECK(cmd)                                                  \
+   do {                                                                 \
+     int e = cmd;                                                       \
+     if (e != MPI_SUCCESS) {                                            \
+       printf("Failed: MPI error %s:%d '%d'\n", __FILE__, __LINE__, e); \
+       exit(EXIT_FAILURE);                                              \
+     }                                                                  \
+   } while (0)
+
+ #define NCCLCHECK(cmd)                                              \
+   do {                                                              \
+     mcclResult_t r = cmd;                                           \
+     if (r != mcclSuccess) {                                         \
+       printf("Failed, NCCL error %s:%d '%s'\n", __FILE__, __LINE__, \
+              mcclGetErrorString(r));                                \
+       exit(EXIT_FAILURE);                                           \
+     }                                                               \
+   } while (0)
+
+ #ifdef USE_ROCM
+ __global__ void dummy_kernel() {
+   for (int i = 0; i < 100; i++) {
+     uint64_t start = wall_clock64();
+     uint64_t cycles_elapsed;
+     do {
+       cycles_elapsed = wall_clock64() - start;
+     } while (cycles_elapsed < 100);
+   }
+   for (int i = 0; i < 100; i++) __nanosleep(1000000);  // 100ms
+ }
+ #else
+ __global__ void dummy_kernel() {
+   #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
+   for (int i = 0; i < 100; i++) __nanosleep(1000000);  // 100ms
+   #else
+   for (int i = 0; i < 100; i++) {
+     long long int start = clock64();
+     while (clock64() - start < 150000000);  // approximately 98.4ms on P40
+   }
+   #endif
+ }
+ #endif
+
+ template <typename T>
+ __global__ void set_data(T* data, int size, int myRank) {
+   for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+        idx += gridDim.x * blockDim.x) {
+     data[idx] = myRank * 0.11f;
+   }
+ }
+
+ template <typename T>
+ __global__ void convert_data(const T* data1, const T* data2, double* fdata1,
+                              double* fdata2, int size) {
+   for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+        idx += gridDim.x * blockDim.x) {
+     fdata1[idx] = data1[idx];
+     fdata2[idx] = data2[idx];
+   }
+ }
+
+ __global__ void init_rand(murandState_t* state, int size, int nRanks) {
+   for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+        idx += gridDim.x * blockDim.x) {
+     for (int i = 0; i < nRanks; i++) {
+       murand_init(i + 1, idx, 0, &state[idx * nRanks + i]);
+     }
+   }
+ }
+
+ template <typename T>
+ __global__ void gen_data(murandState_t* state, T* data, double* ground_truth,
+                          int myRank, int nRanks, int size) {
+   for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < size;
+        idx += gridDim.x * blockDim.x) {
+     double sum = 0.0;
+     for (int i = 0; i < nRanks; i++) {
+       double val = murand_uniform_double(&state[idx * nRanks + i]) * 4;
+       T hval = val;  // downcast first
+       sum += static_cast<double>(hval);
+       if (i == myRank) data[idx] = hval;
+     }
+     ground_truth[idx] = sum;
+   }
+ }
+
+ template <typename T>
+ void run(int myRank, int nRanks, mcclComm_t& comm, int threads, int block_limit,
+          int data_size, bool performance_test) {
+   T* result;
+   musaStream_t stream;
+   MUSACHECK(musaStreamCreateWithFlags(&stream, musaStreamNonBlocking));
+   MUSACHECK(musaMalloc(&result, data_size * sizeof(T)));
+   MUSACHECK(musaMemset(result, 0, data_size * sizeof(T)));
+
+   musaIpcMemHandle_t self_data_handle;
+   musaIpcMemHandle_t data_handles[8];
+   vllm_musa::Signal* buffer;
+   T* self_data_copy;
+   /**
+    * Allocate IPC buffer
+    *
+    * The first section is a temporary buffer for storing intermediate allreduce
+    * results, if a particular algorithm requires it. The second section is for
+    * the input to the allreduce. The actual API takes the input pointer as an
+    * argument (that is, they can and usually should be allocated separately).
+    * But since the input pointers and the temporary buffer all require IPC
+    * registration, they are allocated and registered together in the test for
+    * convenience.
+    */
+ #ifdef USE_ROCM
+   MUSACHECK(hipExtMallocWithFlags(
+       (void**)&buffer, 2 * data_size * sizeof(T) + sizeof(vllm_musa::Signal),
+       hipDeviceMallocUncached));
+ #else
+   MUSACHECK(
+       musaMalloc(&buffer, 2 * data_size * sizeof(T) + sizeof(vllm_musa::Signal)));
+ #endif
+   MUSACHECK(
+       musaMemset(buffer, 0, 2 * data_size * sizeof(T) + sizeof(vllm_musa::Signal)));
+   MUSACHECK(musaMalloc(&self_data_copy, data_size * sizeof(T)));
+   MUSACHECK(musaIpcGetMemHandle(&self_data_handle, buffer));
+
+   MPICHECK(MPI_Allgather(&self_data_handle, sizeof(musaIpcMemHandle_t),
+                          MPI_BYTE, data_handles, sizeof(musaIpcMemHandle_t),
+                          MPI_BYTE, MPI_COMM_WORLD));
+
+   void* rank_data;
+   size_t rank_data_sz = 16 * 1024 * 1024;
+   MUSACHECK(musaMalloc(&rank_data, rank_data_sz));
+   vllm_musa::Signal* ipc_ptrs[8];
+   for (int i = 0; i < nRanks; i++) {
+     if (i == myRank)
+       ipc_ptrs[i] = buffer;
+     else
+       MUSACHECK(musaIpcOpenMemHandle((void**)&ipc_ptrs[i], data_handles[i],
+                                      musaIpcMemLazyEnablePeerAccess));
+   }
+   vllm_musa::CustomAllreduce fa(ipc_ptrs, rank_data, rank_data_sz, myRank, nRanks);
+   auto* self_data =
+       reinterpret_cast<T*>(reinterpret_cast<char*>(buffer) +
+                            sizeof(vllm_musa::Signal) + data_size * sizeof(T));
+   // hack buffer registration
+   {
+     void* data[8];
+     for (int i = 0; i < nRanks; i++) {
+       data[i] =
+           ((char*)ipc_ptrs[i]) + sizeof(vllm_musa::Signal) + data_size * sizeof(T);
+     }
+     fa.register_buffer(data);
+   }
+
+   double* ground_truth;
+   MUSACHECK(musaMallocHost(&ground_truth, data_size * sizeof(double)));
+   murandState_t* states;
+   MUSACHECK(musaMalloc(&states, sizeof(murandState_t) * nRanks * data_size));
+   init_rand<<<108, 1024, 0, stream>>>(states, data_size, nRanks);
+   gen_data<T><<<108, 1024, 0, stream>>>(states, self_data, ground_truth, myRank,
+                                         nRanks, data_size);
+   MUSACHECK(musaMemcpyAsync(self_data_copy, self_data, data_size * sizeof(T),
+                             musaMemcpyDeviceToDevice, stream));
+   musaEvent_t start, stop;
+   MUSACHECK(musaEventCreate(&start));
+   MUSACHECK(musaEventCreate(&stop));
+
+   mcclDataType_t mcclDtype;
+   if (std::is_same<T, half>::value) {
+     mcclDtype = mcclFloat16;
+   } else if (std::is_same<T, mt_bfloat16>::value) {
+     mcclDtype = mcclBfloat16;
+   } else {
+     mcclDtype = mcclFloat;
+   }
+   double *mccl_result, *my_result;
+   MUSACHECK(musaMallocHost(&mccl_result, data_size * sizeof(double)));
+   MUSACHECK(musaMallocHost(&my_result, data_size * sizeof(double)));
+   if (performance_test) {
+     dummy_kernel<<<1, 1, 0, stream>>>();
+     constexpr int warmup_iters = 5;
+     constexpr int num_iters = 100;
+     // warmup
+     for (int i = 0; i < warmup_iters; i++) {
+       NCCLCHECK(mcclAllReduce(result, result, data_size, mcclDtype, mcclSum,
+                               comm, stream));
+     }
+     MUSACHECK(musaEventRecord(start, stream));
+     for (int i = 0; i < num_iters; i++) {
+       NCCLCHECK(mcclAllReduce(result, result, data_size, mcclDtype, mcclSum,
+                               comm, stream));
+     }
+     MUSACHECK(musaEventRecord(stop, stream));
+     MUSACHECK(musaStreamSynchronize(stream));
+     float allreduce_ms = 0;
+     musaEventElapsedTime(&allreduce_ms, start, stop);
+
+     dummy_kernel<<<1, 1, 0, stream>>>();
+     // warm up
+     for (int i = 0; i < warmup_iters; i++) {
+       fa.allreduce<T>(stream, self_data, result, data_size, threads,
+                       block_limit);
+     }
+     MUSACHECK(musaEventRecord(start, stream));
+     for (int i = 0; i < num_iters; i++) {
+       fa.allreduce<T>(stream, self_data, result, data_size, threads,
+                       block_limit);
+     }
+     MUSACHECK(musaEventRecord(stop, stream));
+     MUSACHECK(musaStreamSynchronize(stream));
+
+     float duration_ms = 0;
+     musaEventElapsedTime(&duration_ms, start, stop);
+     if (myRank == 0)
+       printf(
+           "Rank %d done, nGPUs:%d, sz (kb): %d, %d, %d, my time:%.2fus, mccl "
+           "time:%.2fus\n",
+           myRank, nRanks, data_size * sizeof(T) / 1024, threads, block_limit,
+           duration_ms * 1e3 / num_iters, allreduce_ms * 1e3 / num_iters);
+
+     // And wait for all the queued up work to complete
+     MUSACHECK(musaStreamSynchronize(stream));
+
+     NCCLCHECK(mcclAllReduce(self_data_copy, self_data, data_size, mcclDtype,
+                             mcclSum, comm, stream));
+
+     convert_data<T><<<108, 1024, 0, stream>>>(self_data, result, mccl_result,
+                                               my_result, data_size);
+     MUSACHECK(musaStreamSynchronize(stream));
+
+     for (unsigned long j = 0; j < data_size; j++) {
+       auto diff = abs(mccl_result[j] - my_result[j]);
+       if (diff >= 4e-2) {
+         printf("Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n",
+                myRank, j, mccl_result[j], my_result[j], ground_truth[j]);
+         break;
+       }
+     }
+     long double mccl_diffs = 0.0;
+     long double my_diffs = 0.0;
+     for (int j = 0; j < data_size; j++) {
+       mccl_diffs += abs(mccl_result[j] - ground_truth[j]);
+       my_diffs += abs(my_result[j] - ground_truth[j]);
+     }
+     if (myRank == 0)
+       std::cout << "average abs diffs: mccl: " << mccl_diffs / data_size
+                 << " me: " << my_diffs / data_size << std::endl;
+   } else {
+     for (int i = 0; i < 100; i++) {
+       fa.allreduce<T>(stream, self_data, result, data_size, threads,
+                       block_limit);
+       MUSACHECK(musaStreamSynchronize(stream));
+       NCCLCHECK(mcclAllReduce(self_data, self_data_copy, data_size, mcclDtype,
+                               mcclSum, comm, stream));
+       convert_data<T><<<108, 1024, 0, stream>>>(
+           self_data_copy, result, mccl_result, my_result, data_size);
+       MUSACHECK(musaStreamSynchronize(stream));
+
+       for (unsigned long j = 0; j < data_size; j++) {
+         auto diff = abs(mccl_result[j] - my_result[j]);
+         if (diff >= 4e-2) {
+           printf(
+               "Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n",
+               myRank, j, mccl_result[j], my_result[j], ground_truth[j]);
+           break;
+         }
+       }
+     }
+     if (myRank == 0)
+       printf("Test passed: nGPUs:%d, sz (kb): %d, %d, %d\n", nRanks,
+              data_size * sizeof(T) / 1024, threads, block_limit);
+     // long double mccl_diffs = 0.0;
+     // long double my_diffs = 0.0;
+     // for (int j = 0; j < data_size; j++) {
+     //   mccl_diffs += abs(mccl_result[j] - ground_truth[j]);
+     //   my_diffs += abs(my_result[j] - ground_truth[j]);
+     // }
+     // if (myRank == 0)
+     //   std::cout << "average abs diffs: mccl: " << mccl_diffs / data_size
+     //             << " me: " << my_diffs / data_size << std::endl;
+   }
+
+   MUSACHECK(musaFree(result));
+   MUSACHECK(musaFree(self_data_copy));
+   MUSACHECK(musaFree(rank_data));
+   MUSACHECK(musaFree(buffer));
+   MUSACHECK(musaFree(states));
+   MUSACHECK(musaFreeHost(ground_truth));
+   MUSACHECK(musaFreeHost(mccl_result));
+   MUSACHECK(musaFreeHost(my_result));
+   MUSACHECK(musaStreamDestroy(stream));
+ }
+
+ int main(int argc, char** argv) {
+   int nRanks, myRank;
+   MPICHECK(MPI_Init(&argc, &argv));
+   MPICHECK(MPI_Comm_rank(MPI_COMM_WORLD, &myRank));
+   MPICHECK(MPI_Comm_size(MPI_COMM_WORLD, &nRanks));
+   MUSACHECK(musaSetDevice(myRank));
+   mcclUniqueId id;
+   mcclComm_t comm;
+   if (myRank == 0) mcclGetUniqueId(&id);
+   MPICHECK(MPI_Bcast(static_cast<void*>(&id), sizeof(id), MPI_BYTE, 0,
+                      MPI_COMM_WORLD));
+   NCCLCHECK(mcclCommInitRank(&comm, nRanks, id, myRank));
+
+   bool performance_test = true;
+   musaProfilerStart();
+   // Uncomment to scan through different block size configs.
+   for (int threads : {256, 512, 1024}) {
+      for (int block_limit = 16; block_limit < 112; block_limit += 4) {
+        run<half>(myRank, nRanks, comm, threads, block_limit, 1024 * 1024,
+        performance_test);
+      }
+   }
+   const int block_limit = 16;
+   // Scan through different sizes to test performance.
+   for (int sz = 512; sz <= (8 << 20); sz *= 2) {
+     run<half>(myRank, nRanks, comm, 512, 36, sz + 8 * 47, performance_test);
+   }
+
+   musaProfilerStop();
+   MPICHECK(MPI_Finalize());
+   return EXIT_SUCCESS;
+ }
\ No newline at end of file
diff --git a/vllm_musa/csrc/decode_mla/decode_mla.mu b/vllm_musa/csrc/decode_mla/decode_mla.mu
new file mode 100644
index 000000000..4c1e7d446
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/decode_mla.mu
@@ -0,0 +1,384 @@
+#include <torch/all.h>
+#include <torch/torch.h>
+#include <torch_musa/csrc/core/MUSAGuard.h>
+#include <iostream>
+#include "../type_convert.muh"
+#include "../dispatch_utils.h"
+
+#include "./mp22_impl/decode_mla_kernel.hpp"
+#include "./mp31_impl/decode_mla.hpp"
+
+template<typename T>
+struct dependents_false : std::false_type {};
+
+// Macro for MP22 decode MLA kernel launcher
+#define LAUNCH_MP22_DECODE_MLA(TYPE) \
+    mp22_decode_mla_kernel_launcher<TYPE>( \
+        static_cast<TYPE*>(Q.data_ptr()), \
+        static_cast<TYPE*>(K_Buffer.data_ptr()), \
+        static_cast<TYPE*>(V_Buffer.data_ptr()), \
+        static_cast<float*>(Att_out.data_ptr()), \
+        static_cast<int32_t*>(Req_to_tokens.data_ptr()), \
+        static_cast<int32_t*>(B_Seqlen.data_ptr()), \
+        sm_scale, \
+        num_splits_kv, \
+        num_blocks, \
+        1, 1,   /* q=1, kv=1 */ \
+        d_qk, d_vo, num_heads, num_head_kv, batch, \
+        Req_to_tokens.stride(0), \
+        Q.stride(0), Q.stride(1), \
+        K_Buffer.stride(-3), K_Buffer.stride(-2), \
+        V_Buffer.stride(-3), V_Buffer.stride(-2), \
+        Att_out.stride(0), Att_out.stride(1), Att_out.stride(2) \
+    )
+
+// Macro for MP31 decode MLA kernel launcher
+#define LAUNCH_MP31_DECODE_MLA(TYPE, ATT_TYPE, HEAD_DIM_QK, HEAD_DIM_VO, K_STAGES, V_STAGES, SPLITS_KV) \
+    mp31_decode_mla_kernel_launcher<TYPE, ATT_TYPE, HEAD_DIM_QK, HEAD_DIM_VO, K_STAGES, V_STAGES>( \
+        static_cast<TYPE*>(Q.data_ptr()), \
+        static_cast<TYPE*>(K_Buffer.data_ptr()), \
+        static_cast<TYPE*>(V_Buffer.data_ptr()), \
+        static_cast<ATT_TYPE*>(Att_out.data_ptr()), \
+        static_cast<int32_t*>(Req_to_tokens.data_ptr()), \
+        static_cast<int32_t*>(B_Seqlen.data_ptr()), \
+        sm_scale, \
+        num_splits_kv, \
+        seqlen_q_ori, 1, \
+        d_qk, d_vo, num_heads, num_head_kv, batch, num_blocks, \
+        Req_to_tokens.stride(0), \
+        Q.stride(0), Q.stride(1), Q.stride(2), \
+        K_Buffer.stride(0), K_Buffer.stride(1), K_Buffer.stride(2), \
+        V_Buffer.stride(0), V_Buffer.stride(1), V_Buffer.stride(2), \
+        Att_out.stride(3), Att_out.stride(0), Att_out.stride(1), Att_out.stride(2), \
+        causal \
+    )
+
+// Macro for MP31 decode MHA kernel launcher
+#define LAUNCH_MP31_DECODE_MHA(TYPE, ATT_TYPE, HEAD_DIM_QK, HEAD_DIM_VO, K_STAGES, V_STAGES, SPLITS_KV) \
+    mp31_decode_mla_kernel_launcher<TYPE, ATT_TYPE, HEAD_DIM_QK, HEAD_DIM_VO, K_STAGES, V_STAGES>( \
+        static_cast<TYPE*>(Q.data_ptr()), \
+        static_cast<TYPE*>(K_Buffer.data_ptr()), \
+        static_cast<TYPE*>(V_Buffer.data_ptr()), \
+        static_cast<ATT_TYPE*>(Att_out.data_ptr()), \
+        static_cast<int32_t*>(Req_to_tokens.data_ptr()), \
+        static_cast<int32_t*>(B_Seqlen.data_ptr()), \
+        sm_scale, \
+        num_splits_kv, \
+        1, 1,   /* q=1, kv=1 */ \
+        d_qk, d_vo, num_head, num_head_kv, batch, num_blocks, \
+        Req_to_tokens.stride(0), \
+        Q.stride(0), Q.stride(2), Q.stride(1), \
+        K_Buffer.stride(0), K_Buffer.stride(1), K_Buffer.stride(2), \
+        V_Buffer.stride(0), V_Buffer.stride(1), V_Buffer.stride(2), \
+        Att_out.stride(3), Att_out.stride(0), Att_out.stride(2), Att_out.stride(1), \
+        causal \
+    )
+
+template <typename T>
+void mp22_decode_mla_kernel_launcher(
+     T* Q, const T* K_Buffer, const T* V_Buffer, float* Att_out,
+     const int32_t* Req_to_tokens, const int32_t* B_Seqlen, float sm_scale, int num_splits_kv, int num_pages,
+    int q, int kv, int d_qk, int d_vo, int h, int h_kv, int b,
+    int stride_req_to_tokens_b,
+    int stride_qbs, int stride_qh,
+    int stride_buf_kbs, int stride_buf_kh,
+    int stride_buf_vbs, int stride_buf_vh,
+    int stride_mid_ob, int stride_mid_oh, int stride_mid_os
+) {
+    constexpr int TileM = 32;
+    constexpr int TileN = 32;
+    constexpr int HeadDimQK = 576;
+    constexpr int HeadDimVO = 512;
+    constexpr int PageSize = 16;
+    const musaStream_t stream = at::musa::getCurrentMUSAStream();
+    using KernelTraits = DecodeMLAKernelTraits<HeadDimQK, HeadDimVO, TileM, TileN, PageSize,
+                                                T, true, false>;
+    int h_r = h / h_kv;
+    auto problem_shape = mute::make_shape(mute::_1{}, PageSize, d_qk, d_vo, mute::make_shape(h_r, h_kv), b);
+    using DecodeMLAParams = DecodeMLAParams<KernelTraits>;
+    using Arguments = typename DecodeMLAParams::Arguments;
+    Arguments arguments = {
+        Q,
+        K_Buffer,
+        V_Buffer,
+        Req_to_tokens,
+        B_Seqlen,
+        Att_out,
+        sm_scale,
+        num_splits_kv,
+        num_pages,
+        stride_req_to_tokens_b,
+        stride_qbs,
+        stride_qh,
+        stride_buf_kbs,
+        stride_buf_kh,
+        stride_buf_vbs,
+        stride_buf_vh,
+        stride_mid_ob,
+        stride_mid_oh,
+        stride_mid_os
+    };
+
+    auto params = DecodeMLAParams::to_underlying_arguments(problem_shape, arguments);
+    dim3 block_dims(KernelTraits::NumThreads);
+    dim3 grid_dims = {static_cast<uint32_t>(b), static_cast<uint32_t>(ceil_div(h, min(KernelTraits::CTA_Q, h_r))), static_cast<uint32_t>(num_splits_kv)};
+    decode_mla_kernel<KernelTraits><<<grid_dims, block_dims, 0, stream>>>(params);
+}
+
+template <typename T, typename Att_T, int HeadDimQK, int HeadDimVO, int KStages, int VStages>
+void mp31_decode_mla_kernel_launcher(
+     T* Q, const T* K_Buffer, const T* V_Buffer, Att_T* Att_out,
+    const int32_t* Req_to_tokens, const int32_t* B_Seqlen, float sm_scale, int num_splits_kv,
+    int q, int kv, int d_qk, int d_vo, int h, int h_kv, int b, int num_blocks,
+    int stride_req_to_tokens_b,
+    int stride_qbs, int stride_q_seq, int stride_qh,
+    int stride_buf_kblock, int stride_buf_kpage, int stride_buf_kh,
+    int stride_buf_vblock, int stride_buf_vpage,int stride_buf_vh,
+    int stride_mid_o_split, int stride_mid_ob, int stride_mid_os, int stride_mid_oh,
+    bool causal
+) {
+    constexpr int BlockM = 32;
+    constexpr int BlockN = 64;
+    constexpr int PageSize = BlockN;
+    const musaStream_t stream = at::musa::getCurrentMUSAStream();
+
+    // PackGQA
+    int h_r = h / h_kv;
+    int q_seq_per_hkv = q * h_r;
+
+    auto problem_shape = mute::make_shape(q, PageSize, d_qk, d_vo, mute::make_shape(h_r, h_kv), b);
+    if (causal) {
+        using KernelTraits = FmhaDecodeKernelTratis<BlockM, BlockN, HeadDimQK, HeadDimVO, PageSize, T, Att_T, KStages, VStages, true>;
+        typename KernelTraits::PackGQAArguments arguments {
+            Q,
+            K_Buffer,
+            V_Buffer,
+            Req_to_tokens,
+            B_Seqlen,
+            Att_out,
+            sm_scale,
+            num_splits_kv,
+            num_blocks,
+            stride_req_to_tokens_b,
+            stride_qh,
+            stride_q_seq,
+            stride_qbs,
+            stride_buf_kblock,
+            stride_buf_kpage,
+            stride_buf_kh,
+            stride_buf_vblock,
+            stride_buf_vpage,
+            stride_buf_vh,
+            stride_mid_oh,
+            stride_mid_os,
+            stride_mid_ob,
+            stride_mid_o_split
+        };
+
+        auto params = KernelTraits::to_underlying_arguments(problem_shape, arguments);
+        int num_m_block = mute::ceil_div(q_seq_per_hkv, BlockM);
+        int batch_and_split = b * num_splits_kv;
+        dim3 grid = {static_cast<uint32_t>(num_m_block), static_cast<uint32_t>(h_kv), static_cast<uint32_t>(batch_and_split)};
+        mla_decode<KernelTraits><<<grid, KernelTraits::NumThreads, 0, stream>>>(params);
+    }
+    else {
+        using KernelTraits = FmhaDecodeKernelTratis<BlockM, BlockN, HeadDimQK, HeadDimVO, PageSize, T, Att_T, KStages, VStages, false>;
+        typename KernelTraits::PackGQAArguments arguments {
+            Q,
+            K_Buffer,
+            V_Buffer,
+            Req_to_tokens,
+            B_Seqlen,
+            Att_out,
+            sm_scale,
+            num_splits_kv,
+            num_blocks,
+            stride_req_to_tokens_b,
+            stride_qh,
+            stride_q_seq,
+            stride_qbs,
+            stride_buf_kblock,
+            stride_buf_kpage,
+            stride_buf_kh,
+            stride_buf_vblock,
+            stride_buf_vpage,
+            stride_buf_vh,
+            stride_mid_oh,
+            stride_mid_os,
+            stride_mid_ob,
+            stride_mid_o_split
+        };
+
+        auto params = KernelTraits::to_underlying_arguments(problem_shape, arguments);
+        int num_m_block = mute::ceil_div(q_seq_per_hkv, BlockM);
+        int batch_and_split = b * num_splits_kv;
+        dim3 grid = {static_cast<uint32_t>(num_m_block), static_cast<uint32_t>(h_kv), static_cast<uint32_t>(batch_and_split)};
+        mla_decode<KernelTraits><<<grid, KernelTraits::NumThreads, 0, stream>>>(params);
+    }
+}
+
+void decode_mla(
+    torch::Tensor &Q,
+    torch::Tensor &K_Buffer,
+    torch::Tensor &V_Buffer,
+    float sm_scale,
+    int num_splits_kv,
+    torch::Tensor &Req_to_tokens,
+    torch::Tensor &B_Seqlen,
+    torch::Tensor &Att_out
+) {
+    const int batch = B_Seqlen.size(0);
+    const int num_heads = Q.size(1);
+    const int d_qk = Q.size(2);
+    const int num_blocks = K_Buffer.size(0);
+    const int page_size = K_Buffer.size(1);
+    const int num_head_kv = K_Buffer.size(-2);
+    const int d_vo = V_Buffer.size(3);
+
+    constexpr int HeadDimQK = 576;
+    constexpr int HeadDimVO = 512;
+    constexpr int KStages = 1;
+    constexpr int VStages = 1;
+
+    TORCH_CHECK(d_qk == HeadDimQK, "mp31 decode mla d_qk only support ", HeadDimQK);
+    TORCH_CHECK(d_vo == HeadDimVO, "mp31 decode mla d_vo only support ", HeadDimVO);
+
+    const at::musa::OptionalMUSAGuard device_guard(device_of(Q));
+
+    if(at::musa::getMUSAArch() == 220) {
+        TORCH_CHECK(page_size == 16, "mp22 decode mla BLOCK_SIZE only support 16.");
+        if (Q.scalar_type() == at::ScalarType::BFloat16) {
+            LAUNCH_MP22_DECODE_MLA(mutlass::bfloat16_t);
+        } else if (Q.scalar_type() == at::ScalarType::Half) {
+            LAUNCH_MP22_DECODE_MLA(mutlass::half_t);
+        } else {
+            std::cerr << "Decode MLA not impl for this dtype." << std::endl;
+            assert(false);
+        }
+    } else if (at::musa::getMUSAArch() == 310) {
+        TORCH_CHECK(page_size == 64, "mp31 decode mla BLOCK_SIZE only support 64.");
+
+        const int seqlen_q_ori = Q.size(0) / batch;
+        const int num_q_heads_per_hk = num_heads / num_head_kv;
+        const int q_seq_per_hk = seqlen_q_ori * num_q_heads_per_hk;
+        const bool causal = true;
+
+        Q = Q.view({batch, q_seq_per_hk, num_head_kv, d_qk});
+        if (num_splits_kv == 1) {
+            Att_out = Att_out.view({batch, q_seq_per_hk, num_head_kv, num_splits_kv, d_vo});
+        } else {
+            Att_out = Att_out.view({batch, q_seq_per_hk, num_head_kv, num_splits_kv, d_vo + 1});
+        }
+
+        if (Q.scalar_type() == at::ScalarType::BFloat16) {
+            if (num_splits_kv == 1) {
+                LAUNCH_MP31_DECODE_MLA(mutlass::bfloat16_t, mutlass::bfloat16_t, HeadDimQK, HeadDimVO, KStages, VStages, 1);
+            } else {
+                LAUNCH_MP31_DECODE_MLA(mutlass::bfloat16_t, float, HeadDimQK, HeadDimVO, KStages, VStages, 0);
+            }
+        } else if (Q.scalar_type() == at::ScalarType::Half) {
+            if (num_splits_kv == 1) {
+                LAUNCH_MP31_DECODE_MLA(mutlass::half_t, mutlass::half_t, HeadDimQK, HeadDimVO, KStages, VStages, 1);
+            } else {
+                LAUNCH_MP31_DECODE_MLA(mutlass::half_t, float, HeadDimQK, HeadDimVO, KStages, VStages, 0);
+            }
+        } else {
+            std::cerr << "Decode MLA not impl for this dtype." << std::endl;
+            assert(false);
+        }
+    } else {
+        std::cerr << "Decode MLA only support for mp22 & mp31." << std::endl;
+        assert(false);
+    }
+}
+
+void mp31_decode_mha(
+    torch::Tensor &Q,
+    torch::Tensor &K_Buffer,
+    torch::Tensor &V_Buffer,
+    float sm_scale,
+    int num_splits_kv,
+    torch::Tensor &Req_to_tokens,
+    torch::Tensor &B_Seqlen,
+    torch::Tensor &Att_out
+) {
+    int batch = Q.size(0);
+    int num_head = Q.size(1);
+    int num_head_kv = K_Buffer.size(-2);
+    int d_qk = Q.size(2);
+    int d_vo = K_Buffer.size(3);
+    int num_blocks = K_Buffer.size(0);
+    int page_size = K_Buffer.size(1);
+    int head_size_v = V_Buffer.size(-1);
+
+    const int seqlen_q_ori = 1;
+    const int num_q_heads_per_hk = num_head / num_head_kv;
+    const int q_seq_per_hk = seqlen_q_ori * num_q_heads_per_hk;
+    const bool causal = false;
+
+    Q = Q.view({batch, num_head_kv, num_q_heads_per_hk, d_qk});
+
+    constexpr int KStages = 2;
+    constexpr int VStages = 2;
+
+    TORCH_CHECK((d_qk == 64 && d_vo == 64) || (d_qk == 128 && d_vo == 128),
+               "mp31 decode mla d_qk and d_vo only support 64 or 128.");
+
+    TORCH_CHECK(page_size == 64, "mp31 decode mla BLOCK_SIZE only support 64.");
+
+
+    int current_arch = at::musa::getMUSAArch();
+    TORCH_CHECK(current_arch >= 300, "mp31 decode mha not support on MUSA arch ", current_arch);
+
+    if (num_splits_kv == 1) {
+        Att_out = Att_out.view({batch, num_head_kv, q_seq_per_hk, num_splits_kv, head_size_v});
+    } else {
+        Att_out = Att_out.view({batch, num_head_kv, q_seq_per_hk, num_splits_kv, head_size_v + 1});
+    }
+
+    if (Q.scalar_type() == at::ScalarType::BFloat16) {
+        if (num_splits_kv == 1) {
+            if (d_qk == 64 && d_vo == 64) {
+                LAUNCH_MP31_DECODE_MHA(mutlass::bfloat16_t, mutlass::bfloat16_t, 64, 64, KStages, VStages, 1);
+            } else if (d_qk == 128 && d_vo == 128) {
+                LAUNCH_MP31_DECODE_MHA(mutlass::bfloat16_t, mutlass::bfloat16_t, 128, 128, KStages, VStages, 1);
+            } else {
+                TORCH_CHECK(false, "mp31 decode mla d_qk and d_vo only support 64 or 128.");
+            }
+        } else {
+            if (d_qk == 64 && d_vo == 64) {
+                LAUNCH_MP31_DECODE_MHA(mutlass::bfloat16_t, float, 64, 64, KStages, VStages, 0);
+            } else if (d_qk == 128 && d_vo == 128) {
+                LAUNCH_MP31_DECODE_MHA(mutlass::bfloat16_t, float, 128, 128, KStages, VStages, 0);
+            } else {
+                TORCH_CHECK(false, "mp31 decode mla d_qk and d_vo only support 64 or 128.");
+            }
+        }
+    } else if (Q.scalar_type() == at::ScalarType::Half) {
+        if (num_splits_kv == 1) {
+            if (d_qk == 64 && d_vo == 64) {
+                LAUNCH_MP31_DECODE_MHA(mutlass::half_t, mutlass::half_t, 64, 64, KStages, VStages, 1);
+            } else if (d_qk == 128 && d_vo == 128) {
+                LAUNCH_MP31_DECODE_MHA(mutlass::half_t, mutlass::half_t, 128, 128, KStages, VStages, 1);
+            } else {
+                TORCH_CHECK(false, "mp31 decode mla d_qk and d_vo only support 64 or 128.");
+            }
+        } else {
+            if (d_qk == 64 && d_vo == 64) {
+                LAUNCH_MP31_DECODE_MHA(mutlass::half_t, float, 64, 64, KStages, VStages, 0);
+            } else if (d_qk == 128 && d_vo == 128) {
+                LAUNCH_MP31_DECODE_MHA(mutlass::half_t, float, 128, 128, KStages, VStages, 0);
+            } else {
+                TORCH_CHECK(false, "mp31 decode mla d_qk and d_vo only support 64 or 128.");
+            }
+        }
+    } else {
+        std::cerr << "Decode MLA not impl for this dtype." << std::endl;
+    }
+
+    if (num_splits_kv == 1) {
+        Att_out = Att_out.reshape({batch, num_head, num_splits_kv, head_size_v});
+    } else {
+        Att_out = Att_out.reshape({batch, num_head, num_splits_kv, head_size_v + 1});
+    }
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_kernel.hpp b/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_kernel.hpp
new file mode 100644
index 000000000..28812afc7
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_kernel.hpp
@@ -0,0 +1,795 @@
+#pragma once
+
+#include <mutlass/numeric_conversion.h>
+
+#include "decode_mla_params.hpp"
+#include "utils.hpp"
+#include "./online_softmax.hpp"
+
+template <class KernelTraits>
+__global__ __launch_bounds__(KernelTraits::NumThreads, 1)
+void decode_mla_kernel(const typename DecodeMLAParams<KernelTraits>::Params params) {
+  #if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 220)
+  using ElementAccum = typename KernelTraits::ElementAccumulator;
+  constexpr int FragmentSize = KernelTraits::FragmentSize;
+
+  // Get them from traits
+  const int KVGroupNum = get<4,0>(params.problem_shape);
+  const int QHeadNum = size<4>(params.problem_shape);
+  const int NUM_KV_SPLITS = params.NumKVSplits;
+
+  constexpr int CTA_Q = KernelTraits::CTA_Q;
+  constexpr int Alignment = KernelTraits::Alignment;
+
+  const int ValidBlockH = KVGroupNum > CTA_Q ? CTA_Q : KVGroupNum;
+
+  __shared__ char __attribute__((aligned(128))) smem[sizeof(typename KernelTraits::SharedStorage)];
+  auto& shared_storage = *reinterpret_cast<typename KernelTraits::SharedStorage*>(smem);
+
+  auto [cur_batch, cur_head_id, split_kv_id] = static_cast<uint3>(blockIdx);
+
+  int cur_kv_head = cur_head_id / ceil_div(KVGroupNum, CTA_Q);
+
+  int thread_idx = threadIdx.x;
+
+  int cur_batch_seq_len = params.ptr_B_Seqlen[cur_batch];
+  int cur_batch_req_idx = cur_batch;
+
+  int kv_len_per_split = ceil_div(cur_batch_seq_len, NUM_KV_SPLITS);
+  int split_kv_start = kv_len_per_split * split_kv_id;
+  int split_kv_end = min(split_kv_start + kv_len_per_split, cur_batch_seq_len);
+
+  //if (thread_idx == 0) {
+  //      printf("bid:(%d, %d, %d) kv_st:%d kv_end:%d, kv_len_per_split:%d cur_seq_len:%d\n", blockIdx.x, blockIdx.y, blockIdx.z, split_kv_start, split_kv_end, kv_len_per_split, cur_batch_seq_len);
+  //}
+
+  // early return
+  if (split_kv_end <= split_kv_start) return;
+
+  const int* block_table = params.ptr_Req_to_tokens + cur_batch_req_idx * params.stride_req_to_tokens_b;
+
+  // TensorQ ((H_R, H_KV), D_QK, B)
+  auto shape_Q = make_shape(QHeadNum, get<2>(params.problem_shape), get<5>(params.problem_shape));
+  auto stride_Q = make_stride(params.stride_qh, _1{}, params.stride_qbs);
+  using TileShapeQ = conditional_t<KernelTraits::SplitQKSmem,
+                                  Shape<Int<KernelTraits::CTA_Q>, Int<512>>,
+                                  Shape<Int<KernelTraits::CTA_Q>, Int<576>>>;
+
+  auto tile_shape_Q = TileShapeQ{};
+
+  Tensor mQ = make_tensor(make_gmem_ptr(params.ptr_Q), shape_Q, stride_Q);
+  Tensor gQ = local_tile(mQ, tile_shape_Q, make_coord(cur_head_id, _0{}, cur_batch));
+  Tensor sQ  = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), typename KernelTraits::SmemLayoutQ{});
+
+  typename KernelTraits::GmemTiledCopyQ gmem_tiled_copy_Q;
+  ThrCopy gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(thread_idx);
+  Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); // (QCPY, QCPY_M, QCPY_N)
+  Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); // (QCPY, QCPY_M, QCPY_N)
+
+  Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ)));
+  Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ);
+  Tensor tQpQ = make_tensor<bool>(make_shape(size<1>(tQsQ), size<2>(tQsQ)), Stride<_1,_0>{}); // let's only check Head
+
+  constexpr int KVPageNum = KernelTraits::CTA_KV / KernelTraits::PAGE_SIZE;
+  static_assert(KernelTraits::PAGE_SIZE == 16);
+  // TensorK (bs, D_QK, (H_R, H_KV))
+  auto tile_shape_K = select<1, 2>(typename KernelTraits::TileShape_QKD{});
+  Tensor sK  = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), typename KernelTraits::SmemLayoutK{});
+
+  typename KernelTraits::GmemTiledCopyK gmem_tiled_copy_K;
+  ThrCopy gmem_thr_copy_K = gmem_tiled_copy_K.get_thread_slice(thread_idx);
+  Tensor tKsK = gmem_thr_copy_K.partition_D(sK); // (KCPY, KCPY_M, KCPY_N)
+
+  Tensor cK = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK)));
+  Tensor tKcK = gmem_thr_copy_K.partition_S(cK);
+  Tensor tKpK = make_tensor<bool>(make_shape(size<1>(tKcK), size<2>(tKcK)), Stride<_1, _0>{}); // let's only check Seq
+
+  // TensorV (bs, D_VO, (H_R, H_KV))
+  auto tile_shape_V = select<2, 1>(typename KernelTraits::TileShape_PDV{});
+  Tensor sV  = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), typename KernelTraits::SmemLayoutV{});
+  Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), typename KernelTraits::SmemLayoutVt{});
+
+  // Reuse TileCopyK to get same tensor partition...
+  typename KernelTraits::GmemTiledCopyK gmem_tiled_copy_V;
+  ThrCopy gmem_thr_copy_V = gmem_tiled_copy_V.get_thread_slice(thread_idx);
+  Tensor tVsV = gmem_thr_copy_V.partition_D(sV); // (VCPY, VCPY_M, VCPY_N)
+
+  Tensor cV = make_identity_tensor(make_shape(size<0>(sV), size<1>(sV)));
+  Tensor tVcV = gmem_thr_copy_V.partition_S(cV);
+  Tensor tCpV = make_tensor<bool>(make_shape(size<1>(tVcV), size<2>(tVcV)), Stride<_1, _0>{});
+
+  // Store P
+  Tensor sP  = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), typename KernelTraits::SmemLayoutP{});
+
+  // MMA
+  typename KernelTraits::TiledMmaQK tiled_mma_qk;
+  typename KernelTraits::TiledMmaPV tiled_mma_pv;
+
+  auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx);
+  auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx);
+
+  // QK MMA AB
+  Tensor tPrQ = thr_mma_qk.partition_fragment_A(sQ); // (MMA, MMA_M, MMA_K)
+  Tensor tPrK = thr_mma_qk.partition_fragment_B(sK); // (MMA, MMA_M, MMA_K)
+
+  // PV MMA AB
+  Tensor tOrP  = thr_mma_pv.partition_fragment_A(sP);
+  Tensor tOrVt = thr_mma_pv.partition_fragment_B(sVt);
+
+  auto tile_shape_P = select<0, 1>(typename KernelTraits::TileShape_QKD{});
+  auto tile_shape_O = select<0, 1>(typename KernelTraits::TileShape_PDV{});
+  // Allocate accum
+  Tensor accum_p = partition_fragment_C(tiled_mma_qk, tile_shape_P);
+  Tensor accum_o = partition_fragment_C(tiled_mma_pv, tile_shape_O);
+
+  // Allocate softmax
+  Tensor accum_softmax = make_fragment_like(accum_p);
+  Tensor accum_cvt     = make_fragment_like<typename KernelTraits::ElementKV>(accum_softmax);
+
+  Tensor tAcc_frg = recast<mutlass::Array<ElementAccum, FragmentSize>>(accum_softmax);
+  Tensor tCvt_frg = recast<mutlass::Array<typename KernelTraits::ElementKV, FragmentSize>>(accum_cvt);
+
+  // STS P
+  auto tiled_copy_r2s = make_tiled_copy_C(typename KernelTraits::R2SCopyAtomP{}, tiled_mma_qk);
+  auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(thread_idx);
+  Tensor tPrP = thr_copy_r2s.retile_S(accum_cvt);
+  Tensor tPsP = thr_copy_r2s.partition_D(sP);
+
+  // LDS QKPV
+  auto thr_copy_Q = make_tiled_copy_A(typename KernelTraits::SmemCopyAtomQ{}, tiled_mma_qk).get_thread_slice(thread_idx);
+  auto thr_copy_K = make_tiled_copy_B(typename KernelTraits::SmemCopyAtomK{}, tiled_mma_qk).get_thread_slice(thread_idx);
+  auto thr_copy_P = make_tiled_copy_A(typename KernelTraits::SmemCopyAtomP{}, tiled_mma_pv).get_thread_slice(thread_idx);
+  auto thr_copy_V = make_tiled_copy_B(typename KernelTraits::SmemCopyAtomV{}, tiled_mma_pv).get_thread_slice(thread_idx);
+
+  // QK OOB check
+  Tensor cP = make_identity_tensor(tile_shape_P);
+  Tensor tPcP = thr_mma_qk.partition_C(cP);
+
+  // QK GEMM
+  Tensor tPsQ           = thr_copy_Q.partition_S(sQ);
+  Tensor tPrQ_copy_view = thr_copy_Q.retile_D(tPrQ);
+  MUTE_STATIC_ASSERT_V(size<1>(tPsQ) == size<1>(tPrQ_copy_view));
+
+  Tensor tPsK           = thr_copy_K.partition_S(sK);
+  Tensor tPrK_copy_view = thr_copy_K.retile_D(tPrK);
+  MUTE_STATIC_ASSERT_V(size<1>(tPsK) == size<1>(tPrK_copy_view));
+
+  // Softmax
+  constexpr int RowsPerThread = size<0>(tile_shape_Q) / tile_size<0>(tiled_mma_qk) * 4;
+  Tensor alpha       = make_tensor<ElementAccum>(Shape<Int<RowsPerThread>>{});
+  Tensor lse         = make_fragment_like(alpha);
+  Tensor row_sum     = make_fragment_like(alpha);
+  Tensor row_max     = make_fragment_like(alpha);
+  Tensor old_row_max = make_fragment_like(alpha);
+
+  constexpr bool DoSoftmax = true;
+
+  if constexpr (DoSoftmax) {
+	fill(lse,         ElementAccum(0.0f));
+	fill(alpha,       ElementAccum(0.0f));
+	fill(old_row_max, -mutlass::platform::numeric_limits<ElementAccum>::infinity());
+  } else {
+	fill(lse, ElementAccum(1.0f));
+  }
+
+  // PV GEMM
+  Tensor tOsP           = thr_copy_P.partition_S(sP);
+  Tensor tOrP_copy_view = thr_copy_P.retile_D(tOrP);
+
+  Tensor tOsVt           = thr_copy_V.partition_S(sVt);
+  Tensor tOrV_copy_view  = thr_copy_V.retile_D(tOrVt);
+
+  clear(accum_p);
+  clear(accum_o);
+
+  // Output Tensor (Heads, D_VO + 1, SPLITS, B)
+
+  auto shape_O = make_shape(QHeadNum, get<3>(params.problem_shape), NUM_KV_SPLITS, get<5>(params.problem_shape));
+  auto stride_O = make_stride(params.stride_mid_oh, _1{}, params.stride_mid_os, params.stride_mid_ob);
+  Tensor mO = make_tensor(make_gmem_ptr(params.ptr_O), shape_O, stride_O);
+  Tensor gO = local_tile(mO, tile_shape_O, make_coord(cur_head_id, _0{}, split_kv_id, cur_batch));
+
+  Tensor tOgO = thr_mma_pv.partition_C(gO);
+  Tensor cO   = make_identity_tensor(tile_shape_O);
+  Tensor tOcO = thr_mma_pv.partition_C(cO);
+
+
+  //
+  // PREDICATES
+  //
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < size<0>(tQpQ); ++i) {
+    int cur_head = cur_head_id * ValidBlockH + get<0>(tQcQ(0,i,0));
+    tQpQ(i, 0) = (cur_head < (cur_head_id + 1) * ValidBlockH) && (cur_head < QHeadNum);
+  }
+
+  //if (blockIdx.z == 1) {
+  //      //printf("bidx:%d batch:%d block table value:%d\n", blockIdx.x, cur_batch_req_idx, block_table[0]);
+  //      printf("bidx:%d start_n:%d offset_n:%d, off:%d\n", blockIdx.x, split_kv_start, split_kv_start + get<0>(tKcK(0,0,0)), (split_kv_start + get<0>(tKcK(0,0,0))) / KernelTraits::PAGE_SIZE);
+  //}
+
+
+  //if (blockIdx.z == 0) {
+  //      printf("tid:%d tKcK major:%d %d\n", threadIdx.x, get<1>(tKcK(0,0,0)), get<1>(tKcK(0,0,1)));
+  //}
+  if (thread0()) {
+    //print("mQ:"); print(mQ); print("\n");
+    //print("gQ:"); print(gQ); print("\n");
+    //print("sK:"); print(sK); print("\n");
+    //print("tKsK:"); print(tKsK); print("\n");
+    //print("tKpK:"); print(tKpK); print("\n");
+    //print("tQgQ:"); print(tQgQ); print("\n");
+    //print("tQsQ:"); print(tQsQ); print("\n");
+    //print("tQpQ:"); print(tQpQ); print("\n");
+    //print("tKgK:"); print(tKgK_); print("\n");
+    //print("tKsK':"); print(tKsK(_,0,_)); print("\n");
+    //print("accum_p:"); print(accum_p); print("\n");
+
+    //print("mO:"); print(mO); print("\n");
+    //print("gO:"); print(gO); print("\n");
+  }
+
+#if 1
+  // LDGSTS copy Q
+  copy_if(gmem_tiled_copy_Q.with(params.desc_Q), tQpQ, tQgQ, tQsQ);
+
+  Tensor sQPE = make_tensor(make_smem_ptr(shared_storage.smem_q_pe.data()), typename KernelTraits::SmemLayoutQPE{});
+  Tensor sKPE = make_tensor(make_smem_ptr(shared_storage.smem_k_pe.data()), typename KernelTraits::SmemLayoutKPE{});
+
+  if constexpr (KernelTraits::SplitQKSmem) {
+    // lgsts qpe
+    auto tile_shape_QPE = make_shape(get<0>(typename KernelTraits::TileShape_QKD{}), Int<64>{});
+    Tensor gQPE = local_tile(mQ, tile_shape_QPE, make_coord(cur_head_id, _8{}, cur_batch));
+
+    Tensor tQgQPE = gmem_thr_copy_Q.partition_S(gQPE);
+    Tensor tQsQPE = gmem_thr_copy_Q.partition_D(sQPE);
+
+    copy_if(gmem_tiled_copy_Q.with(params.desc_Q), tQpQ, tQgQPE, tQsQPE);
+  }
+
+  __syncthreads();
+
+  // if we keep Q in register && reuse Q, then lds Q
+  if constexpr (KernelTraits::IsQInReg) {
+    copy(tPsQ, tPrQ_copy_view);
+
+    // Reuse SmemQ
+    __syncthreads();
+  }
+
+#if 0
+  for (int start_n = split_kv_start; start_n < split_kv_end; start_n += KernelTraits::CTA_KV) {
+    clear(accum_p);
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size<0>(tKpK); ++i) {
+      tKpK(i, 0) = start_n + get<0>(tKcK(0,i,0)) < split_kv_end;
+    }
+
+    int kv_loc[KVPageNum];
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      int offset_n = start_n + get<0>(tKcK(0,i,0));
+      int kv_page_number = tKpK(i, 0) ? block_table[offset_n / KernelTraits::PAGE_SIZE] : 0;
+      //kv_loc[i] = kv_page_number * KernelTraits::PAGE_SIZE + start_n % KernelTraits::PAGE_SIZE;
+      kv_loc[i] = kv_page_number * KernelTraits::PAGE_SIZE + offset_n % KernelTraits::PAGE_SIZE;
+      //if (blockIdx.z == 1) {
+      //if (tKpK(i,0)) {
+      //        printf("bid:%d tid:%d #:%d off_n:%d, kv_page:%d kv_loc:%d\n", blockIdx.z, threadIdx.x, i, offset_n, kv_page_number, kv_loc[i]);
+      //}
+      //}
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      using TileShapeK = conditional_t<KernelTraits::SplitQKSmem,
+                                       Shape<Int<KernelTraits::PAGE_SIZE>, Int<512>>,
+                                       Shape<Int<KernelTraits::PAGE_SIZE>, Int<576>>>;
+      auto tile_shape_K = TileShapeK{};
+      Tensor gK = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + cur_kv_head * params.stride_buf_kh),
+                              tile_shape_K,
+                              make_stride(params.stride_buf_kbs, _1{}));
+      Tensor tKgK = gmem_thr_copy_K.partition_S(gK);
+
+      // tKgK already offset tKcK(0,i,0)
+      tKgK.data() = tKgK.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_kbs;
+      copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgK(_,_0{},_), tKsK(_, i, _));
+
+
+
+
+      //MUTLASS_PRAGMA_UNROLL
+      //for (int j = 0; j < 8; ++j) {
+      //	//tKgK.data() = params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + get<1>(tKcK(0,0,j));
+      //  Tensor tKgK_ = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + get<1>(tKcK(0,0,j))), tKgK.layout());
+      //	//if (thread0()) {
+      //	//        ///printf("m,n(%d %d): base:%p, cur:%p\n", i, j, params.ptr_K_Buffer, tKgK.data());
+      //  //	print("tKgK:"); print(tKgK_); print("\n");
+      //	//}
+      //  copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgK_(_,_0{},j), tKsK(_, i, j));
+      //}
+    }
+
+    __syncthreads();
+
+    //if (blockIdx.z == 1) {
+    //        printf("k:%f %f %f %f\n", float(shared_storage.smem_k[0]),float(shared_storage.smem_k[1]),float(shared_storage.smem_k[2]),float(shared_storage.smem_k[3]));
+    //}
+
+    // P = QK
+    gemm<KernelTraits::IsQInReg>(accum_p, tPrQ, tPrK, tPsQ, tPsK, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+
+
+    if constexpr (KernelTraits::SplitQKSmem) {
+      // make sure we already finish reading smem_k
+      __syncthreads();
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < KVPageNum; ++i) {
+        // ldgsts kpe
+        auto tile_shape_KPE = make_shape(Int<KernelTraits::PAGE_SIZE>{}, Int<64>{});
+        Tensor gKPE = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + cur_kv_head * params.stride_buf_kh + Int<512>{}),
+                                  tile_shape_KPE,
+                                  make_stride(params.stride_buf_kbs, _1{}));
+
+        Tensor tKgKPE = gmem_thr_copy_K.partition_S(gKPE);
+        Tensor tKsKPE = gmem_thr_copy_K.partition_D(sKPE);
+
+	tKgKPE.data() = tKgKPE.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_kbs;
+        copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgKPE(_,_0{},_), tKsKPE(_,i,_));
+
+	//tKgKPE.data() = params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + Int<512>{} + get<1>(tKcK(0,0,0));
+	//Tensor tKgKPE_ = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + Int<512>{} + get<1>(tKcK(0,0,0))), tKgKPE.layout());
+        //copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgKPE_(_,_0{},_), tKsKPE(_,i,_));
+      }
+      // construct qpe/kpe
+      Tensor tPrQPE = thr_mma_qk.partition_fragment_A(sQPE);
+      Tensor tPrKPE = thr_mma_qk.partition_fragment_B(sKPE);
+      Tensor tPsQPE = thr_copy_Q.partition_S(sQPE);
+      Tensor tPsKPE = thr_copy_K.partition_S(sKPE);
+      // finish lgsts
+      __syncthreads();
+      gemm<KernelTraits::IsQInReg>(accum_p, tPrQPE, tPrKPE, tPsQPE, tPsKPE, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+    }
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_scale(accum_p, accum_softmax, params.sm_scale);
+    } else {
+	mp22_online_softmax_scale(accum_p, accum_softmax, float(1.0f));
+    }
+
+
+    //if constexpr (KernelTraits::logit_cap > 0) {
+    //  // TODO
+    //}
+
+    mp22_process_qk_oob<!DoSoftmax, true>(accum_softmax, tPcP, cur_head_id * ValidBlockH, ValidBlockH, QHeadNum, start_n, split_kv_end);
+
+    //if (blockIdx.z == 0) {
+    //    printf("t:%d softmax:%f %f %f %f\n", thread_idx, accum_softmax(0), accum_softmax(1), accum_softmax(2), accum_softmax(3));
+    //}
+
+    // if we keep Q in smem, then KV shared the smem, sync to make sure QK is done
+    if constexpr (!KernelTraits::IsQInReg) {
+      __syncthreads();
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      Tensor gV = make_tensor(make_gmem_ptr(params.ptr_V_Buffer + cur_kv_head * params.stride_buf_vh),
+                              make_shape(Int<KernelTraits::PAGE_SIZE>{}, get<1>(typename KernelTraits::TileShape_PDV{})),
+                              make_stride(params.stride_buf_vbs, _1{}));
+      Tensor tVgV = gmem_thr_copy_V.partition_S(gV);
+      tVgV.data() = tVgV.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_vbs;
+      copy_if(gmem_tiled_copy_V.with(params.desc_V), tKpK(i,_), tVgV(_, _0{}, _), tVsV(_, i, _));
+      //MUTLASS_PRAGMA_UNROLL
+      //for (int j = 0; j < 8; ++j) {
+      //	//tVgV.data() = params.ptr_V_Buffer + kv_loc[i] * params.stride_buf_vbs + cur_kv_head * params.stride_buf_vh + get<1>(tKcK(0,0,j));
+      //  Tensor tVgV_ = make_tensor(make_gmem_ptr(params.ptr_V_Buffer + kv_loc[i] * params.stride_buf_vbs + cur_kv_head * params.stride_buf_vh + get<1>(tKcK(0,0,j))), tVgV.layout());
+      //	copy_if(gmem_tiled_copy_V.with(params.desc_V), tKpK(i,_), tVgV_(_, _0{}, j), tVsV(_, i, j));
+      //}
+    }
+
+    __syncthreads();
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_row_max<ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+	//if (blockIdx.z == 0) {
+	//        printf("tid:%d acc_softmax:%f %f %f %f\n",  thread_idx, accum_softmax(0), accum_softmax(1), accum_softmax(2), accum_softmax(3));
+	//	//if (thread_idx == 1) printf("soft:%f exp2:%f\n", accum_softmax(0), exp2f(accum_softmax(0)));
+	//	//printf("tid:%d acc:%f %f %f %f\n", thread_idx, accum_o(0), accum_o(1), accum_o(2), accum_o(3));
+	//}
+
+	online_softmax_exp2(accum_softmax);
+	mp22_update_accum(accum_o, alpha);
+    }
+    MUTE_UNROLL
+    for (int i = 0; i < size(tCvt_frg); ++i) {
+      tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::ElementKV, ElementAccum, FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+    }
+
+    //if (blockIdx.z == 0) {
+    //        printf("t:%d cvt:%f %f %f %f\n", thread_idx, float(accum_cvt(0)), float(accum_cvt(1)), float(accum_cvt(2)), float(accum_cvt(3)));
+    //}
+
+    copy(tiled_copy_r2s, tPrP, tPsP);
+
+    __syncthreads();
+
+    //if (thread0()) {
+    //    printf("vbuf:%f %f pbuf:%f %f\n", float(shared_storage.smem_v[0]), float(shared_storage.smem_v[1]), float(shared_storage.smem_p[0]), float(shared_storage.smem_p[1]));
+    //}
+
+    gemm(accum_o, tOrP, tOrVt, tOsP, tOsVt, tiled_mma_pv, thr_copy_P, thr_copy_V);
+
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_row_sum<ElementAccum>(accum_softmax, row_sum, lse, alpha);
+    }
+  }
+
+#else
+  int start_n = split_kv_start;
+
+  for (; start_n < split_kv_end - KernelTraits::CTA_KV; start_n += KernelTraits::CTA_KV) {
+    clear(accum_p);
+
+    int kv_loc[KVPageNum];
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      int offset_n = start_n + get<0>(tKcK(0,i,0));
+      int kv_page_number = block_table[offset_n / KernelTraits::PAGE_SIZE];
+      kv_loc[i] = kv_page_number * KernelTraits::PAGE_SIZE + offset_n % KernelTraits::PAGE_SIZE;
+      //if (blockIdx.z == 1) {
+      //if (tKpK(i,0)) {
+      //        printf("bid:%d tid:%d #:%d off_n:%d, kv_page:%d kv_loc:%d\n", blockIdx.z, threadIdx.x, i, offset_n, kv_page_number, kv_loc[i]);
+      //}
+      //}
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      using TileShapeK = conditional_t<KernelTraits::SplitQKSmem,
+                                       Shape<Int<KernelTraits::PAGE_SIZE>, Int<512>>,
+                                       Shape<Int<KernelTraits::PAGE_SIZE>, Int<576>>>;
+      auto tile_shape_K = TileShapeK{};
+      Tensor gK = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + cur_kv_head * params.stride_buf_kh),
+                              tile_shape_K,
+                              make_stride(params.stride_buf_kbs, _1{}));
+      Tensor tKgK = gmem_thr_copy_K.partition_S(gK);
+
+      // tKgK already offset tKcK(0,i,0)
+      tKgK.data() = tKgK.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_kbs;
+      copy(gmem_tiled_copy_K.with(params.desc_K), tKgK(_,_0{},_), tKsK(_, i, _));
+    }
+
+    __syncthreads();
+
+    //if (blockIdx.z == 1) {
+    //        printf("k:%f %f %f %f\n", float(shared_storage.smem_k[0]),float(shared_storage.smem_k[1]),float(shared_storage.smem_k[2]),float(shared_storage.smem_k[3]));
+    //}
+
+    // P = QK
+    gemm<KernelTraits::IsQInReg>(accum_p, tPrQ, tPrK, tPsQ, tPsK, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+
+
+    if constexpr (KernelTraits::SplitQKSmem) {
+      // make sure we already finish reading smem_k
+      __syncthreads();
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < KVPageNum; ++i) {
+        // ldgsts kpe
+        auto tile_shape_KPE = make_shape(Int<KernelTraits::PAGE_SIZE>{}, Int<64>{});
+        Tensor gKPE = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + cur_kv_head * params.stride_buf_kh + Int<512>{}),
+                                  tile_shape_KPE,
+                                  make_stride(params.stride_buf_kbs, _1{}));
+
+        Tensor tKgKPE = gmem_thr_copy_K.partition_S(gKPE);
+        Tensor tKsKPE = gmem_thr_copy_K.partition_D(sKPE);
+
+	tKgKPE.data() = tKgKPE.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_kbs;
+        copy(gmem_tiled_copy_K.with(params.desc_K), tKgKPE(_,_0{},_), tKsKPE(_,i,_));
+
+      //if (blockIdx.z == 0) {
+      //  printf("tid:%d i:%d kv_loc:%d base:%p ", thread_idx, i, kv_loc[i], params.ptr_K_Buffer + cur_kv_head * params.stride_buf_kh);
+      //  print("tKgKPE:");print(tKgKPE.data()); print("\n");
+      //}
+
+	//tKgKPE.data() = params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + Int<512>{} + get<1>(tKcK(0,0,0));
+	//Tensor tKgKPE_ = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + Int<512>{} + get<1>(tKcK(0,0,0))), tKgKPE.layout());
+        //copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgKPE_(_,_0{},_), tKsKPE(_,i,_));
+      }
+      // construct qpe/kpe
+      Tensor tPrQPE = thr_mma_qk.partition_fragment_A(sQPE);
+      Tensor tPrKPE = thr_mma_qk.partition_fragment_B(sKPE);
+      Tensor tPsQPE = thr_copy_Q.partition_S(sQPE);
+      Tensor tPsKPE = thr_copy_K.partition_S(sKPE);
+      // finish lgsts
+      __syncthreads();
+      gemm<KernelTraits::IsQInReg>(accum_p, tPrQPE, tPrKPE, tPsQPE, tPsKPE, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+    }
+
+
+    //if (blockIdx.z == 0) {
+    //    printf("t:%d qk:%f %f %f %f\n", thread_idx, accum_p(0), accum_p(1), accum_p(2), accum_p(3));
+    //}
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_scale(accum_p, accum_softmax, params.sm_scale);
+    } else {
+	mp22_online_softmax_scale(accum_p, accum_softmax, float(1.0f));
+    }
+
+
+    //if constexpr (KernelTraits::logit_cap > 0) {
+    //  // TODO
+    //}
+
+    mp22_process_qk_oob<!DoSoftmax, false>(accum_softmax, tPcP, cur_head_id * ValidBlockH, ValidBlockH, QHeadNum, start_n, split_kv_end);
+
+    //if (blockIdx.z == 0) {
+    //    printf("t:%d softmax:%f %f %f %f\n", thread_idx, accum_softmax(0), accum_softmax(1), accum_softmax(2), accum_softmax(3));
+    //}
+
+    // if we keep Q in smem, then KV shared the smem, sync to make sure QK is done
+    if constexpr (!KernelTraits::IsQInReg) {
+      __syncthreads();
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      Tensor gV = make_tensor(make_gmem_ptr(params.ptr_V_Buffer + cur_kv_head * params.stride_buf_vh),
+                              make_shape(Int<KernelTraits::PAGE_SIZE>{}, get<1>(typename KernelTraits::TileShape_PDV{})),
+                              make_stride(params.stride_buf_vbs, _1{}));
+      Tensor tVgV = gmem_thr_copy_V.partition_S(gV);
+      tVgV.data() = tVgV.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_vbs;
+      copy(gmem_tiled_copy_V.with(params.desc_V), tVgV(_, _0{}, _), tVsV(_, i, _));
+    }
+
+    __syncthreads();
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_row_max<ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+	//if (blockIdx.z == 0) {
+	//        printf("tid:%d acc_softmax:%f %f %f %f\n",  thread_idx, accum_softmax(0), accum_softmax(1), accum_softmax(2), accum_softmax(3));
+	//	//if (thread_idx == 1) printf("soft:%f exp2:%f\n", accum_softmax(0), exp2f(accum_softmax(0)));
+	//	//printf("tid:%d acc:%f %f %f %f\n", thread_idx, accum_o(0), accum_o(1), accum_o(2), accum_o(3));
+	//}
+
+	online_softmax_exp2(accum_softmax);
+	mp22_update_accum(accum_o, alpha);
+    }
+    MUTE_UNROLL
+    for (int i = 0; i < size(tCvt_frg); ++i) {
+      tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::ElementKV, ElementAccum, FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+    }
+
+    //if (blockIdx.z == 0) {
+    //        printf("t:%d cvt:%f %f %f %f\n", thread_idx, float(accum_cvt(0)), float(accum_cvt(1)), float(accum_cvt(2)), float(accum_cvt(3)));
+    //}
+
+    copy(tiled_copy_r2s, tPrP, tPsP);
+
+    __syncthreads();
+
+    gemm(accum_o, tOrP, tOrVt, tOsP, tOsVt, tiled_mma_pv, thr_copy_P, thr_copy_V);
+
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_row_sum<ElementAccum>(accum_softmax, row_sum, lse, alpha);
+    }
+
+  } // end hotloop
+
+  // hot loop free smem
+  __syncthreads();
+  {
+    clear(accum_p);
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size<0>(tKpK); ++i) {
+      tKpK(i, 0) = start_n + get<0>(tKcK(0,i,0)) < split_kv_end;
+    }
+
+    int kv_loc[KVPageNum];
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      int offset_n = start_n + get<0>(tKcK(0,i,0));
+      int kv_page_number = tKpK(i, 0) ? block_table[offset_n / KernelTraits::PAGE_SIZE] : 0;
+      //kv_loc[i] = kv_page_number * KernelTraits::PAGE_SIZE + start_n % KernelTraits::PAGE_SIZE;
+      kv_loc[i] = kv_page_number * KernelTraits::PAGE_SIZE + offset_n % KernelTraits::PAGE_SIZE;
+      //if (blockIdx.z == 1) {
+      //if (tKpK(i,0)) {
+      //        printf("bid:%d tid:%d #:%d off_n:%d, kv_page:%d kv_loc:%d\n", blockIdx.z, threadIdx.x, i, offset_n, kv_page_number, kv_loc[i]);
+      //}
+      //}
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      using TileShapeK = conditional_t<KernelTraits::SplitQKSmem,
+                                       Shape<Int<KernelTraits::PAGE_SIZE>, Int<512>>,
+                                       Shape<Int<KernelTraits::PAGE_SIZE>, Int<576>>>;
+      auto tile_shape_K = TileShapeK{};
+      Tensor gK = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + cur_kv_head * params.stride_buf_kh),
+                              tile_shape_K,
+                              make_stride(params.stride_buf_kbs, _1{}));
+      Tensor tKgK = gmem_thr_copy_K.partition_S(gK);
+
+      // tKgK already offset tKcK(0,i,0)
+      tKgK.data() = tKgK.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_kbs;
+      copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgK(_,_0{},_), tKsK(_, i, _));
+
+
+
+
+      //MUTLASS_PRAGMA_UNROLL
+      //for (int j = 0; j < 8; ++j) {
+      //	//tKgK.data() = params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + get<1>(tKcK(0,0,j));
+      //  Tensor tKgK_ = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + get<1>(tKcK(0,0,j))), tKgK.layout());
+      //	//if (thread0()) {
+      //	//        ///printf("m,n(%d %d): base:%p, cur:%p\n", i, j, params.ptr_K_Buffer, tKgK.data());
+      //  //	print("tKgK:"); print(tKgK_); print("\n");
+      //	//}
+      //  copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgK_(_,_0{},j), tKsK(_, i, j));
+      //}
+    }
+
+    __syncthreads();
+
+    //if (blockIdx.z == 1) {
+    //        printf("k:%f %f %f %f\n", float(shared_storage.smem_k[0]),float(shared_storage.smem_k[1]),float(shared_storage.smem_k[2]),float(shared_storage.smem_k[3]));
+    //}
+
+    // P = QK
+    gemm<KernelTraits::IsQInReg>(accum_p, tPrQ, tPrK, tPsQ, tPsK, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+
+
+    if constexpr (KernelTraits::SplitQKSmem) {
+      // make sure we already finish reading smem_k
+      __syncthreads();
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < KVPageNum; ++i) {
+        // ldgsts kpe
+        auto tile_shape_KPE = make_shape(Int<KernelTraits::PAGE_SIZE>{}, Int<64>{});
+        Tensor gKPE = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + cur_kv_head * params.stride_buf_kh + Int<512>{}),
+                                  tile_shape_KPE,
+                                  make_stride(params.stride_buf_kbs, _1{}));
+
+        Tensor tKgKPE = gmem_thr_copy_K.partition_S(gKPE);
+        Tensor tKsKPE = gmem_thr_copy_K.partition_D(sKPE);
+
+	tKgKPE.data() = tKgKPE.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_kbs;
+        copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgKPE(_,_0{},_), tKsKPE(_,i,_));
+
+      //if (blockIdx.z == 0) {
+      //  printf("tid:%d i:%d kv_loc:%d base:%p ", thread_idx, i, kv_loc[i], params.ptr_K_Buffer + cur_kv_head * params.stride_buf_kh);
+      //  print("tKgKPE:");print(tKgKPE.data()); print("\n");
+      //}
+
+	//tKgKPE.data() = params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + Int<512>{} + get<1>(tKcK(0,0,0));
+	//Tensor tKgKPE_ = make_tensor(make_gmem_ptr(params.ptr_K_Buffer + kv_loc[i] * params.stride_buf_kbs + cur_kv_head * params.stride_buf_kh + Int<512>{} + get<1>(tKcK(0,0,0))), tKgKPE.layout());
+        //copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK(i, _), tKgKPE_(_,_0{},_), tKsKPE(_,i,_));
+      }
+      // construct qpe/kpe
+      Tensor tPrQPE = thr_mma_qk.partition_fragment_A(sQPE);
+      Tensor tPrKPE = thr_mma_qk.partition_fragment_B(sKPE);
+      Tensor tPsQPE = thr_copy_Q.partition_S(sQPE);
+      Tensor tPsKPE = thr_copy_K.partition_S(sKPE);
+      // finish lgsts
+      __syncthreads();
+      gemm<KernelTraits::IsQInReg>(accum_p, tPrQPE, tPrKPE, tPsQPE, tPsKPE, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+    }
+
+
+    //if (blockIdx.z == 0) {
+    //    printf("t:%d qk:%f %f %f %f\n", thread_idx, accum_p(0), accum_p(1), accum_p(2), accum_p(3));
+    //}
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_scale(accum_p, accum_softmax, params.sm_scale);
+    } else {
+	mp22_online_softmax_scale(accum_p, accum_softmax, float(1.0f));
+    }
+
+
+    //if constexpr (KernelTraits::logit_cap > 0) {
+    //  // TODO
+    //}
+
+    mp22_process_qk_oob<!DoSoftmax, true>(accum_softmax, tPcP, cur_head_id * ValidBlockH, ValidBlockH, QHeadNum, start_n, split_kv_end);
+
+    //if (blockIdx.z == 0) {
+    //    printf("t:%d softmax:%f %f %f %f\n", thread_idx, accum_softmax(0), accum_softmax(1), accum_softmax(2), accum_softmax(3));
+    //}
+
+    // if we keep Q in smem, then KV shared the smem, sync to make sure QK is done
+    if constexpr (!KernelTraits::IsQInReg) {
+      __syncthreads();
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < KVPageNum; ++i) {
+      Tensor gV = make_tensor(make_gmem_ptr(params.ptr_V_Buffer + cur_kv_head * params.stride_buf_vh),
+                              make_shape(Int<KernelTraits::PAGE_SIZE>{}, get<1>(typename KernelTraits::TileShape_PDV{})),
+                              make_stride(params.stride_buf_vbs, _1{}));
+      Tensor tVgV = gmem_thr_copy_V.partition_S(gV);
+      tVgV.data() = tVgV.data() + (-get<0>(tKcK(0,0,0)) + kv_loc[i]) * params.stride_buf_vbs;
+      copy_if(gmem_tiled_copy_V.with(params.desc_V), tKpK(i,_), tVgV(_, _0{}, _), tVsV(_, i, _));
+      //MUTLASS_PRAGMA_UNROLL
+      //for (int j = 0; j < 8; ++j) {
+      //	//tVgV.data() = params.ptr_V_Buffer + kv_loc[i] * params.stride_buf_vbs + cur_kv_head * params.stride_buf_vh + get<1>(tKcK(0,0,j));
+      //  Tensor tVgV_ = make_tensor(make_gmem_ptr(params.ptr_V_Buffer + kv_loc[i] * params.stride_buf_vbs + cur_kv_head * params.stride_buf_vh + get<1>(tKcK(0,0,j))), tVgV.layout());
+      //	copy_if(gmem_tiled_copy_V.with(params.desc_V), tKpK(i,_), tVgV_(_, _0{}, j), tVsV(_, i, j));
+      //}
+    }
+
+    __syncthreads();
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_row_max<ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+	//if (blockIdx.z == 0) {
+	//        printf("tid:%d acc_softmax:%f %f %f %f\n",  thread_idx, accum_softmax(0), accum_softmax(1), accum_softmax(2), accum_softmax(3));
+	//	//if (thread_idx == 1) printf("soft:%f exp2:%f\n", accum_softmax(0), exp2f(accum_softmax(0)));
+	//	//printf("tid:%d acc:%f %f %f %f\n", thread_idx, accum_o(0), accum_o(1), accum_o(2), accum_o(3));
+	//}
+
+	online_softmax_exp2(accum_softmax);
+	mp22_update_accum(accum_o, alpha);
+    }
+    MUTE_UNROLL
+    for (int i = 0; i < size(tCvt_frg); ++i) {
+      tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::ElementKV, ElementAccum, FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+    }
+
+    //if (blockIdx.z == 0) {
+    //        printf("t:%d cvt:%f %f %f %f\n", thread_idx, float(accum_cvt(0)), float(accum_cvt(1)), float(accum_cvt(2)), float(accum_cvt(3)));
+    //}
+
+    copy(tiled_copy_r2s, tPrP, tPsP);
+
+    __syncthreads();
+
+    //if (thread0()) {
+    //    printf("vbuf:%f %f pbuf:%f %f\n", float(shared_storage.smem_v[0]), float(shared_storage.smem_v[1]), float(shared_storage.smem_p[0]), float(shared_storage.smem_p[1]));
+    //}
+
+    gemm(accum_o, tOrP, tOrVt, tOsP, tOsVt, tiled_mma_pv, thr_copy_P, thr_copy_V);
+
+
+    if constexpr (DoSoftmax) {
+	mp22_online_softmax_row_sum<ElementAccum>(accum_softmax, row_sum, lse, alpha);
+    }
+
+  }
+
+#endif
+
+  if constexpr (DoSoftmax) {
+    mp22_apply_softmax_normalizer(accum_o, lse);
+  }
+  // Store
+
+  /// STG Accum
+  Tensor accum_o_mn = make_tensor(accum_o.data(), mp22_convert_layout_acc_rowcol(accum_o.layout()));
+  Tensor tOcO_mn = make_tensor(tOcO.data(), mp22_convert_layout_acc_rowcol(tOcO.layout()));
+  Tensor tOgO_mn = make_tensor(tOgO.data(), mp22_convert_layout_acc_rowcol(tOgO.layout()));
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_o_mn); ++mi) {
+    const bool mask_h = get<0>(tOcO_mn(mi, 0)) < ValidBlockH && (cur_head_id * ValidBlockH + get<0>(tOcO_mn(mi, 0)) < QHeadNum);
+
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_o_mn); ++ni) {
+      if (mask_h) {
+        tOgO_mn(mi, ni) = accum_o_mn(mi,ni);
+      }
+    }
+
+    if (mask_h) {
+      params.ptr_O[cur_batch * params.stride_mid_ob + (cur_head_id * CTA_Q + get<0>(tOcO_mn(mi,0))) * params.stride_mid_oh + split_kv_id * params.stride_mid_os + 512] = row_max(mi) * params.rescale + std::logf(lse(mi));
+    }
+  }
+  #endif
+  #endif
+}
diff --git a/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_kernel_traits.hpp b/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_kernel_traits.hpp
new file mode 100644
index 000000000..d8e6b35bc
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_kernel_traits.hpp
@@ -0,0 +1,236 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mute/atom/mma_atom.hpp>
+#include <mute/atom/copy_atom.hpp>
+
+#include <mutlass/mutlass.h>
+#include <mutlass/gemm/gemm.h>
+#include <mutlass/gemm/collective/collective_builder.hpp>
+
+#include "mutlass_extension.hpp"
+
+using namespace mute;
+
+template <
+  class Element,
+  class SmemLayoutQ,
+  class SmemLayoutK,
+  class SmemLayoutV,
+  class SmemLayoutP
+>
+struct SharedStorageReuseQ {
+  union {
+    mute::array_aligned<Element,  mute::cosize_v<SmemLayoutQ>> smem_q;
+    struct {
+      mute::array_aligned<Element, mute::cosize_v<SmemLayoutK>> smem_k;
+      mute::array_aligned<Element, mute::cosize_v<SmemLayoutV>> smem_v;
+      mute::array_aligned<Element, mute::cosize_v<SmemLayoutP>> smem_p;
+    };
+    mute::array_aligned<Element, 1> smem_q_pe; // dummy
+    mute::array_aligned<Element, 1> smem_k_pe; // dummy
+  };
+};
+
+template <
+  class Element,
+  class SmemLayoutQ,
+  class SmemLayoutQPE,
+  class SmemLayoutK,
+  class SmemLayoutKPE,
+  class SmemLayoutV,
+  class SmemLayoutP
+>
+struct SharedStorageSplitHeadDim {
+  // Q use 36KB
+  struct {
+    mute::array_aligned<Element, mute::cosize_v<SmemLayoutQ>> smem_q; // q_model
+    mute::array_aligned<Element, mute::cosize_v<SmemLayoutQPE>> smem_q_pe;
+  };
+  // Reuse K_Model&K_PE && V --> 32KB
+  union {
+    union {
+      mute::array_aligned<Element, mute::cosize_v<SmemLayoutK>> smem_k;
+      mute::array_aligned<Element, mute::cosize_v<SmemLayoutKPE>> smem_k_pe;
+    };
+    mute::array_aligned<Element, mute::cosize_v<SmemLayoutV>> smem_v;
+  };
+  // P 2KB
+  mute::array_aligned<Element, mute::cosize_v<SmemLayoutP>> smem_p;
+};
+
+template <class TileN, class PermuteTileShape = TileN>
+constexpr auto
+get_mp22_permute_tile_for_operand_B() {
+  using MmaAtomN = _16;
+  static_assert(TileN{} % PermuteTileShape{} == 0);
+  static_assert(PermuteTileShape{} % MmaAtomN{} == 0);
+
+  using RepeatInsidePermuteTile = decltype(PermuteTileShape{} / MmaAtomN{});
+  using RepeatAcrossPermuteTile = decltype(TileN{} / PermuteTileShape{});
+
+  using PermuteTileN = Layout<Shape <               MmaAtomN, RepeatInsidePermuteTile, RepeatAcrossPermuteTile>,
+                              Stride<RepeatInsidePermuteTile,                      _1,        PermuteTileShape>>;
+
+  return PermuteTileN{};
+}
+
+template <
+  int HEAD_DIM_QK_,
+  int HEAD_DIM_VO_,
+  int CTA_Q_,
+  int CTA_KV_,
+  int PAGE_SIZE_,
+  class Element_,
+  bool SplitKV_ = true,
+  bool IsQInReg_ = true
+>
+struct DecodeMLAKernelTraits {
+  using Element = Element_;
+  using ElementQ = Element;
+  using ElementKV = Element;
+  using ElementAccumulator = float;
+  using ElementO = conditional_t<SplitKV_, float, Element>;
+
+  static constexpr int CTA_Q         = CTA_Q_;
+  static constexpr int CTA_KV        = CTA_KV_;
+  static constexpr int HEAD_DIM_QK   = HEAD_DIM_QK_;
+  static constexpr int HEAD_DIM_VO   = HEAD_DIM_VO_;
+  static constexpr int IsQInReg      = IsQInReg_;
+  static constexpr int SplitKV       = SplitKV_;
+  static constexpr int PAGE_SIZE     = PAGE_SIZE_;
+  static constexpr int SplitQKSmem   = !IsQInReg;
+
+  static_assert(SplitKV, "Only support SplitKV.");
+
+  static constexpr int FragmentSize = 2;
+
+  using TileShape_QKD = Shape<Int<CTA_Q>, Int<CTA_KV>, Int<HEAD_DIM_QK>>;
+  using TileShape_PDV = Shape<Int<CTA_Q>, Int<HEAD_DIM_VO>, Int<CTA_KV>>;
+
+  // assume (1,1,1) for decode mla
+  using AtomLayoutMNK = Layout<Shape<Int<CTA_Q/32>, _1, _1>>;
+
+  // Internal Mma Stride type
+  using MmaStrideMNMajor = Stride<_1, int>;
+  using MmaStrideKMajor  = Stride<int, _1>;
+
+  using PermutedTileN = decltype(get_mp22_permute_tile_for_operand_B<Int<CTA_KV>, Int<32>>());
+  using PermuteLayoutMNK = Tile<Underscore, PermutedTileN, Underscore>;
+
+  using TiledMmaQK = decltype(mute::make_tiled_mma(
+    mutlass::gemm::collective::detail::mp22_mma_operation_select<
+                                        ElementQ, MmaStrideKMajor,
+                                        ElementKV,MmaStrideKMajor>(),
+    AtomLayoutMNK{},
+    PermuteLayoutMNK{}));
+
+  using TiledMmaPV =decltype(mute::make_tiled_mma(
+    mutlass::gemm::collective::detail::mp22_mma_operation_select<
+                                        ElementKV, MmaStrideKMajor,
+                                        ElementKV, MmaStrideMNMajor>(),
+    AtomLayoutMNK{}));
+
+  static constexpr int NumThreads = size(TiledMmaQK{});
+  static_assert(size(TiledMmaQK{}) == size(TiledMmaPV{}));
+
+  using SmemLayoutAtomQ = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementQ,  MmaStrideKMajor>());
+  using SmemLayoutQ_ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0,2>(TileShape_QKD{})));
+
+  using SmemLayoutAtomK = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutK_ = decltype(tile_to_shape(SmemLayoutAtomK{}, select<1,2>(TileShape_QKD{})));
+
+  // Split QK 576 = 512 + 64
+  static_assert(HEAD_DIM_QK == 576, "Only support HEAD_DIM_QK=576");
+  using SmemLayoutQMODEL = decltype(tile_to_shape(SmemLayoutAtomQ{}, make_shape(get<0>(TileShape_QKD{}), Int<512>{})));
+  using SmemLayoutQPE    = decltype(tile_to_shape(SmemLayoutAtomQ{}, make_shape(get<0>(TileShape_QKD{}), Int<64>{})));
+
+  using SmemLayoutKMODEL = decltype(tile_to_shape(SmemLayoutAtomK{}, make_shape(get<1>(TileShape_QKD{}), Int<512>{})));
+  using SmemLayoutKPE    = decltype(tile_to_shape(SmemLayoutAtomK{}, make_shape(get<1>(TileShape_QKD{}), Int<64>{})));
+
+  using SmemLayoutQ = mute::conditional_t<SplitQKSmem, SmemLayoutQMODEL, SmemLayoutQ_>;
+  using SmemLayoutK = mute::conditional_t<SplitQKSmem, SmemLayoutKMODEL, SmemLayoutK_>;
+
+  using SmemLayoutAtomP = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutP = decltype(tile_to_shape(SmemLayoutAtomP{}, select<0,1>(TileShape_QKD{})));
+
+  using SmemLayoutAtomV = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutV  = decltype(tile_to_shape(SmemLayoutAtomV{}, select<2,1>(TileShape_PDV{})));
+  using SmemLayoutVt = decltype(composition(SmemLayoutV{},
+                        make_layout(select<1,2>(TileShape_PDV{}), LayoutRight{})));
+
+  using SmemCopyAtomQ = Copy_Atom<DefaultCopy, ElementQ>;
+  using SmemCopyAtomK = Copy_Atom<DefaultCopy, ElementKV>;
+  using SmemCopyAtomP = Copy_Atom<DefaultCopy, ElementKV>;
+  using SmemCopyAtomV = Copy_Atom<DefaultCopy, ElementKV>;
+
+  using R2SCopyAtomP  = Copy_Atom<UniversalCopy<mute::uint_byte_t<FragmentSize * sizeof(ElementKV)>>, ElementKV>;
+
+  // TileCopy
+#if 0
+  static constexpr int AlignmentQ = 128 / sizeof_bits_v<ElementQ>;
+  static constexpr int ThreadsMajorQ = mute::gcd(HEAD_DIM_QK / AlignmentQ, NumThreads);
+  static constexpr int ThreadsMinorQ = NumThreads / ThreadsMajorQ;
+  using GmemCopyQThrLayout = Layout<Shape <Int<ThreadsMinorQ>, Int<ThreadsMajorQ>>,
+                                    Stride<Int<ThreadsMajorQ>,                _1>>;
+  using GmemCopyQValLayout = Layout<Shape<_1, Int<AlignmentQ>>>;
+  using GmemCopyOpQ = MP22_ROBUST_LDGSTS<uint_bit_t<AlignmentQ * sizeof_bits_v<ElementQ>>>;
+
+  using GmemTiledCopyQ = decltype(make_tiled_copy(
+                            Copy_Atom<GmemCopyOpQ, ElementQ>{},
+                            GmemCopyQThrLayout{}, GmemCopyQValLayout{}));
+
+  static constexpr int AlignmentK = 128 / sizeof_bits_v<ElementKV>;
+  static constexpr int ThreadsMajorK = mute::gcd(HEAD_DIM_QK / AlignmentK, NumThreads);
+  static constexpr int ThreadsMinorK = NumThreads / ThreadsMajorK;
+  using GmemCopyKThrLayout = Layout<Shape <Int<ThreadsMinorK>, Int<ThreadsMajorK>>,
+                                    Stride<Int<ThreadsMajorK>,                _1>>;
+  using GmemCopyKValLayout = Layout<Shape<_1, Int<AlignmentK>>>;
+  using GmemCopyOpK = MP22_LDGSTS<uint_bit_t<AlignmentK * sizeof_bits_v<ElementKV>>>;
+
+  using GmemTiledCopyK = decltype(make_tiled_copy(
+                            Copy_Atom<GmemCopyOpK, ElementKV>{},
+                            GmemCopyKThrLayout{}, GmemCopyKValLayout{}));
+
+  static constexpr int AlignmentV = 128 / sizeof_bits_v<ElementKV>;
+  static constexpr int ThreadsMajorV = mute::gcd(HEAD_DIM_VO / AlignmentV, NumThreads);
+  static constexpr int ThreadsMinorV = NumThreads / ThreadsMajorV;
+  using GmemCopyVThrLayout = Layout<Shape <Int<ThreadsMinorV>, Int<ThreadsMajorV>>,
+                                    Stride<Int<ThreadsMajorV>,                _1>>;
+  using GmemCopyVValLayout = Layout<Shape<_1, Int<AlignmentV>>>;
+  using GmemCopyOpV = MP22_LDGSTS<uint_bit_t<AlignmentV * sizeof_bits_v<ElementKV>>>;
+
+  using GmemTiledCopyV = decltype(make_tiled_copy(
+                            Copy_Atom<GmemCopyOpV, ElementKV>{},
+                            GmemCopyVThrLayout{}, GmemCopyVValLayout{}));
+#else
+  static constexpr int Alignment = 128 / sizeof_bits_v<Element>;
+  // hard code threads... 8x16 to load 64x16...
+  static constexpr int ThreadsMajorQ = 8;
+  static constexpr int ThreadsMinorQ = NumThreads / ThreadsMajorQ;
+
+  using GmemCopyQThrLayout = Layout<Shape <Int<ThreadsMinorQ>, Int<ThreadsMajorQ>>,
+                                    Stride<Int<ThreadsMajorQ>,                _1>>;
+  using GmemCopyQValLayout = Layout<Shape<_1, Int<Alignment>>>;
+  using GmemCopyOpQ = MP22_ROBUST_LDGSTS<uint_bit_t<Alignment * sizeof_bits_v<ElementQ>>>;
+
+  using GmemTiledCopyQ = decltype(make_tiled_copy(
+                            Copy_Atom<GmemCopyOpQ, ElementQ>{},
+                            GmemCopyQThrLayout{}, GmemCopyQValLayout{}));
+
+  using GmemCopyOpK = MP22_ROBUST_LDGSTS<uint_bit_t<Alignment * sizeof_bits_v<ElementKV>>>;
+  using GmemTiledCopyK = decltype(make_tiled_copy(
+                            Copy_Atom<GmemCopyOpK, ElementKV>{},
+                            GmemCopyQThrLayout{}, GmemCopyQValLayout{}));
+
+  using GmemTiledCopyV = GmemTiledCopyK;
+#endif
+
+  using SharedStorage  = mute::conditional_t<IsQInReg,
+                                SharedStorageReuseQ<Element, SmemLayoutQ, SmemLayoutK, SmemLayoutV, SmemLayoutP>,
+                                SharedStorageSplitHeadDim<Element, SmemLayoutQ, SmemLayoutQPE, SmemLayoutK, SmemLayoutKPE, SmemLayoutV, SmemLayoutP>>;
+};
diff --git a/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_params.hpp b/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_params.hpp
new file mode 100644
index 000000000..5701925c6
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp22_impl/decode_mla_params.hpp
@@ -0,0 +1,152 @@
+#pragma once
+
+#include "decode_mla_kernel_traits.hpp"
+#include "mutlass_extension.hpp"
+
+template <class KernelTraits>
+struct DecodeMLAParams {
+  using Element = typename KernelTraits::Element;
+  using ElementAccumulator = typename KernelTraits::ElementAccumulator;
+
+  using ElementO = typename KernelTraits::ElementO;
+
+  using IndexType = int32_t;
+
+  // Q, K, D_QK, D_VO, (H_R, H_KV), B
+  using ProblemShape = Shape<_1, int, int, int, Shape<int, int>, int>;
+
+  // Q, D_QK, (H_R, H_KV), B
+  using StrideQ = Stride<_0, _1, Stride<int, int>, int>;
+
+  // K, D_QK, (H_R, H_KV), B
+  using StrideK = Stride<int, _1, Stride<_0, int>, int>;
+
+  // V, D_VO, (H_R, H_KV), B
+  using StrideV = StrideK;
+
+  // Q, D_VO, (H_R, H_KV), B
+  using StrideO = StrideQ;
+
+  struct Arguments {
+    Element const* ptr_Q;
+    Element const* ptr_K_Buffer;
+    Element const* ptr_V_Buffer;
+
+    int32_t const* ptr_Req_to_tokens;
+
+    int32_t const* ptr_B_Seqlen;
+
+    ElementO* ptr_O;
+
+    float sm_scale;
+    int NumKVSplits;
+    int NumPages;
+
+    IndexType stride_req_to_tokens_b;
+    IndexType stride_qbs;
+    IndexType stride_qh;
+    IndexType stride_buf_kbs;
+    IndexType stride_buf_kh;
+    IndexType stride_buf_vbs;
+    IndexType stride_buf_vh;
+    IndexType stride_mid_ob;
+    IndexType stride_mid_oh;
+    IndexType stride_mid_os;
+  };
+
+  struct Params {
+    ProblemShape problem_shape;
+
+    Element const* ptr_Q;
+    StrideQ stride_Q;
+    RobustDescriptor desc_Q;
+
+    Element const* ptr_K_Buffer;
+    StrideK stride_K;
+    RobustDescriptor desc_K;
+
+    Element const* ptr_V_Buffer;
+    StrideV stride_V;
+    RobustDescriptor desc_V;
+
+    int32_t const* ptr_Req_to_tokens;
+
+    int32_t const* ptr_B_Seqlen;
+
+    ElementO* ptr_O;
+    StrideO stride_O;
+
+    float sm_scale;
+    int NumKVSplits;
+    IndexType stride_req_to_tokens_b;
+    IndexType stride_qbs;
+    IndexType stride_qh;
+    IndexType stride_buf_kbs;
+    IndexType stride_buf_kh;
+    IndexType stride_buf_vbs;
+    IndexType stride_buf_vh;
+    IndexType stride_mid_ob;
+    IndexType stride_mid_oh;
+    IndexType stride_mid_os;
+
+    float rescale;
+  };
+
+  static Params
+  to_underlying_arguments(ProblemShape problem_shape, Arguments const& args) {
+    int DimQK = mute::size<2>(problem_shape);
+    int DimVO = mute::size<3>(problem_shape);
+    int H = mute::size<4>(problem_shape);
+    int H_KV = mute::size<4,1>(problem_shape);
+    int B = mute::size<5>(problem_shape);
+
+    StrideQ stride_Q = make_stride(_0{}, _1{}, make_stride(H_KV*args.stride_qh, args.stride_qh), args.stride_qbs);
+
+    // handle batch&seq in kernel
+    StrideK stride_K = make_stride(args.stride_buf_kbs, _1{}, make_stride(_0{}, args.stride_buf_kh), 0);
+
+    StrideV stride_V = make_stride(args.stride_buf_vbs, _1{}, make_stride(_0{}, args.stride_buf_vh), 0);
+
+    StrideO stride_O = make_stride(_0{}, _1{}, make_stride(H_KV*args.stride_mid_oh, args.stride_mid_oh), args.stride_mid_ob);
+
+    // Just wrapper robust Q
+    RobustDescriptor desc_Q = make_robust_desc(args.ptr_Q, DimQK * H * B);
+    RobustDescriptor desc_K = make_robust_desc(args.ptr_K_Buffer, DimQK * H_KV * KernelTraits::PAGE_SIZE * args.NumPages);
+    RobustDescriptor desc_V = make_robust_desc(args.ptr_V_Buffer, DimVO * H_KV * KernelTraits::PAGE_SIZE * args.NumPages);
+
+    constexpr float kLog2e = float(1.4426950408889634074);
+    constexpr float rescale = 1.0f / kLog2e;
+
+    float sm_scale = args.sm_scale * kLog2e;
+
+    return {
+      problem_shape,
+      args.ptr_Q,
+      stride_Q,
+      desc_Q,
+      args.ptr_K_Buffer,
+      stride_K,
+      desc_K,
+      args.ptr_V_Buffer,
+      stride_V,
+      desc_V,
+      args.ptr_Req_to_tokens,
+      args.ptr_B_Seqlen,
+      args.ptr_O,
+      stride_O,
+      sm_scale,
+      args.NumKVSplits,
+      args.stride_req_to_tokens_b,
+      args.stride_qbs,
+      args.stride_qh,
+      args.stride_buf_kbs,
+      args.stride_buf_kh,
+      args.stride_buf_vbs,
+      args.stride_buf_vh,
+      args.stride_mid_ob,
+      args.stride_mid_oh,
+      args.stride_mid_os,
+      rescale
+    };
+  }
+};
diff --git a/vllm_musa/csrc/decode_mla/mp22_impl/mutlass_extension.hpp b/vllm_musa/csrc/decode_mla/mp22_impl/mutlass_extension.hpp
new file mode 100644
index 000000000..136bbac4f
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp22_impl/mutlass_extension.hpp
@@ -0,0 +1,410 @@
+#pragma once
+
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 220)
+#define MUTE_MP22_ROBUST_EXT
+#endif
+
+#include <mute/config.hpp>
+#include <mute/atom/copy_atom.hpp>
+#include <mute/arch/copy_mp31_desc.hpp> // RobustDescriptor
+
+namespace mute {
+
+//
+// Robust LD
+//
+template <class TS, class TD = TS>
+struct MP22_ROBUST_LOAD
+{
+  using SRegisters = TS[1];
+  using DRegisters = TD[1];
+
+  static_assert(sizeof(TS) == sizeof(TD),
+                "MP22_ROBUST_LOAD requires sizeof(src_value_type) == sizeof(dst_value_type)");
+
+  MUTE_HOST_DEVICE static void
+  copy(TS               const& gmem_src,
+       TD                    & rmem_dst,
+       bool                    pred,
+       RobustDescriptor const& src_desc)
+
+  {
+#if defined(MUTE_MP22_ROBUST_EXT)
+    void const* gmem_ptr = pred ? reinterpret_cast<void const*>(&gmem_src)
+                         : reinterpret_cast<void const*>(src_desc.oob_addr);
+    // meet the intrinsic requirement
+    using RobustRegV3 = uint32_t __attribute__((vector_size(12)));
+    union ExtractRobustReg {
+      RobustReg reg_v4;
+      struct {
+        RobustRegV3 reg;
+        uint32_t dummy;
+      };
+    };
+
+    ExtractRobustReg desc {src_desc.reg};
+
+    if constexpr (sizeof(TS) == 1) {
+      auto val = __musa_ld_robust_i8(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 2) {
+      auto val = __musa_ld_robust_i16(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 4) {
+      auto val = __musa_ld_robust_i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 8) {
+      auto val = __musa_ld_robust_i64(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 16) {
+      auto val = __musa_ld_robust_v4i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 32) {
+      auto val = __musa_ld_robust_v8i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 64) {
+      auto val = __musa_ld_robust_v16i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 128) {
+      auto val = __musa_ld_robust_v32i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else {
+      static_assert(sizeof(TS) <= 128, "Currently we only support up-to 1024-bit load");
+    }
+#else
+    MUTE_INVALID_CONTROL_PATH(
+        "Trying to use MP22_ROBUST_LOAD without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+  }
+};
+
+
+//
+// Robust load traits
+//
+template <class S, class D>
+struct MP22_ROBUST_LOAD_OP : MP22_ROBUST_LOAD<S, D> {};
+
+template <class S, class D>
+struct MP22_ROBUST_PRED_LOAD_OP : MP22_ROBUST_LOAD<S, D> {};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_PRED_LOAD_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  // Predicate value that determines whether to load or return zero
+  bool pred;
+
+  // Overload copy_unpack to pass robust buffer access descriptor
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst)
+  {
+    static_assert(is_gmem<TS>::value, "Expected gmem src for MP22_ROBUST_LOAD");
+    static_assert(is_rmem<TD>::value, "Expected rmem dst for MP22_ROBUST_LOAD");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    MUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+        "This src layout is incompatible with this tiled copy.");
+    MUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+        "This dst layout is incompatible with this tiled copy.");
+
+    MP22_ROBUST_LOAD<S, D>::copy(rS[0], rD[0], traits.pred, traits.src_desc);
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_LOAD_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP22_ROBUST_PRED_LOAD_OP<S, D>>
+  with(bool pred) const {
+    return {src_desc, pred};
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_LOAD<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Construct an executable MP22_ROBUST_LOAD with robust buffer access descriptor
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP22_ROBUST_LOAD_OP<S, D>>
+  with(RobustDescriptor desc) const {
+    return {desc};
+  }
+
+  // Don't try to execute a copy with MP22_ROBUST_LOAD before calling .with()
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst) = delete;
+
+};
+
+
+//
+// Robust LDGSTS
+//
+template <class TS, class TD = TS>
+struct MP22_ROBUST_LDGSTS
+{
+  using SRegisters = TS[1];
+  using DRegisters = TD[1];
+
+  static_assert(sizeof(TS) == sizeof(TD),
+                "MP22_ROBUST_LDGSTS requires sizeof(src_value_type) == sizeof(dst_value_type)");
+  static_assert(sizeof(TS) == 1  || sizeof(TS) == 2  || sizeof(TS) == 4  || sizeof(TS) == 8 ||
+                sizeof(TS) == 16 || sizeof(TS) == 32 || sizeof(TS) == 64 || sizeof(TS) == 128,
+                "MP22_ROBUST_LDGSTS sizeof(TS) is not supported");
+
+  MUTE_HOST_DEVICE static void
+  copy(TS               const& gmem_src,
+       TD                    & smem_dst,
+       bool                    pred,
+       RobustDescriptor const& src_desc)
+  {
+#if defined(MUTE_MP22_ROBUST_EXT)
+    uint64_t gmem_int_ptr = reinterpret_cast<uint64_t>(&gmem_src);
+    uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst);
+    auto gmem_ptr = make_ptr_with_address_space<AddressSpace::Global>(
+                            pred ? gmem_int_ptr : src_desc.oob_addr);
+    auto smem_ptr = make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr);
+
+    // meet the intrinsic requirement
+    using RobustRegV3 = uint32_t __attribute__((vector_size(12)));
+    union ExtractRobustReg {
+      RobustReg reg_v4;
+      struct {
+        RobustRegV3 reg;
+        uint32_t dummy;
+      };
+    };
+
+    ExtractRobustReg desc {src_desc.reg};
+
+    __musa_memcpy_g2s_robust(smem_ptr, gmem_ptr, sizeof(TS), desc.reg, 0);
+#else
+    MUTE_INVALID_CONTROL_PATH(
+        "Trying to use MP22_ROBUST_LDGSTS without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+  }
+};
+
+MUTE_HOST_DEVICE
+void
+ldlms_wait() {
+#if defined(MUTE_MP22_ROBUST_EXT)
+  __musa_memcpy_g2s_wait();
+#else
+  MUTE_INVALID_CONTROL_PATH("Trying to use MemcpyG2S Wait without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+}
+
+
+//
+// Robust LDGSTS traits
+//
+template <class S, class D>
+struct MP22_ROBUST_LDGSTS_OP : MP22_ROBUST_LDGSTS<S, D> {};
+
+template <class S, class D>
+struct MP22_ROBUST_PRED_LDGSTS_OP : MP22_ROBUST_LDGSTS<S, D> {};
+
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_PRED_LDGSTS_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  // Predicate value that determines whether to load or return zero
+  bool pred;
+
+  // Overload copy_unpack to pass robust buffer access descriptor
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst)
+  {
+    static_assert(is_gmem<TS>::value, "Expected gmem src for MP22_ROBUST_LDGSTS");
+    static_assert(is_smem<TD>::value, "Expected smem dst for MP22_ROBUST_LDGSTS");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    MUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+        "This src layout is incompatible with this tiled copy.");
+    MUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+        "This dst layout is incompatible with this tiled copy.");
+
+    MP22_ROBUST_LDGSTS<S, D>::copy(rS[0], rD[0], traits.pred, traits.src_desc);
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_LDGSTS_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP22_ROBUST_PRED_LDGSTS_OP<S, D>>
+  with(bool pred) const {
+    return {src_desc, pred};
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_LDGSTS<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Construct an executable MP22_ROBUST_LDGSTS with robust buffer access descriptor
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP22_ROBUST_LDGSTS_OP<S, D>>
+  with(RobustDescriptor desc) const {
+    return {desc};
+  }
+
+  // Don't try to execute a copy with MP22_ROBUST_LDGSTS before calling .with()
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst) = delete;
+};
+
+
+//
+// LDGSTS
+//
+
+template <class TS, class TD = TS>
+struct MP22_LDGSTS
+{
+  using SRegisters = TS[1];
+  using DRegisters = TD[1];
+
+  static_assert(sizeof(TS) == sizeof(TD),
+                "MP22_ROBUST_LDGSTS requires sizeof(src_value_type) == sizeof(dst_value_type)");
+  static_assert(sizeof(TS) == 1  || sizeof(TS) == 2  || sizeof(TS) == 4  || sizeof(TS) == 8 ||
+                sizeof(TS) == 16 || sizeof(TS) == 32 || sizeof(TS) == 64 || sizeof(TS) == 128,
+                "MP22_ROBUST_LDGSTS sizeof(TS) is not supported");
+
+  MUTE_HOST_DEVICE static void
+  copy(TS               const& gmem_src,
+       TD                    & smem_dst)
+  {
+#if defined(MUTE_MP22_ROBUST_EXT)
+    uint64_t gmem_int_ptr = reinterpret_cast<uint64_t>(&gmem_src);
+    uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst);
+    auto gmem_ptr = make_ptr_with_address_space<AddressSpace::Global>(gmem_int_ptr);
+    auto smem_ptr = make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr);
+
+    __musa_memcpy_g2s(smem_ptr, gmem_ptr, sizeof(TS), 0);
+#else
+    MUTE_INVALID_CONTROL_PATH(
+        "Trying to use MP22_ROBUST_LDGSTS without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_LDGSTS<S,D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+};
+
+} // namespace mute
diff --git a/vllm_musa/csrc/decode_mla/mp22_impl/online_softmax.hpp b/vllm_musa/csrc/decode_mla/mp22_impl/online_softmax.hpp
new file mode 100644
index 000000000..5ace39980
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp22_impl/online_softmax.hpp
@@ -0,0 +1,278 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+**************************************************************************************************/
+
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mutlass/fast_math.h>
+#include "../online_softmax_ops.hpp"
+
+using namespace mute;
+
+template <class LayoutAccum>
+MUTE_DEVICE
+auto
+mp22_convert_layout_acc_rowcol(LayoutAccum acc_layout) {
+  if constexpr (decltype(rank<0>(acc_layout))::value == 3) { // MP31
+    // from ((2, VN, VM), MMA_M, MMA_N) to ((VM, MMA_M),(2, VN, MMA_N));
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(rank(acc_layout))::value == 3);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0,2>(l), get<1>(l)), make_layout(get<0,0>(l), get<0,1>(l), get<2>(l)));
+  } else {  // MP22
+    // from((2,4), MMA_M, MMA_N) to ((4, MMA_M), (2, MMA_N))
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(size<0, 1>(acc_layout))::value == 4);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0,1>(l), get<1>(l)), make_layout(get<0,0>(l), get<2>(l)));
+  }
+}
+
+template <typename Tensor0, typename Tensor1, typename Operator>
+MUTE_DEVICE void quad_allreduce_(Tensor0& dst, Tensor1& src, Operator& op) {
+  MUTE_STATIC_ASSERT_V(size(dst) == size(src));
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(dst); ++i) {
+    dst(i) = Allreduce<4>::run(src(i), op);
+  }
+}
+
+template <typename Tensor0, typename Tensor1, typename Operator>
+MUTE_DEVICE void extra_allreduce_(Tensor0& dst, Tensor1& src, Operator& op) {
+  MUTE_STATIC_ASSERT_V(size(dst) == size(src));
+  constexpr int mask = ((~(128 - 1)) & 0x7f) << 7 | (128 - 1);
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(dst); ++i) {
+    int val = *reinterpret_cast<int32_t*>(&src(i));
+    int ret = __musa_shfl_xor_sync_i32(val, 32 & 0x7f, mask);
+    // TODO: assume float here...get type from tensor
+    dst(i) = op(src(i), *reinterpret_cast<float*>(&ret));
+  }
+}
+
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class Fragment0, class Fragment1, class AccumType>
+MUTE_HOST_DEVICE
+void
+mp22_online_softmax_scale(Fragment0& accum, Fragment1& accum_softmax, AccumType rln2_scale) {
+  MUTE_UNROLL
+  for (int i = 0; i < size(accum); ++i) {
+    accum_softmax(i) = accum(i) * rln2_scale;
+  }
+}
+
+template <class Fragment0, class Fragment1, class AccumType, class Fragment2>
+MUTE_HOST_DEVICE
+void
+mp22_online_softmax_scale(Fragment0& accum, Fragment1& accum_softmax, AccumType rln2_scale, Fragment2& tPcP, int oob_n, bool cond) {
+  Tensor accum_dst = make_tensor(accum_softmax.data(), mp22_convert_layout_acc_rowcol(accum_softmax.layout()));
+  Tensor accum_src = make_tensor(accum.data(), mp22_convert_layout_acc_rowcol(accum.layout()));
+  Tensor cP_view   = make_tensor(tPcP.data(), mp22_convert_layout_acc_rowcol(tPcP.layout()));
+
+  MUTE_STATIC_ASSERT_V(size<0>(accum_dst) == size<0>(accum_src));
+  MUTE_STATIC_ASSERT_V(size<1>(accum_dst) == size<1>(accum_src));
+
+  //if (threadIdx.x == 0) {
+  //  printf("bid:%d cond:%d\n", blockIdx.x, cond);
+  //}
+
+  MUTE_UNROLL
+  for (int m = 0; m < size<0>(accum_dst); ++m) {
+    MUTE_UNROLL
+    for (int n = 0; n < size<1>(accum_dst); ++n) {
+      accum_dst(m, n) = (cond && get<1>(cP_view(0,n)) >= oob_n) ?
+                                -mutlass::platform::numeric_limits<AccumType>::infinity()
+                              : accum_src(m, n) * rln2_scale;
+    }
+  }
+
+}
+
+template <class ElementAccum, class Fragment0, class Fragmnet1>
+MUTE_HOST_DEVICE
+void
+mp22_online_softmax_row_max(Fragment0& accum_softmax,
+                       Fragmnet1& row_max,
+                       Fragmnet1& old_row_max,
+                       Fragmnet1& alpha) {
+  Tensor accum_s = make_tensor(accum_softmax.data(), mp22_convert_layout_acc_rowcol(accum_softmax.layout()));
+
+  MUTE_STATIC_ASSERT_V(size(row_max) == size<0>(accum_s));
+  MUTE_STATIC_ASSERT_V(size(old_row_max) == size<0>(accum_s));
+  MUTE_STATIC_ASSERT_V(size(alpha) == size<0>(accum_s));
+  //static_assert(size<1>(accum_s) == 16);
+
+  MaxOp<ElementAccum> op;
+  ThreadReduce<size<1>(accum_s)> thread_reduce;
+
+  thread_reduce.template run<true>(accum_s, row_max, op);
+
+  octave_allreduce_(row_max, row_max, op);
+  // use width=128 shulff xor
+  extra_allreduce_(row_max, row_max, op);
+
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    row_max(mi) = mutlass::fast_max(row_max(mi), old_row_max(mi));
+
+    alpha(mi) = old_row_max(mi) - row_max(mi);
+    alpha(mi) = exp2f(alpha(mi));
+    old_row_max(mi) = row_max(mi);
+  }
+
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_s); ++ni) {
+      accum_s(mi, ni) = accum_s(mi, ni) - row_max(mi);
+    }
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+mp22_update_accum(Fragment0& accum, Fragment1& alpha) {
+  Tensor accum_o = make_tensor(accum.data(), mp22_convert_layout_acc_rowcol(accum.layout()));
+  MUTE_STATIC_ASSERT_V(size(alpha) == size<0>(accum_o));
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_o); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_o); ++ni) {
+      accum_o(mi, ni) *= alpha(mi);
+    }
+  }
+}
+
+template <class ElementAccum, class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+mp22_online_softmax_row_sum(Fragment0& accum_softmax, Fragment1& row_sum,
+                       Fragment1& lse, Fragment1& alpha) {
+  Tensor accum_s = make_tensor(accum_softmax.data(), mp22_convert_layout_acc_rowcol(accum_softmax.layout()));
+  //static_assert(size<1>(accum_s) == 16);
+  MUTE_STATIC_ASSERT_V(size<0>(accum_s) == size(row_sum));
+  SumOp<ElementAccum> op;
+  ThreadReduce<size<1>(accum_s)> thread_reduce;
+  thread_reduce.template run<true>(accum_s, row_sum, op);
+  octave_allreduce_(row_sum, row_sum, op);
+  // use width=128 shulff xor
+  extra_allreduce_(row_sum, row_sum, op);
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    lse(mi) = lse(mi) * alpha(mi) + row_sum(mi);
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+mp22_apply_softmax_normalizer(Fragment0& accum, Fragment1& lse) {
+  Tensor accum_o = make_tensor(accum.data(), mp22_convert_layout_acc_rowcol(accum.layout()));
+  MUTE_STATIC_ASSERT_V(size<0>(accum_o) == size(lse));
+
+  Tensor rlse = make_fragment_like(lse);
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(lse); ++i) {
+    rlse(i) = 1.0f / lse(i);
+  }
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_o); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_o); ++ni) {
+      accum_o(mi, ni) *= rlse(mi);
+    }
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+mp22_apply_causal_mask(Fragment0& accum_, Fragment1& tPcP_, int m_idx, int n_idx, bool debug) {
+  Tensor accum = make_tensor(accum_.data(), mp22_convert_layout_acc_rowcol(accum_.layout()));
+  Tensor tPcP  = make_tensor(tPcP_.data(), mp22_convert_layout_acc_rowcol(tPcP_.layout()));
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum); ++mi) {
+    const int row_idx = m_idx + get<0>(tPcP(mi, 0));
+    const int col_idx_limit = row_idx + 1;
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum); ++ni) {
+      const int col_idx = n_idx + get<1>(tPcP(0, ni));
+      //if (debug && blockIdx.x == 1) {
+      //  printf("tid:%d col_idx_limit:%d col_idx:%d\n", threadIdx.x, col_idx_limit, col_idx);
+      //}
+      if (col_idx >= col_idx_limit) {
+        accum(mi, ni) = -std::numeric_limits<float>::infinity();
+      }
+    }
+  }
+}
+
+template <bool zero_clean, bool check_n, class Tensor0, class Tensor1>
+MUTE_HOST_DEVICE
+void
+mp22_process_qk_oob(Tensor0& accum_, Tensor1& tPcP_, int m_base, int VALID_BLOCK_H, int Q_HEAD_NUM, int n_base, int n_bound) {
+
+  Tensor accum = make_tensor(accum_.data(), mp22_convert_layout_acc_rowcol(accum_.layout()));
+  Tensor tPcP  = make_tensor(tPcP_.data(), mp22_convert_layout_acc_rowcol(tPcP_.layout()));
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum); ++mi) {
+    // don't need?
+    //const bool valid_h = get<0>(tPcP(mi, 0)) < VALID_BLOCK_H && (m_base + get<0>(tPcP(mi, 0)) < Q_HEAD_NUM);
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum); ++ni) {
+      //bool valid_n = n_base + get<1>(tPcP(0,ni)) < n_bound;
+      //
+
+      bool valid_n = true;
+      if constexpr (check_n) {
+        valid_n = n_base + get<1>(tPcP(0,ni)) < n_bound;
+      }
+
+      float oob_value = -std::numeric_limits<float>::infinity();
+      if constexpr (zero_clean) {
+        oob_value = 0.0f;
+      }
+      accum(mi, ni) = valid_n ? accum(mi, ni) : oob_value;
+    }
+  }
+}
diff --git a/vllm_musa/csrc/decode_mla/mp22_impl/utils.hpp b/vllm_musa/csrc/decode_mla/mp22_impl/utils.hpp
new file mode 100644
index 000000000..4450c5528
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp22_impl/utils.hpp
@@ -0,0 +1,41 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+
+template <
+  bool A_in_regs = false,
+  bool B_in_regs = false,
+  typename Tensor0,
+  typename Tensor1,
+  typename Tensor2,
+  typename Tensor3,
+  typename Tensor4,
+  typename TiledMma,
+  typename ThrCopyA,
+  typename ThrCopyB
+>
+MUTE_DEVICE
+void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsA,
+          Tensor4 const& tCsB, TiledMma tiled_mma,
+          ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) {
+  MUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+  MUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+  MUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
+  Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
+  MUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
+  Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+  MUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
+
+  if constexpr (!A_in_regs) { copy(tCsA(_,_,_0{}), tCrA_copy_view(_,_,_0{})); }
+  if constexpr (!B_in_regs) { copy(tCsB(_,_,_0{}), tCrB_copy_view(_,_,_0{})); }
+
+  MUTE_UNROLL
+  for (int i = 0; i < size<2>(tCrA); ++i) {
+    if (i < size<2>(tCrA) - 1) {
+      if constexpr (!A_in_regs) { copy(tCsA(_,_,i + 1), tCrA_copy_view(_,_,i + 1)); }
+      if constexpr (!B_in_regs) { copy(tCsB(_,_,i + 1), tCrB_copy_view(_,_,i + 1)); }
+    }
+
+    mute::gemm(tiled_mma, tCrA(_,_,i), tCrB(_,_,i), acc);
+  }
+}
diff --git a/vllm_musa/csrc/decode_mla/mp31_impl/decode_mla.hpp b/vllm_musa/csrc/decode_mla/mp31_impl/decode_mla.hpp
new file mode 100644
index 000000000..fa8d29ac1
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp31_impl/decode_mla.hpp
@@ -0,0 +1,600 @@
+#pragma once
+
+#include "kernel_tratis.hpp"
+#include "online_softmax.hpp"
+
+#include <mutlass/numeric_conversion.h>
+
+template <class KernelTraits>
+__global__ __launch_bounds__(KernelTraits::NumThreads, 1)
+void mla_decode(typename KernelTraits::PackGQAParams const params) {
+  #if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 310)
+  __shared__ char __attribute__((aligned(256))) smem_buf[sizeof(typename KernelTraits::SharedStorage)];
+  auto& shared_storage = *reinterpret_cast<typename KernelTraits::SharedStorage*>(smem_buf);
+
+  auto [m_coord, h_kv_coord, batch_and_split_coord] = static_cast<uint3>(blockIdx);
+
+  auto split_coord = batch_and_split_coord % params.num_splits;
+  auto b_coord = batch_and_split_coord / params.num_splits;
+
+  int cur_batch_seq_len = params.ptr_seqlen[b_coord];
+
+  constexpr int BlockN = KernelTraits::BlockN;
+
+  int n_tiles = ceil_div(cur_batch_seq_len, BlockN);
+  int n_tiles_per_split = ceil_div(n_tiles, params.num_splits);
+
+  int split_kv_start = n_tiles_per_split * split_coord;
+  int split_kv_end   = min(split_kv_start + n_tiles_per_split, n_tiles);
+
+  // early return
+  if (split_kv_end <= split_kv_start) return;
+
+  int thread_idx = threadIdx.x;
+
+  //if (thread_idx == 0 && h_q_coord == 0) {
+  //  printf("batch_id:%d split_id:%d cur_seqlen:%d, n_tiles:%d split:(%d %d)\n", b_coord, split_coord, cur_batch_seq_len, n_tiles, split_kv_start, split_kv_end);
+  //}
+
+  const int* block_table = params.ptr_block_table + b_coord * params.stride_block_table_b;
+
+  Tensor sQ  = make_tensor(make_smem_ptr(shared_storage.smem_q.data()),  typename KernelTraits::SmemLayoutQ{});
+  Tensor sK  = make_tensor(make_smem_ptr(shared_storage.smem_k.data()),  typename KernelTraits::SmemLayoutK{});
+  Tensor sP  = make_tensor(make_smem_ptr(shared_storage.smem_p.data()),  typename KernelTraits::SmemLayoutP{});
+  Tensor sV  = make_tensor(make_smem_ptr(shared_storage.smem_v.data()),  typename KernelTraits::SmemLayoutV{});
+  Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_vt.data()), typename KernelTraits::SmemLayoutVt{});
+  Tensor sAlpha = make_tensor(make_smem_ptr(shared_storage.smem_score.data()), typename KernelTraits::SmemLayoutScore{});
+  Tensor sLse = make_tensor(make_smem_ptr(shared_storage.smem_lse.data()), typename KernelTraits::SmemLayoutLse{});
+
+
+  typename KernelTraits::TiledMmaQK tiled_mma_qk;
+  typename KernelTraits::TiledMmaPV tiled_mma_pv;
+  typename KernelTraits::TiledMmaQKPermute tiled_mma_qk_perm;
+
+  auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx % KernelTraits::NumThreadsQK);
+  auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx % KernelTraits::NumThreadsPV);
+  auto thr_mma_qk_perm = tiled_mma_qk_perm.get_thread_slice(thread_idx % KernelTraits::NumThreadsQK);
+
+  // MMA QK
+  Tensor tCsQ = thr_mma_qk.partition_A(sQ);
+  Tensor tCrQ = thr_mma_qk.make_fragment_A(tCsQ);
+
+  Tensor tCsK = thr_mma_qk.partition_B(sK);
+  Tensor tCrK = thr_mma_qk.make_fragment_B(tCsK);
+
+  // MMA PV
+  Tensor tCsP = thr_mma_pv.partition_A(sP);
+  Tensor tCrP = thr_mma_pv.make_fragment_A(tCsP);
+
+  Tensor tCsV = thr_mma_pv.partition_B(sVt);
+  Tensor tCrV = thr_mma_pv.make_fragment_B(tCsV);
+
+  // R2S
+  using PermuteAlignmentType = mute::uint_bit_t<sizeof_bits_v<typename KernelTraits::Element> * KernelTraits::FragmentSize>;
+
+  auto tiled_copy_r2s = make_tiled_copy_C(Copy_Atom<UniversalCopy<PermuteAlignmentType>, typename KernelTraits::Element>{},
+                                          typename KernelTraits::TiledMmaQKPermute{});
+  auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(thread_idx % KernelTraits::NumThreadsQK);
+  Tensor tPsP = thr_copy_r2s.partition_D(sP);
+
+  constexpr int RowsPerThread = KernelTraits::BlockM / (KernelTraits::NumThreadsQK / 128) / 16;
+
+  Tensor alpha       = make_tensor<typename KernelTraits::ElementAccum>(Shape<Int<RowsPerThread>>{});
+  Tensor lse         = make_fragment_like(alpha);
+  Tensor row_sum     = make_fragment_like(alpha);
+  Tensor row_max     = make_fragment_like(alpha);
+  Tensor old_row_max = make_fragment_like(alpha);
+
+
+  fill(lse, typename KernelTraits::ElementAccum(0.0f));
+  fill(alpha, typename KernelTraits::ElementAccum(0.0f));
+  fill(old_row_max, -mutlass::platform::numeric_limits<typename KernelTraits::ElementAccum>::infinity());
+
+  constexpr int NumWarpsQK = KernelTraits::NumThreadsQK / mutlass::NumThreadsPerWarp;
+  constexpr int NumWarpsPV = KernelTraits::NumThreadsPV / mutlass::NumThreadsPerWarp;
+
+  int warp_idx = mutlass::canonical_warp_idx();
+  int warp_squad_idx = mutlass::canonical_warp_squad_idx();
+
+  enum class WarpSquadRole {
+    ProducerConsumerPV,
+    ProducerConsumerQK,
+  };
+
+  auto warp_squad_role = WarpSquadRole(warp_squad_idx);
+
+  //bool is_producer_consumer_pv = warp_idx < NumWarpsPV;
+  //bool is_producer_consumer_qk = !is_producer_consumer_pv;
+
+  mutlass::arch::allocate_async_barriers(sizeof(typename KernelTraits::BarrierStorage));
+  typename KernelTraits::BarrierStorage* barrier_ids = reinterpret_cast<typename KernelTraits::BarrierStorage*>(0);
+
+  typename KernelTraits::BarrierQ barrierQ(reinterpret_cast<uint64_t>(&barrier_ids->BarrierQ));
+
+  if (warp_idx == 0) {
+    barrierQ.init(1);
+  }
+
+  typename KernelTraits::MainloopPipelineK::Params k_pipe_params;
+  k_pipe_params.producer_arv_count = NumWarpsQK;
+  k_pipe_params.consumer_arv_count = NumWarpsPV;
+  typename KernelTraits::MainloopPipelineK pipeline_k(k_pipe_params, reinterpret_cast<uint64_t>(&barrier_ids->PipelineK));
+
+  typename KernelTraits::MainloopPipelineK::PipelineState pipeline_k_producer_state = mutlass::make_producer_start_state<typename KernelTraits::MainloopPipelineK>();
+  typename KernelTraits::MainloopPipelineK::PipelineState pipeline_k_consumer_state;
+
+
+  typename KernelTraits::MainloopPipelineV::Params v_pipe_params;
+  v_pipe_params.transaction_bytes = KernelTraits::TmeTransactionBytesV;
+  v_pipe_params.num_consumers = NumWarpsPV;
+  v_pipe_params.num_producers = 1;
+  typename KernelTraits::MainloopPipelineV pipeline_v(v_pipe_params, reinterpret_cast<uint64_t>(&barrier_ids->PipelineV));
+
+  typename KernelTraits::MainloopPipelineV::PipelineState pipeline_v_producer_state = mutlass::make_producer_start_state<typename KernelTraits::MainloopPipelineV>();
+  typename KernelTraits::MainloopPipelineV::PipelineState pipeline_v_consumer_state;
+
+  typename KernelTraits::MainloopPipelineP::Params p_pipe_params;
+  p_pipe_params.producer_arv_count = NumWarpsQK;
+  p_pipe_params.consumer_arv_count = NumWarpsPV;
+  typename KernelTraits::MainloopPipelineP pipeline_p(p_pipe_params, reinterpret_cast<uint64_t>(&barrier_ids->PipelineP));
+
+  typename KernelTraits::MainloopPipelineP::PipelineState pipeline_p_producer_state = mutlass::make_producer_start_state<typename KernelTraits::MainloopPipelineP>();
+  typename KernelTraits::MainloopPipelineP::PipelineState pipeline_p_consumer_state;
+
+  typename KernelTraits::MainloopPipelineAlpha::Params alpha_pipe_params;
+  alpha_pipe_params.producer_arv_count = NumWarpsQK;
+  alpha_pipe_params.consumer_arv_count = NumWarpsPV;
+  typename KernelTraits::MainloopPipelineAlpha pipeline_alpha(alpha_pipe_params, reinterpret_cast<uint64_t>(&barrier_ids->PipelineAlpha));
+
+  typename KernelTraits::MainloopPipelineAlpha::PipelineState pipeline_alpha_producer_state = mutlass::make_producer_start_state<typename KernelTraits::MainloopPipelineAlpha>();
+  typename KernelTraits::MainloopPipelineAlpha::PipelineState pipeline_alpha_consumer_state;
+
+  __syncthreads();
+
+  using TileShapeQKD = typename KernelTraits::TileShapeQKD;
+  using TileShapePDV = typename KernelTraits::TileShapePDV;
+
+  // QK
+  if (warp_squad_role == WarpSquadRole::ProducerConsumerQK) {
+    auto accum_p = partition_fragment_C(tiled_mma_qk, select<0, 1>(typename KernelTraits::TileShapeQKD{}));
+    auto accum_softmax = make_fragment_like(accum_p);
+    auto accum_cvt = make_fragment_like<typename KernelTraits::Element>(accum_softmax);
+    clear(accum_p);
+
+    Tensor tPcP = thr_mma_qk_perm.partition_C(make_identity_tensor(select<0, 1>(typename KernelTraits::TileShapeQKD{})));
+    Tensor tPcP_mn = make_tensor(tPcP.data(), mp31_convert_layout_acc_rowcol(tPcP.layout()));
+    Tensor accum_p_mn = make_tensor(accum_p.data(), mp31_convert_layout_acc_rowcol(accum_p.layout()));
+    Tensor accum_softmax_mn = make_tensor(accum_softmax.data(), mp31_convert_layout_acc_rowcol(accum_softmax.layout()));
+
+    Tensor tPrP = thr_copy_r2s.retile_S(accum_cvt);
+    Tensor tCvt_frg = recast<mutlass::Array<typename KernelTraits::Element, KernelTraits::FragmentSize>>(accum_cvt);
+    Tensor tAcc_frg = recast<mutlass::Array<typename KernelTraits::ElementAccum, KernelTraits::FragmentSize>>(accum_softmax);
+
+    // Assume PageSize == TileN
+    // TODO: relax it
+    static_assert(KernelTraits::PageSize == KernelTraits::BlockN);
+
+    constexpr int PagesPerBlockN = KernelTraits::BlockN / KernelTraits::PageSize;
+
+    // Reverse to early process OOB
+    int cur_tile_idx = split_kv_end - 1;
+
+    int kv_page_number = block_table[cur_tile_idx];
+
+    TiledCopy gmem_tiled_copy_K = typename KernelTraits::GmemTiledCopyK{};
+    ThrCopy gmem_thr_copy_K = gmem_tiled_copy_K.get_thread_slice(thread_idx % KernelTraits::NumThreadsQK);
+
+    Tensor cK = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK)));
+    Tensor tKcK = gmem_thr_copy_K.partition_S(cK);
+    Tensor tKpK = make_tensor<bool>(make_shape(size<1>(tKcK), size<2>(tKcK)), Stride<_1, _0>{});
+
+    Tensor tKsK = gmem_thr_copy_K.partition_D(sK);
+
+    MUTE_UNROLL
+    for (int i = 0; i < size<0>(tKpK); ++i) {
+      tKpK(i) = cur_tile_idx * KernelTraits::BlockN + get<0>(tKcK(0, i, 0)) < cur_batch_seq_len;
+    }
+
+    Tensor mK = make_tensor(make_gmem_ptr(params.ptr_K), params.layout_K);
+    // TODO:permute tile maybe less than block tile
+    Tensor gK_pages = local_tile(mK, typename KernelTraits::LdgstsTileK{}, make_coord(_0{}, _0{}, h_kv_coord, _)); // (PageSizeTile, HeadDimTile, pages)
+
+    bool is_oob_tile = (cur_tile_idx + 1) * KernelTraits::BlockN > cur_batch_seq_len;
+
+    {
+      Tensor gK = gK_pages(_, _, kv_page_number);
+      Tensor tKgK = gmem_thr_copy_K.partition_S(gK);
+
+      pipeline_k.producer_acquire(pipeline_k_producer_state);
+
+      copy(gmem_tiled_copy_K.with(params.desc_K), tKgK, tKsK(_,_,_,pipeline_k_producer_state.index()));
+
+
+      //if (thread_idx == KernelTraits::NumThreadsPV && b_coord == 0 && split_coord == 0 && h_q_coord == 0) {
+      //  print("mK:"); print(mK); print("\n");
+      //  print("gK_pages:"); print(gK_pages); print("\n");
+      //  print("gK:"); print(gK); print("\n");
+      //  print("tKgK:"); print(tKgK); print("\n");
+      //}
+    }
+
+    ldgsts_wait();
+    pipeline_k.producer_commit(pipeline_k_producer_state);
+    ++pipeline_k_producer_state;
+
+    barrierQ.wait(0);
+    pipeline_k.consumer_wait(pipeline_k_consumer_state);
+
+    mute::gemm(tiled_mma_qk, tCrQ, tCrK(_,_,_,pipeline_k_consumer_state.index()), accum_p);
+
+
+    warpsquad_wait();
+
+    constexpr int n_masking_steps = !KernelTraits::IsCausal ? 1 : mute::ceil_div(KernelTraits::BlockM, KernelTraits::BlockN) + 1;
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_p_mn); ++ni) {
+      //bool is_oob_token = is_oob_tile && (cur_tile_idx * KernelTraits::BlockN + get<1>(tPcP_mn(0, ni))) >= cur_batch_seq_len;
+      bool is_oob_token = get<1>(tPcP_mn(0, ni)) >= (cur_batch_seq_len - cur_tile_idx * KernelTraits::BlockN);
+      MUTE_UNROLL
+      for (int mi = 0; mi < size<0>(accum_p_mn); ++mi) {
+        if constexpr (!KernelTraits::IsCausal) {
+          accum_p_mn(mi, ni) = is_oob_token ? -mutlass::platform::numeric_limits<typename KernelTraits::ElementAccum>::infinity() : accum_p_mn(mi ,ni);
+        } else {
+          int row = get<0>(tPcP_mn(mi, 0));
+          // int col_limit_right = cur_batch_seq_len -1  - cur_tile_idx * KernelTraits::BlockN - (params.q_seq_per_hkv - 1 - (m_coord * KernelTraits::BlockM + row)) / params.h_r;
+          int col_limit_right = (m_coord * KernelTraits::BlockM + row) / params.h_r + cur_batch_seq_len - params.q_seq_per_hkv / params.h_r - (cur_tile_idx) * KernelTraits::BlockN;
+          if (get<1>(tPcP_mn(0, ni)) > col_limit_right) {
+            accum_p_mn(mi, ni) = -mutlass::platform::numeric_limits<typename KernelTraits::ElementAccum>::infinity();
+          }
+        }
+      }
+    }
+
+    int thread_idx_qk = thread_idx - KernelTraits::NumThreadsPV;
+
+    --cur_tile_idx;
+
+    for (int masking_step = n_masking_steps - 1; cur_tile_idx >= split_kv_start; --masking_step, --cur_tile_idx) {
+      if constexpr (KernelTraits::KStages == 1) {
+        warpsquad_wait();
+        pipeline_k.consumer_release(pipeline_k_consumer_state);
+        ++pipeline_k_consumer_state;
+      }
+      // load K
+      {
+        kv_page_number = block_table[cur_tile_idx];
+        Tensor gK = gK_pages(_, _, kv_page_number);
+
+        Tensor tKgK = gmem_thr_copy_K.partition_S(gK);
+
+        pipeline_k.producer_acquire(pipeline_k_producer_state);
+
+        copy(gmem_tiled_copy_K.with(params.desc_K), tKgK, tKsK(_,_,_,pipeline_k_producer_state.index()));
+      }
+
+      if constexpr (KernelTraits::KStages > 1) {
+        warpsquad_wait();
+        pipeline_k.consumer_release(pipeline_k_consumer_state);
+        ++pipeline_k_consumer_state;
+      }
+
+      mp31_online_softmax_scale(accum_p, accum_softmax, params.sm_scale);
+      mp31_online_softmax_row_max<typename KernelTraits::ElementAccum>(accum_softmax, row_max, old_row_max, alpha); // alpha is updated
+
+      clear(accum_p);
+
+      ldgsts_wait();
+      pipeline_k.producer_commit(pipeline_k_producer_state);
+      ++pipeline_k_producer_state;
+
+
+      pipeline_k.consumer_wait(pipeline_k_consumer_state);
+
+      mute::gemm(tiled_mma_qk, tCrQ, tCrK(_,_,_,pipeline_k_consumer_state.index()), accum_p);
+
+      online_softmax_exp2(accum_softmax);
+      mp31_online_softmax_row_sum<typename KernelTraits::ElementAccum>(accum_softmax, row_sum, lse, alpha);
+
+      // STS alpha
+      pipeline_alpha.producer_acquire(pipeline_alpha_producer_state);
+
+      MUTE_UNROLL
+      for (int i = 0; i < size(alpha); ++i) {
+        sAlpha(get<0>(tPcP_mn(i, _0{})), pipeline_alpha_producer_state.index()) = alpha(i);
+      }
+
+      MUTE_UNROLL
+      for (int i = 0; i < size(tCvt_frg); ++i) {
+        tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::Element, typename KernelTraits::ElementAccum, KernelTraits::FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+      }
+
+      // STS P
+      pipeline_p.producer_acquire(pipeline_p_producer_state);
+      copy(tiled_copy_r2s, tPrP, tPsP);
+
+      __threadfence_block();
+      pipeline_p.producer_commit(pipeline_p_producer_state);
+      pipeline_alpha.producer_commit(pipeline_alpha_producer_state);
+      ++pipeline_p_producer_state;
+      ++pipeline_alpha_producer_state;
+
+      warpsquad_wait();
+      const bool is_masking_step = masking_step > 0;
+      if (is_masking_step) {
+        MUTE_UNROLL
+        for (int ni = 0; ni < size<1>(accum_p_mn); ++ni) {
+          //bool is_oob_token = is_oob_tile && (cur_tile_idx * KernelTraits::BlockN + get<1>(tPcP_mn(0, ni))) >= cur_batch_seq_len;
+          bool is_oob_token = get<1>(tPcP_mn(0, ni)) >= (cur_batch_seq_len - cur_tile_idx * KernelTraits::BlockN);
+          MUTE_UNROLL
+          for (int mi = 0; mi < size<0>(accum_p_mn); ++mi) {
+            if constexpr (!KernelTraits::IsCausal) {
+              accum_p_mn(mi, ni) = is_oob_token ? -mutlass::platform::numeric_limits<typename KernelTraits::ElementAccum>::infinity() : accum_p_mn(mi ,ni);
+            } else {
+              int row = get<0>(tPcP_mn(mi, 0));
+              // int col_limit_right = cur_batch_seq_len - cur_tile_idx * KernelTraits::BlockN - (params.q_seq_per_hkv - 1 - (m_coord * KernelTraits::BlockM + row)) / params.h_r;
+              int col_limit_right = (m_coord * KernelTraits::BlockM + row) / params.h_r + cur_batch_seq_len - params.q_seq_per_hkv / params.h_r - (cur_tile_idx) * KernelTraits::BlockN;
+              if (get<1>(tPcP_mn(0, ni)) > col_limit_right) {
+                accum_p_mn(mi, ni) = -mutlass::platform::numeric_limits<typename KernelTraits::ElementAccum>::infinity();
+              }
+            }
+          }
+        }
+      }
+    } // end loop
+
+    warpsquad_wait();
+    mp31_online_softmax_scale(accum_p, accum_softmax, params.sm_scale);
+    mp31_online_softmax_row_max<typename KernelTraits::ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+
+    //if (thread_idx_qk == 0 && b_coord == 0 && split_coord == 0 && h_q_coord == 0) {
+    //  printf("gpu qk acc:%f, row_max:%f\n", accum_softmax(0), row_max(0));
+    //}
+
+    online_softmax_exp2(accum_softmax);
+    mp31_online_softmax_row_sum<typename KernelTraits::ElementAccum>(accum_softmax, row_sum, lse, alpha);
+
+    //if (b_coord == 0 && split_coord == 0 && h_kv_coord == 0 && get<0>(tPcP_mn(0, 0)) == 0) {
+    //  printf("tid:%d v:%f %f %f %f %f %f %f %f\n", thread_idx_qk,
+    //                                               accum_softmax_mn(0, 0),
+    //                                               accum_softmax_mn(0, 1),
+    //                                               accum_softmax_mn(0, 2),
+    //                                               accum_softmax_mn(0, 3),
+    //                                               accum_softmax_mn(0, 4),
+    //                                               accum_softmax_mn(0, 5),
+    //                                               accum_softmax_mn(0, 6),
+    //                                               accum_softmax_mn(0, 7));
+    //}
+
+    // STS alpha & lse
+    pipeline_alpha.producer_acquire(pipeline_alpha_producer_state);
+
+    MUTE_UNROLL
+    for (int i = 0; i < size(alpha); ++i) {
+      sAlpha(get<0>(tPcP_mn(i, _0{})), pipeline_alpha_producer_state.index()) = alpha(i);
+    }
+
+    MUTE_UNROLL
+    for (int i = 0; i < size(lse); ++i) {
+      sLse(get<0>(tPcP_mn(i, _0{}))) = lse(i);
+    }
+
+    MUTE_UNROLL
+    for (int i = 0; i < size(tCvt_frg); ++i) {
+      tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::Element, typename KernelTraits::ElementAccum, KernelTraits::FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+    }
+
+    // STS P
+    pipeline_p.producer_acquire(pipeline_p_producer_state);
+    copy(tiled_copy_r2s, tPrP, tPsP);
+
+    __threadfence_block();
+    pipeline_p.producer_commit(pipeline_p_producer_state);
+    pipeline_alpha.producer_commit(pipeline_alpha_producer_state);
+
+    //// STG
+    //typename KernelTraits::ElementOutput* ptr_res = params.ptr_O + b_coord * params.stride_mid_ob + split_coord * params.stride_mid_os + h_q_coord * params.stride_mid_oh + KernelTraits::HeadDimVO;
+    typename KernelTraits::ElementOutput* ptr_res = params.ptr_O + b_coord * params.stride_mid_batch + split_coord * params.stride_mid_split + h_kv_coord * params.stride_mid_hkv + m_coord * KernelTraits::BlockM * params.stride_mid_seq + KernelTraits::HeadDimVO;
+
+    //if (get<0>(tPcP_mn(0, 0)) == 0) {
+    //  ptr_res[0] = row_max(0) * params.rescale + std::logf(lse(0));
+    //}
+    if (params.num_splits > 1) {
+      MUTE_UNROLL
+      for (int mi = 0; mi < size<0>(tPcP_mn); ++mi) {
+        int m_offset = get<0>(tPcP_mn(mi, 0)) * params.stride_mid_seq;
+        bool is_valid = m_coord * KernelTraits::BlockM + get<0>(tPcP_mn(mi, 0)) < params.q_seq_per_hkv;
+        if (is_valid) {
+          ptr_res[m_offset] = typename KernelTraits::ElementOutput(row_max(mi) * params.rescale + std::logf(lse(mi)));
+        }
+      }
+    }
+  }
+  // PV
+  else if (warp_squad_role == WarpSquadRole::ProducerConsumerPV) {
+
+    auto accum_o = partition_fragment_C(tiled_mma_pv, select<0, 1>(typename KernelTraits::TileShapePDV{}));
+
+    clear(accum_o);
+
+    Tensor tOcO = thr_mma_pv.partition_C(make_identity_tensor(select<0,1>(typename KernelTraits::TileShapePDV{})));
+    Tensor tOcO_mn = make_tensor(tOcO.data(), mp31_convert_layout_acc_rowcol(tOcO.layout()));
+
+    Tensor mQ = params.tme_load_Q.get_tme_tensor(params.shape_q);
+    auto cta_tme_q = params.tme_load_Q.get_slice(0);
+    Tensor gQ = local_tile(mQ, select<0, 2>(typename KernelTraits::TileShapeQKD{}), make_coord(m_coord, _0{}, h_kv_coord, b_coord));
+
+    Tensor tQgQ = cta_tme_q.partition_S(gQ);
+    Tensor tQsQ = cta_tme_q.partition_D(sQ);
+
+    if (warp_idx == 0) {
+      barrierQ.arrive_and_expect_tx(KernelTraits::TmeTransactionBytesQ);
+      uint32_t bar_id = barrierQ.get_barrier_id();
+      copy(params.tme_load_Q.with(bar_id), tQgQ, tQsQ);
+    }
+
+    // (page, d, h, num_blocks)
+    Tensor mV = params.tme_load_V.get_tme_tensor(params.shape_v);
+    auto cta_tme_v = params.tme_load_V.get_slice(0);
+
+    static_assert(KernelTraits::PageSize == KernelTraits::BlockN);
+    // TODO: tme dim should meet the page size
+    Tensor gV_pages = local_tile(mV, select<2, 1>(typename KernelTraits::TileShapePDV{}), make_coord(_0{}, _0{}, h_kv_coord, _));
+
+    // Reverse to early process OOB
+    int cur_tile_idx = split_kv_end - 1;
+
+    int kv_page_number = block_table[cur_tile_idx];
+
+    bool is_oob_tile = (cur_tile_idx + 1) * KernelTraits::BlockN > cur_batch_seq_len;
+    int residue_seq = is_oob_tile ? cur_batch_seq_len - n_tiles * KernelTraits::BlockN: 0;
+
+    Tensor tVsV = cta_tme_v.partition_D(sV);
+
+    // last tile
+    {
+      Tensor gV = gV_pages(_, _, kv_page_number);
+
+      Tensor tVgV = cta_tme_v.partition_S(gV);
+
+      if (warp_idx == 0) {
+        pipeline_v.producer_acquire(pipeline_v_producer_state);
+        uint32_t bar_id = pipeline_v.producer_get_barrier_id(pipeline_v_producer_state);
+        copy(params.tme_load_V.with(bar_id), tVgV, tVsV(_,_,_,pipeline_v_producer_state.index()));
+        ++pipeline_v_producer_state;
+      }
+
+      //if (thread_idx == 0 && h_q_coord == 0 && b_coord == 0) {
+      //  printf("split:%d cur_tile:%d kv_page_num:%d\n", split_coord, cur_tile_idx, kv_page_number);
+      //  //print("mV:"); print(mV); print("\n");
+      //  //print("gV_pages:"); print(gV_pages); print("\n");
+      //  print("gV:"); print(gV); print("\n");
+      //  //print("tVgV:"); print(tVgV); print("\n");
+      //  //print("gV_oob:"); print(gV_oob); print("\n");
+      //}
+    }
+
+    // LDS alpha
+    pipeline_alpha.consumer_wait(pipeline_alpha_consumer_state);
+    MUTE_UNROLL
+    for (int i = 0; i < size(alpha); ++i) {
+      alpha(i) = sAlpha(get<0>(tOcO_mn(i, _0{})), pipeline_alpha_consumer_state.index());
+    }
+
+    mp31_update_accum(accum_o, alpha);
+
+    pipeline_alpha.consumer_release(pipeline_alpha_consumer_state);
+    ++pipeline_alpha_consumer_state;
+
+    pipeline_p.consumer_wait(pipeline_p_consumer_state);
+    pipeline_v.consumer_wait(pipeline_v_consumer_state);
+
+
+    mute::gemm(tiled_mma_pv, tCrP, tCrV(_,_,_,pipeline_v_consumer_state.index()), accum_o);
+
+    --cur_tile_idx;
+    for (; cur_tile_idx >= split_kv_start; --cur_tile_idx) {
+      kv_page_number = block_table[cur_tile_idx];
+
+      if constexpr (KernelTraits::VStages == 1) {
+        warpsquad_wait();
+        pipeline_v.consumer_release(pipeline_v_consumer_state);
+        ++pipeline_v_consumer_state;
+        pipeline_p.consumer_release(pipeline_p_consumer_state);
+        ++pipeline_p_consumer_state;
+      }
+
+      {
+        Tensor gV = gV_pages(_, _, kv_page_number);
+        Tensor tVgV = cta_tme_v.partition_S(gV);
+
+        if (warp_idx == 0) {
+          pipeline_v.producer_acquire(pipeline_v_producer_state);
+          uint32_t bar_id = pipeline_v.producer_get_barrier_id(pipeline_v_producer_state);
+          copy(params.tme_load_V.with(bar_id), tVgV, tVsV(_,_,_,pipeline_v_producer_state.index()));
+          ++pipeline_v_producer_state;
+        }
+        //if (thread_idx == 0 && h_q_coord == 0 && b_coord == 0) {
+        //  printf("split:%d cur_tile:%d kv_page_num:%d\n", split_coord, cur_tile_idx, kv_page_number);
+        //  //print("mV:"); print(mV); print("\n");
+        //  //print("gV_pages:"); print(gV_pages); print("\n");
+        //  print("gV:"); print(gV); print("\n");
+        //  //print("tVgV:"); print(tVgV); print("\n");
+        //  //print("gV_oob:"); print(gV_oob); print("\n");
+        //}
+      }
+
+      if constexpr (KernelTraits::VStages > 1) {
+        warpsquad_wait();
+        pipeline_v.consumer_release(pipeline_v_consumer_state);
+        ++pipeline_v_consumer_state;
+        pipeline_p.consumer_release(pipeline_p_consumer_state);
+        ++pipeline_p_consumer_state;
+      }
+
+      // LDS alpha
+      pipeline_alpha.consumer_wait(pipeline_alpha_consumer_state);
+      MUTE_UNROLL
+      for (int i = 0; i < size(alpha); ++i) {
+        alpha(i) = sAlpha(get<0>(tOcO_mn(i, _0{})), pipeline_alpha_consumer_state.index());
+      }
+
+      mp31_update_accum(accum_o, alpha);
+
+      pipeline_alpha.consumer_release(pipeline_alpha_consumer_state);
+      ++pipeline_alpha_consumer_state;
+
+      pipeline_p.consumer_wait(pipeline_p_consumer_state);
+      pipeline_v.consumer_wait(pipeline_v_consumer_state);
+
+      mute::gemm(tiled_mma_pv, tCrP, tCrV(_,_,_,pipeline_v_consumer_state.index()), accum_o);
+
+    }
+
+    warpsquad_wait();
+
+    Tensor accum_o_mn = make_tensor(accum_o.data(), mp31_convert_layout_acc_rowcol(accum_o.layout()));
+
+    // LDS lse
+    MUTE_UNROLL
+    for (int i = 0; i < size(lse); ++i) {
+      lse(i) = sLse(get<0>(tOcO_mn(i, _0{})));
+    }
+
+    //if (b_coord == 0 && split_coord == 0 && h_q_coord == 0) {
+    //  printf("tid:%d p:%f v:%f\n", thread_idx, float(sP(thread_idx)), float(sV(thread_idx)));
+    //}
+
+
+    mp31_apply_softmax_normalizer(accum_o, lse);
+
+    //if (b_coord == 0 && split_coord == 0 && h_kv_coord == 0 && get<0>(tOcO_mn(0, 0)) == 0) {
+    //  printf("tid:%d lse:%f v:%f %f %f %f %f %f %f %f\n", thread_idx,
+    //                                               lse(0),
+    //                                               accum_o_mn(0, 0),
+    //                                               accum_o_mn(0, 1),
+    //                                               accum_o_mn(0, 2),
+    //                                               accum_o_mn(0, 3),
+    //                                               accum_o_mn(0, 4),
+    //                                               accum_o_mn(0, 5),
+    //                                               accum_o_mn(0, 6),
+    //                                               accum_o_mn(0, 7));
+    //}
+
+    // STG
+    typename KernelTraits::ElementOutput* ptr_res = params.ptr_O + b_coord * params.stride_mid_batch + split_coord * params.stride_mid_split + h_kv_coord * params.stride_mid_hkv + m_coord * KernelTraits::BlockM * params.stride_mid_seq;
+
+    MUTE_UNROLL
+    for (int mi = 0; mi < size<0>(accum_o_mn); ++mi) {
+      int m_offset = get<0>(tOcO_mn(mi, 0)) * params.stride_mid_seq;
+      bool is_valid = m_coord * KernelTraits::BlockM + get<0>(tOcO_mn(mi, 0)) < params.q_seq_per_hkv;
+      if (is_valid) {
+        MUTE_UNROLL
+        for (int ni = 0; ni < size<1>(accum_o_mn); ++ni) {
+          int offset = m_offset + get<1>((tOcO_mn(0, ni)));
+          ptr_res[offset] = typename KernelTraits::ElementOutput(accum_o_mn(mi, ni));
+        }
+      }
+    }
+  }
+  #endif
+}
diff --git a/vllm_musa/csrc/decode_mla/mp31_impl/fmha_decode.hpp b/vllm_musa/csrc/decode_mla/mp31_impl/fmha_decode.hpp
new file mode 100644
index 000000000..f87cb5fd6
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp31_impl/fmha_decode.hpp
@@ -0,0 +1,559 @@
+#pragma once
+
+#include "kernel_tratis.hpp"
+#include "./online_softmax.hpp"
+
+#include <mutlass/numeric_conversion.h>
+
+template <class KernelTraits>
+__global__ __launch_bounds__(KernelTraits::NumThreads, 1)
+void fmha_decode(typename KernelTraits::Params const params) {
+  #if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 310)
+  __shared__ char __attribute__((aligned(256))) smem_buf[sizeof(typename KernelTraits::SharedStorage)];
+  auto& shared_storage = *reinterpret_cast<typename KernelTraits::SharedStorage*>(smem_buf);
+
+  auto [b_coord, h_q_coord, split_coord] = static_cast<uint3>(blockIdx);
+
+  auto h_kv_coord = h_q_coord / get<4, 0>(params.problem_shape);
+
+
+  int cur_batch_seq_len = params.ptr_seqlen[b_coord];
+
+  constexpr int BlockN = KernelTraits::BlockN;
+
+  int n_tiles = ceil_div(cur_batch_seq_len, BlockN);
+  int n_tiles_per_split = ceil_div(n_tiles, params.num_splits);
+
+  int split_kv_start = n_tiles_per_split * split_coord;
+  int split_kv_end   = min(split_kv_start + n_tiles_per_split, n_tiles);
+
+  // early return
+  if (split_kv_end <= split_kv_start) return;
+
+  int thread_idx = threadIdx.x;
+
+  //if (thread_idx == 0 && h_q_coord == 0) {
+  //  printf("batch_id:%d split_id:%d cur_seqlen:%d, n_tiles:%d split:(%d %d)\n", b_coord, split_coord, cur_batch_seq_len, n_tiles, split_kv_start, split_kv_end);
+  //}
+
+  const int* block_table = params.ptr_block_table + b_coord * params.stride_block_table_b;
+
+  Tensor sQ  = make_tensor(make_smem_ptr(shared_storage.smem_q.data()),  typename KernelTraits::SmemLayoutQ{});
+  Tensor sK  = make_tensor(make_smem_ptr(shared_storage.smem_k.data()),  typename KernelTraits::SmemLayoutK{});
+  Tensor sP  = make_tensor(make_smem_ptr(shared_storage.smem_p.data()),  typename KernelTraits::SmemLayoutP{});
+  Tensor sV  = make_tensor(make_smem_ptr(shared_storage.smem_v.data()),  typename KernelTraits::SmemLayoutV{});
+  Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_vt.data()), typename KernelTraits::SmemLayoutVt{});
+  Tensor sAlpha = make_tensor(make_smem_ptr(shared_storage.smem_score.data()), typename KernelTraits::SmemLayoutScore{});
+  Tensor sLse = make_tensor(make_smem_ptr(shared_storage.smem_lse.data()), typename KernelTraits::SmemLayoutLse{});
+
+
+  typename KernelTraits::TiledMmaQK tiled_mma_qk;
+  typename KernelTraits::TiledMmaPV tiled_mma_pv;
+  typename KernelTraits::TiledMmaQKPermute tiled_mma_qk_perm;
+
+  auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx % KernelTraits::NumThreadsQK);
+  auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx % KernelTraits::NumThreadsPV);
+  auto thr_mma_qk_perm = tiled_mma_qk_perm.get_thread_slice(thread_idx % KernelTraits::NumThreadsQK);
+
+  // MMA QK
+  Tensor tCsQ = thr_mma_qk.partition_A(sQ);
+  Tensor tCrQ = thr_mma_qk.make_fragment_A(tCsQ);
+
+  Tensor tCsK = thr_mma_qk.partition_B(sK);
+  Tensor tCrK = thr_mma_qk.make_fragment_B(tCsK);
+
+  // MMA PV
+  Tensor tCsP = thr_mma_pv.partition_A(sP);
+  Tensor tCrP = thr_mma_pv.make_fragment_A(tCsP);
+
+  Tensor tCsV = thr_mma_pv.partition_B(sVt);
+  Tensor tCrV = thr_mma_pv.make_fragment_B(tCsV);
+
+  // R2S
+  using PermuteAlignmentType = mute::uint_bit_t<sizeof_bits_v<typename KernelTraits::Element> * KernelTraits::FragmentSize>;
+
+  auto tiled_copy_r2s = make_tiled_copy_C(Copy_Atom<UniversalCopy<PermuteAlignmentType>, typename KernelTraits::Element>{},
+                                          typename KernelTraits::TiledMmaQKPermute{});
+  auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(thread_idx % KernelTraits::NumThreadsQK);
+  Tensor tPsP = thr_copy_r2s.partition_D(sP);
+
+  constexpr int RowsPerThread = KernelTraits::BlockM / (KernelTraits::NumThreadsQK / 128) / 16;
+
+  Tensor alpha       = make_tensor<typename KernelTraits::ElementAccum>(Shape<Int<RowsPerThread>>{});
+  Tensor lse         = make_fragment_like(alpha);
+  Tensor row_sum     = make_fragment_like(alpha);
+  Tensor row_max     = make_fragment_like(alpha);
+  Tensor old_row_max = make_fragment_like(alpha);
+
+
+  fill(lse, typename KernelTraits::ElementAccum(0.0f));
+  fill(alpha, typename KernelTraits::ElementAccum(0.0f));
+  fill(old_row_max, -mutlass::platform::numeric_limits<typename KernelTraits::ElementAccum>::infinity());
+
+  constexpr int NumWarpsQK = KernelTraits::NumThreadsQK / mutlass::NumThreadsPerWarp;
+  constexpr int NumWarpsPV = KernelTraits::NumThreadsPV / mutlass::NumThreadsPerWarp;
+
+  int warp_idx = mutlass::canonical_warp_idx();
+  int warp_squad_idx = mutlass::canonical_warp_squad_idx();
+
+  enum class WarpSquadRole {
+    ProducerConsumerPV,
+    ProducerConsumerQK,
+  };
+
+  auto warp_squad_role = WarpSquadRole(warp_squad_idx);
+
+  //bool is_producer_consumer_pv = warp_idx < NumWarpsPV;
+  //bool is_producer_consumer_qk = !is_producer_consumer_pv;
+
+  mutlass::arch::allocate_async_barriers(sizeof(typename KernelTraits::BarrierStorage));
+  typename KernelTraits::BarrierStorage* barrier_ids = reinterpret_cast<typename KernelTraits::BarrierStorage*>(0);
+
+  typename KernelTraits::BarrierQ barrierQ(reinterpret_cast<uint64_t>(&barrier_ids->BarrierQ));
+
+  if (warp_idx == 0) {
+    barrierQ.init(1);
+  }
+
+  typename KernelTraits::MainloopPipelineK::Params k_pipe_params;
+  k_pipe_params.producer_arv_count = NumWarpsQK;
+  k_pipe_params.consumer_arv_count = NumWarpsPV;
+  typename KernelTraits::MainloopPipelineK pipeline_k(k_pipe_params, reinterpret_cast<uint64_t>(&barrier_ids->PipelineK));
+
+  typename KernelTraits::MainloopPipelineK::PipelineState pipeline_k_producer_state = mutlass::make_producer_start_state<typename KernelTraits::MainloopPipelineK>();
+  typename KernelTraits::MainloopPipelineK::PipelineState pipeline_k_consumer_state;
+
+
+  typename KernelTraits::MainloopPipelineV::Params v_pipe_params;
+  v_pipe_params.transaction_bytes = KernelTraits::TmeTransactionBytesV;
+  v_pipe_params.num_consumers = NumWarpsPV;
+  v_pipe_params.num_producers = 1;
+  typename KernelTraits::MainloopPipelineV pipeline_v(v_pipe_params, reinterpret_cast<uint64_t>(&barrier_ids->PipelineV));
+
+  typename KernelTraits::MainloopPipelineV::PipelineState pipeline_v_producer_state = mutlass::make_producer_start_state<typename KernelTraits::MainloopPipelineV>();
+  typename KernelTraits::MainloopPipelineV::PipelineState pipeline_v_consumer_state;
+
+  typename KernelTraits::MainloopPipelineP::Params p_pipe_params;
+  p_pipe_params.producer_arv_count = NumWarpsQK;
+  p_pipe_params.consumer_arv_count = NumWarpsPV;
+  typename KernelTraits::MainloopPipelineP pipeline_p(p_pipe_params, reinterpret_cast<uint64_t>(&barrier_ids->PipelineP));
+
+  typename KernelTraits::MainloopPipelineP::PipelineState pipeline_p_producer_state = mutlass::make_producer_start_state<typename KernelTraits::MainloopPipelineP>();
+  typename KernelTraits::MainloopPipelineP::PipelineState pipeline_p_consumer_state;
+
+  typename KernelTraits::MainloopPipelineAlpha::Params alpha_pipe_params;
+  alpha_pipe_params.producer_arv_count = NumWarpsQK;
+  alpha_pipe_params.consumer_arv_count = NumWarpsPV;
+  typename KernelTraits::MainloopPipelineAlpha pipeline_alpha(alpha_pipe_params, reinterpret_cast<uint64_t>(&barrier_ids->PipelineAlpha));
+
+  typename KernelTraits::MainloopPipelineAlpha::PipelineState pipeline_alpha_producer_state = mutlass::make_producer_start_state<typename KernelTraits::MainloopPipelineAlpha>();
+  typename KernelTraits::MainloopPipelineAlpha::PipelineState pipeline_alpha_consumer_state;
+
+  __syncthreads();
+
+  using TileShapeQKD = typename KernelTraits::TileShapeQKD;
+  using TileShapePDV = typename KernelTraits::TileShapePDV;
+
+  // QK
+  if (warp_squad_role == WarpSquadRole::ProducerConsumerQK) {
+    auto accum_p = partition_fragment_C(tiled_mma_qk, select<0, 1>(typename KernelTraits::TileShapeQKD{}));
+    auto accum_softmax = make_fragment_like(accum_p);
+    auto accum_cvt = make_fragment_like<typename KernelTraits::Element>(accum_softmax);
+    clear(accum_p);
+
+    Tensor tPcP = thr_mma_qk_perm.partition_C(make_identity_tensor(select<0, 1>(typename KernelTraits::TileShapeQKD{})));
+    Tensor tPcP_mn = make_tensor(tPcP.data(), mp31_convert_layout_acc_rowcol(tPcP.layout()));
+    Tensor accum_p_mn = make_tensor(accum_p.data(), mp31_convert_layout_acc_rowcol(accum_p.layout()));
+    Tensor accum_softmax_mn = make_tensor(accum_softmax.data(), mp31_convert_layout_acc_rowcol(accum_softmax.layout()));
+
+    Tensor tPrP = thr_copy_r2s.retile_S(accum_cvt);
+    Tensor tCvt_frg = recast<mutlass::Array<typename KernelTraits::Element, KernelTraits::FragmentSize>>(accum_cvt);
+    Tensor tAcc_frg = recast<mutlass::Array<typename KernelTraits::ElementAccum, KernelTraits::FragmentSize>>(accum_softmax);
+
+    // Assume PageSize == TileN
+    // TODO: relax it
+    static_assert(KernelTraits::PageSize == KernelTraits::BlockN);
+
+    constexpr int PagesPerBlockN = KernelTraits::BlockN / KernelTraits::PageSize;
+
+    // Reverse to early process OOB
+    int cur_tile_idx = split_kv_end - 1;
+
+    int kv_page_number = block_table[cur_tile_idx];
+
+    TiledCopy gmem_tiled_copy_K = typename KernelTraits::GmemTiledCopyK{};
+    ThrCopy gmem_thr_copy_K = gmem_tiled_copy_K.get_thread_slice(thread_idx % KernelTraits::NumThreadsQK);
+
+    Tensor cK = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK)));
+    Tensor tKcK = gmem_thr_copy_K.partition_S(cK);
+    Tensor tKpK = make_tensor<bool>(make_shape(size<1>(tKcK), size<2>(tKcK)), Stride<_1, _0>{});
+
+    Tensor tKsK = gmem_thr_copy_K.partition_D(sK);
+
+    MUTE_UNROLL
+    for (int i = 0; i < size<0>(tKpK); ++i) {
+      tKpK(i) = cur_tile_idx * KernelTraits::BlockN + get<0>(tKcK(0, i, 0)) < cur_batch_seq_len;
+    }
+
+    Tensor mK = make_tensor(make_gmem_ptr(params.ptr_K), params.layout_K);
+    // TODO:permute tile maybe less than block tile
+    Tensor gK_pages = local_tile(mK, typename KernelTraits::LdgstsTileK{}, make_coord(_0{}, _0{}, h_kv_coord, _)); // (PageSizeTile, HeadDimTile, pages)
+
+    bool is_oob_tile = (cur_tile_idx + 1) * KernelTraits::BlockN > cur_batch_seq_len;
+
+    {
+      Tensor gK = gK_pages(_, _, kv_page_number);
+      Tensor tKgK = gmem_thr_copy_K.partition_S(gK);
+
+      pipeline_k.producer_acquire(pipeline_k_producer_state);
+
+      copy(gmem_tiled_copy_K.with(params.desc_K), tKgK, tKsK(_,_,_,pipeline_k_producer_state.index()));
+
+
+      //if (thread_idx == KernelTraits::NumThreadsPV && b_coord == 0 && split_coord == 0 && h_q_coord == 0) {
+      //  print("mK:"); print(mK); print("\n");
+      //  print("gK_pages:"); print(gK_pages); print("\n");
+      //  print("gK:"); print(gK); print("\n");
+      //  print("tKgK:"); print(tKgK); print("\n");
+      //}
+    }
+
+    ldgsts_wait();
+    pipeline_k.producer_commit(pipeline_k_producer_state);
+    ++pipeline_k_producer_state;
+
+    barrierQ.wait(0);
+    pipeline_k.consumer_wait(pipeline_k_consumer_state);
+
+    mute::gemm(tiled_mma_qk, tCrQ, tCrK(_,_,_,pipeline_k_consumer_state.index()), accum_p);
+
+
+    warpsquad_wait();
+
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_p_mn); ++ni) {
+      //bool is_oob_token = is_oob_tile && (cur_tile_idx * KernelTraits::BlockN + get<1>(tPcP_mn(0, ni))) >= cur_batch_seq_len;
+      bool is_oob_token = get<1>(tPcP_mn(0, ni)) >= (cur_batch_seq_len - cur_tile_idx * KernelTraits::BlockN);
+      MUTE_UNROLL
+      for (int mi = 0; mi < size<0>(accum_p_mn); ++mi) {
+        accum_p_mn(mi, ni) = is_oob_token ? -mutlass::platform::numeric_limits<typename KernelTraits::ElementAccum>::infinity() : accum_p_mn(mi ,ni);
+      }
+    }
+
+    int thread_idx_qk = thread_idx - KernelTraits::NumThreadsPV;
+
+
+    --cur_tile_idx;
+
+    for (; cur_tile_idx >= split_kv_start; --cur_tile_idx) {
+      if constexpr (KernelTraits::KStages == 1) {
+        warpsquad_wait();
+        pipeline_k.consumer_release(pipeline_k_consumer_state);
+        ++pipeline_k_consumer_state;
+      }
+      // load K
+      {
+        kv_page_number = block_table[cur_tile_idx];
+        Tensor gK = gK_pages(_, _, kv_page_number);
+
+        Tensor tKgK = gmem_thr_copy_K.partition_S(gK);
+
+        pipeline_k.producer_acquire(pipeline_k_producer_state);
+
+        copy(gmem_tiled_copy_K.with(params.desc_K), tKgK, tKsK(_,_,_,pipeline_k_producer_state.index()));
+      }
+
+      if constexpr (KernelTraits::KStages > 1) {
+        warpsquad_wait();
+        pipeline_k.consumer_release(pipeline_k_consumer_state);
+        ++pipeline_k_consumer_state;
+      }
+
+      mp31_online_softmax_scale(accum_p, accum_softmax, params.sm_scale);
+      mp31_online_softmax_row_max<typename KernelTraits::ElementAccum>(accum_softmax, row_max, old_row_max, alpha); // alpha is updated
+
+      clear(accum_p);
+
+      ldgsts_wait();
+      pipeline_k.producer_commit(pipeline_k_producer_state);
+      ++pipeline_k_producer_state;
+
+
+      pipeline_k.consumer_wait(pipeline_k_consumer_state);
+
+      mute::gemm(tiled_mma_qk, tCrQ, tCrK(_,_,_,pipeline_k_consumer_state.index()), accum_p);
+
+      online_softmax_exp2(accum_softmax);
+      mp31_online_softmax_row_sum<typename KernelTraits::ElementAccum>(accum_softmax, row_sum, lse, alpha);
+
+      // STS alpha
+      pipeline_alpha.producer_acquire(pipeline_alpha_producer_state);
+
+      MUTE_UNROLL
+      for (int i = 0; i < size(alpha); ++i) {
+        sAlpha(get<0>(tPcP_mn(i, _0{})), pipeline_alpha_producer_state.index()) = alpha(i);
+      }
+
+      MUTE_UNROLL
+      for (int i = 0; i < size(tCvt_frg); ++i) {
+        tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::Element, typename KernelTraits::ElementAccum, KernelTraits::FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+      }
+
+      // STS P
+      pipeline_p.producer_acquire(pipeline_p_producer_state);
+      copy(tiled_copy_r2s, tPrP, tPsP);
+
+      __threadfence_block();
+      pipeline_p.producer_commit(pipeline_p_producer_state);
+      pipeline_alpha.producer_commit(pipeline_alpha_producer_state);
+      ++pipeline_p_producer_state;
+      ++pipeline_alpha_producer_state;
+    } // end loop
+
+    warpsquad_wait();
+    mp31_online_softmax_scale(accum_p, accum_softmax, params.sm_scale);
+    mp31_online_softmax_row_max<typename KernelTraits::ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+
+    //if (thread_idx_qk == 0 && b_coord == 0 && split_coord == 0 && h_q_coord == 0) {
+    //  printf("gpu qk acc:%f, row_max:%f\n", accum_softmax(0), row_max(0));
+    //}
+
+    online_softmax_exp2(accum_softmax);
+    mp31_online_softmax_row_sum<typename KernelTraits::ElementAccum>(accum_softmax, row_sum, lse, alpha);
+
+    //if (b_coord == 0 && split_coord == 0 && h_q_coord == 0 && get<0>(tPcP_mn(0, 0)) == 0) {
+    //  printf("tid:%d v:%f %f %f %f %f %f %f %f\n", thread_idx_qk,
+    //                                               accum_softmax_mn(0, 0),
+    //                                               accum_softmax_mn(0, 1),
+    //                                               accum_softmax_mn(0, 2),
+    //                                               accum_softmax_mn(0, 3),
+    //                                               accum_softmax_mn(0, 4),
+    //                                               accum_softmax_mn(0, 5),
+    //                                               accum_softmax_mn(0, 6),
+    //                                               accum_softmax_mn(0, 7));
+    //}
+
+    // STS alpha & lse
+    pipeline_alpha.producer_acquire(pipeline_alpha_producer_state);
+
+    MUTE_UNROLL
+    for (int i = 0; i < size(alpha); ++i) {
+      sAlpha(get<0>(tPcP_mn(i, _0{})), pipeline_alpha_producer_state.index()) = alpha(i);
+    }
+
+    MUTE_UNROLL
+    for (int i = 0; i < size(lse); ++i) {
+      sLse(get<0>(tPcP_mn(i, _0{}))) = lse(i);
+    }
+
+    MUTE_UNROLL
+    for (int i = 0; i < size(tCvt_frg); ++i) {
+      tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::Element, typename KernelTraits::ElementAccum, KernelTraits::FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+    }
+
+    // STS P
+    pipeline_p.producer_acquire(pipeline_p_producer_state);
+    copy(tiled_copy_r2s, tPrP, tPsP);
+
+    __threadfence_block();
+    pipeline_p.producer_commit(pipeline_p_producer_state);
+    pipeline_alpha.producer_commit(pipeline_alpha_producer_state);
+
+    // STG
+    typename KernelTraits::ElementOutput* ptr_res = params.ptr_O + b_coord * params.stride_mid_ob + split_coord * params.stride_mid_os + h_q_coord * params.stride_mid_oh + KernelTraits::HeadDimVO;
+
+    if (get<0>(tPcP_mn(0, 0)) == 0) {
+      ptr_res[0] = row_max(0) * params.rescale + std::logf(lse(0));
+    }
+  }
+  // PV
+  else if (warp_squad_role == WarpSquadRole::ProducerConsumerPV) {
+
+    auto accum_o = partition_fragment_C(tiled_mma_pv, select<0, 1>(typename KernelTraits::TileShapePDV{}));
+
+    clear(accum_o);
+
+    Tensor tOcO = thr_mma_pv.partition_C(make_identity_tensor(select<0,1>(typename KernelTraits::TileShapePDV{})));
+    Tensor tOcO_mn = make_tensor(tOcO.data(), mp31_convert_layout_acc_rowcol(tOcO.layout()));
+
+    Tensor mQ = params.tme_load_Q.get_tme_tensor(params.shape_q);
+    auto cta_tme_q = params.tme_load_Q.get_slice(0);
+    Tensor gQ = local_tile(mQ, select<0, 2>(typename KernelTraits::TileShapeQKD{}), make_coord(_0{}, _0{}, h_q_coord, b_coord));
+
+    Tensor tQgQ = cta_tme_q.partition_S(gQ);
+    Tensor tQsQ = cta_tme_q.partition_D(sQ);
+
+    if (warp_idx == 0) {
+      barrierQ.arrive_and_expect_tx(KernelTraits::TmeTransactionBytesQ);
+      uint32_t bar_id = barrierQ.get_barrier_id();
+      copy(params.tme_load_Q.with(bar_id), tQgQ, tQsQ);
+    }
+
+    // (page, d, h, num_blocks)
+    Tensor mV = params.tme_load_V.get_tme_tensor(params.shape_v);
+    auto cta_tme_v = params.tme_load_V.get_slice(0);
+
+    static_assert(KernelTraits::PageSize == KernelTraits::BlockN);
+    // TODO: tme dim should meet the page size
+    Tensor gV_pages = local_tile(mV, select<2, 1>(typename KernelTraits::TileShapePDV{}), make_coord(_0{}, _0{}, h_kv_coord, _));
+
+    // Reverse to early process OOB
+    int cur_tile_idx = split_kv_end - 1;
+
+    int kv_page_number = block_table[cur_tile_idx];
+
+    bool is_oob_tile = (cur_tile_idx + 1) * KernelTraits::BlockN > cur_batch_seq_len;
+    int residue_seq = is_oob_tile ? cur_batch_seq_len - n_tiles * KernelTraits::BlockN: 0;
+
+    Tensor tVsV = cta_tme_v.partition_D(sV);
+
+    // last tile
+    {
+      Tensor gV = gV_pages(_, _, kv_page_number);
+
+      Tensor tVgV = cta_tme_v.partition_S(gV);
+
+      if (warp_idx == 0) {
+        pipeline_v.producer_acquire(pipeline_v_producer_state);
+        uint32_t bar_id = pipeline_v.producer_get_barrier_id(pipeline_v_producer_state);
+        copy(params.tme_load_V.with(bar_id), tVgV, tVsV(_,_,_,pipeline_v_producer_state.index()));
+        ++pipeline_v_producer_state;
+      }
+
+      //if (thread_idx == 0 && h_q_coord == 0 && b_coord == 0) {
+      //  printf("split:%d cur_tile:%d kv_page_num:%d\n", split_coord, cur_tile_idx, kv_page_number);
+      //  //print("mV:"); print(mV); print("\n");
+      //  //print("gV_pages:"); print(gV_pages); print("\n");
+      //  print("gV:"); print(gV); print("\n");
+      //  //print("tVgV:"); print(tVgV); print("\n");
+      //  //print("gV_oob:"); print(gV_oob); print("\n");
+      //}
+    }
+
+    // LDS alpha
+    pipeline_alpha.consumer_wait(pipeline_alpha_consumer_state);
+    MUTE_UNROLL
+    for (int i = 0; i < size(alpha); ++i) {
+      alpha(i) = sAlpha(get<0>(tOcO_mn(i, _0{})), pipeline_alpha_consumer_state.index());
+    }
+
+    mp31_update_accum(accum_o, alpha);
+
+    pipeline_alpha.consumer_release(pipeline_alpha_consumer_state);
+    ++pipeline_alpha_consumer_state;
+
+    pipeline_p.consumer_wait(pipeline_p_consumer_state);
+    pipeline_v.consumer_wait(pipeline_v_consumer_state);
+
+
+    mute::gemm(tiled_mma_pv, tCrP, tCrV(_,_,_,pipeline_v_consumer_state.index()), accum_o);
+
+    --cur_tile_idx;
+    for (; cur_tile_idx >= split_kv_start; --cur_tile_idx) {
+      kv_page_number = block_table[cur_tile_idx];
+
+      if constexpr (KernelTraits::VStages == 1) {
+        warpsquad_wait();
+        pipeline_v.consumer_release(pipeline_v_consumer_state);
+        ++pipeline_v_consumer_state;
+        pipeline_p.consumer_release(pipeline_p_consumer_state);
+        ++pipeline_p_consumer_state;
+      }
+
+      {
+        Tensor gV = gV_pages(_, _, kv_page_number);
+        Tensor tVgV = cta_tme_v.partition_S(gV);
+
+        if (warp_idx == 0) {
+          pipeline_v.producer_acquire(pipeline_v_producer_state);
+          uint32_t bar_id = pipeline_v.producer_get_barrier_id(pipeline_v_producer_state);
+          copy(params.tme_load_V.with(bar_id), tVgV, tVsV(_,_,_,pipeline_v_producer_state.index()));
+          ++pipeline_v_producer_state;
+        }
+        //if (thread_idx == 0 && h_q_coord == 0 && b_coord == 0) {
+        //  printf("split:%d cur_tile:%d kv_page_num:%d\n", split_coord, cur_tile_idx, kv_page_number);
+        //  //print("mV:"); print(mV); print("\n");
+        //  //print("gV_pages:"); print(gV_pages); print("\n");
+        //  print("gV:"); print(gV); print("\n");
+        //  //print("tVgV:"); print(tVgV); print("\n");
+        //  //print("gV_oob:"); print(gV_oob); print("\n");
+        //}
+      }
+
+      if constexpr (KernelTraits::VStages > 1) {
+        warpsquad_wait();
+        pipeline_v.consumer_release(pipeline_v_consumer_state);
+        ++pipeline_v_consumer_state;
+        pipeline_p.consumer_release(pipeline_p_consumer_state);
+        ++pipeline_p_consumer_state;
+      }
+
+      // LDS alpha
+      pipeline_alpha.consumer_wait(pipeline_alpha_consumer_state);
+      MUTE_UNROLL
+      for (int i = 0; i < size(alpha); ++i) {
+        alpha(i) = sAlpha(get<0>(tOcO_mn(i, _0{})), pipeline_alpha_consumer_state.index());
+      }
+
+      mp31_update_accum(accum_o, alpha);
+
+      pipeline_alpha.consumer_release(pipeline_alpha_consumer_state);
+      ++pipeline_alpha_consumer_state;
+
+      pipeline_p.consumer_wait(pipeline_p_consumer_state);
+      pipeline_v.consumer_wait(pipeline_v_consumer_state);
+
+      mute::gemm(tiled_mma_pv, tCrP, tCrV(_,_,_,pipeline_v_consumer_state.index()), accum_o);
+
+    }
+
+    warpsquad_wait();
+
+    Tensor accum_o_mn = make_tensor(accum_o.data(), mp31_convert_layout_acc_rowcol(accum_o.layout()));
+
+    // LDS lse
+    MUTE_UNROLL
+    for (int i = 0; i < size(lse); ++i) {
+      lse(i) = sLse(get<0>(tOcO_mn(i, _0{})));
+    }
+
+    //if (b_coord == 0 && split_coord == 0 && h_q_coord == 0) {
+    //  printf("tid:%d p:%f v:%f\n", thread_idx, float(sP(thread_idx)), float(sV(thread_idx)));
+    //}
+
+
+    mp31_apply_softmax_normalizer(accum_o, lse);
+
+    //if (split_coord == 0 && h_q_coord == 0 && get<0>(tOcO_mn(0, 0)) == 0) {
+    //  printf("tid:%d bs:%d lse:%f v:%f %f %f %f %f %f %f %f\n", thread_idx, b_coord,
+    //                                               lse(0),
+    //                                               accum_o_mn(0, 0),
+    //                                               accum_o_mn(0, 1),
+    //                                               accum_o_mn(0, 2),
+    //                                               accum_o_mn(0, 3),
+    //                                               accum_o_mn(0, 4),
+    //                                               accum_o_mn(0, 5),
+    //                                               accum_o_mn(0, 6),
+    //                                               accum_o_mn(0, 7));
+    //}
+
+    // STG
+
+    // dst shape (d+1, h, s, b)
+    typename KernelTraits::ElementOutput* ptr_res = params.ptr_O + b_coord * params.stride_mid_ob + split_coord * params.stride_mid_os + h_q_coord * params.stride_mid_oh;
+
+    //if (b_coord == 0 && split_coord == 0 && h_q_coord == 0) {
+    //  print("tOcO_mn:"); print(tOcO_mn); print("\n");
+    //  print("accum_o_mn:"); print(accum_o_mn); print("\n");
+    //}
+
+    // now only mi = 0 is valid result
+    if (get<0>(tOcO_mn(0, 0)) == 0) {
+      MUTE_UNROLL
+      for (int ni = 0; ni < size<1>(accum_o_mn); ++ni) {
+        ptr_res[get<1>(tOcO_mn(_0{}, ni))] = accum_o_mn(_0{}, ni);
+      }
+    }
+  }
+  #endif
+}
diff --git a/vllm_musa/csrc/decode_mla/mp31_impl/kernel_tratis.hpp b/vllm_musa/csrc/decode_mla/mp31_impl/kernel_tratis.hpp
new file mode 100644
index 000000000..3230b342f
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp31_impl/kernel_tratis.hpp
@@ -0,0 +1,518 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mute/atom/mma_atom.hpp>
+#include <mute/atom/copy_atom.hpp>
+
+#include <mutlass/mutlass.h>
+#include <mutlass/gemm/gemm.h>
+#include <mutlass/gemm/collective/collective_builder.hpp>
+
+using namespace mute;
+
+
+template <class TileN, class PermuteTileShape>
+constexpr auto
+get_mp31_permute_tile_for_operand_B() {
+  using SqmmaAtomN = _8;
+  static_assert(TileN{} % PermuteTileShape{} == 0);
+  static_assert(PermuteTileShape{} % SqmmaAtomN{} == 0);
+  using RepeatInsidePermuteTile = decltype(PermuteTileShape{} / SqmmaAtomN{});
+  using RepeatAcrossPermuteTile = decltype(TileN{} / PermuteTileShape{});
+
+  using PermuteTileN = Layout<Shape <             SqmmaAtomN, RepeatInsidePermuteTile, RepeatAcrossPermuteTile>,
+                              Stride<RepeatInsidePermuteTile,                      _1,        PermuteTileShape>>;
+  return PermuteTileN{};
+}
+
+
+template <
+  int BlockM_,
+  int BlockN_,
+  int HeadDimQK_,
+  int HeadDimVO_,
+  int PageSize_,
+  class Element_,
+  class ElementOutput_,
+  int KStages_,
+  int VStages_,
+  bool IsCasual_ = false,
+  bool PackGQA_ = false
+>
+struct FmhaDecodeKernelTratis {
+  using Element = Element_;
+  using ElementAccum = float;
+
+  using ElementOutput = ElementOutput_;
+
+  static constexpr int HeadDimQK = HeadDimQK_;
+  static constexpr int HeadDimVO = HeadDimVO_;
+  static constexpr int BlockM  = BlockM_;
+  static constexpr int BlockN  = BlockN_;
+
+  static constexpr int PageSize = PageSize_;
+  static constexpr bool PackGQA = PackGQA_;
+  static constexpr bool IsCausal = IsCasual_;
+
+  // TODO: relax it
+  static_assert(PageSize % 64 == 0 && PageSize == BlockN);
+
+  static constexpr int KStages = KStages_;
+  static constexpr int VStages = VStages_;
+
+  using TileShapeQKD = Shape<Int<BlockM>, Int<BlockN>,    Int<HeadDimQK>>;
+  using TileShapePDV = Shape<Int<BlockM>, Int<HeadDimVO>, Int<BlockN>>;
+
+  using TiledMmaQK = decltype(mute::make_tiled_mma(
+      mute::MP31::SQMMA::ss_op_selector<Element, Element, ElementAccum, TileShapeQKD, TCE::Major::K, TCE::Major::K>()));
+
+  using TiledMmaPV = decltype(mute::make_tiled_mma(
+      mute::MP31::SQMMA::ss_op_selector<Element, Element, ElementAccum, TileShapePDV,
+                                        TCE::Major::K, TCE::Major::MN>()));
+
+  static constexpr int NumThreadsQK = size(TiledMmaQK{});
+  static constexpr int NumThreadsPV = size(TiledMmaPV{});
+  static constexpr int NumThreads = NumThreadsQK + NumThreadsPV;
+
+  static_assert(NumThreads == (NumThreadsQK + NumThreadsPV));
+
+  using SmemLayoutAtomQ = decltype(mute::MP31::SQMMA::Layout_SL256_SS256_SG16_Atom<
+                            get<0>(typename TiledMmaQK::AtomShape_MNK{}), get<2>(typename TiledMmaQK::AtomShape_MNK{}), Element, TCE::Major::K>());
+  using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, select<0, 2>(TileShapeQKD{})));
+
+  using SmemLayoutAtomK = decltype(mute::MP31::SQMMA::Layout_SL256_SS256_SG16_Atom<
+                            get<1>(typename TiledMmaQK::AtomShape_MNK{}), get<2>(typename TiledMmaQK::AtomShape_MNK{}), Element, TCE::Major::K>());
+  using SmemLayoutK = decltype(tile_to_shape(SmemLayoutAtomK{},
+                                             make_shape(get<1>(TileShapeQKD{}), get<2>(TileShapeQKD{}), Int<KStages>{})));
+
+  using SmemLayoutAtomV = decltype(mute::MP31::SQMMA::Layout_SL256_SS256_SG32_Atom<
+                            get<2>(typename TiledMmaPV::AtomShape_MNK{}), get<1>(typename TiledMmaPV::AtomShape_MNK{}), Element, TCE::Major::K>());
+
+  using SmemLayoutV = decltype(tile_to_shape(SmemLayoutAtomV{},
+                                             make_shape(get<2>(TileShapePDV{}), get<1>(TileShapePDV{}), Int<VStages>{})));
+
+  using SmemLayoutVt = decltype(composition(SmemLayoutV{},
+                                  make_ordered_layout(
+                                    make_shape(get<1>(TileShapePDV{}), get<2>(TileShapePDV{}), Int<VStages>{}),
+                                    Step<_2, _1, _3>{})));
+
+  using SmemLayoutAtomP = decltype(mute::MP31::SQMMA::Layout_SL256_SS256_SG16_Atom<
+                            get<0>(typename TiledMmaPV::AtomShape_MNK{}), get<2>(typename TiledMmaPV::AtomShape_MNK{}), Element, TCE::Major::K>());
+  using SmemLayoutP = decltype(tile_to_shape(SmemLayoutAtomP{}, select<0, 1>(TileShapeQKD{})));
+
+  using SmemLayoutScore = Layout<Shape<Int<BlockM>, Int<KStages>>>;
+
+  using SmemLayoutLse = Layout<Shape<Int<BlockM>>>;
+
+  static constexpr int FragmentSize = 128 / sizeof_bits_v<Element>;
+
+  // Permute to make STS happy
+  using GmemTiledCopyOpK = MP31_ROBUST_LDGSTS<mute::uint_bit_t<sizeof_bits_v<Element> * FragmentSize>>;
+  using GmemTiledCopyK =
+      decltype(mutlass::gemm::collective::detail::make_simt_tiled_copy<
+                  NumThreadsQK, Element, FragmentSize, Stride<int, _1> /* HeadDim-major */,
+                  shape<0>(SmemLayoutAtomK{}), shape<1>(SmemLayoutAtomK{}), GmemTiledCopyOpK>());
+
+  using GmemTiledCopyOpOut = MP31_ROBUST_LDGSTS<mute::uint_bit_t<128>>;
+  using GmemTiledCopyOut =
+      decltype(mutlass::gemm::collective::detail::make_simt_tiled_copy<
+                  NumThreads, ElementOutput, 128 / sizeof_bits_v<ElementOutput>, Stride<int, _1>,
+                  BlockM, HeadDimVO, GmemTiledCopyOpOut>());
+
+  using LdgstsTileK =
+      decltype(make_tile(get_mp31_permute_tile_for_operand_B<
+                            decltype(get<1>(TileShapeQKD{})), decltype(get<2>(typename TiledMmaQK::AtomShape_MNK{}))>(),
+                         get<2>(TileShapeQKD{})));
+
+  // used to get proper QKAcc layout
+  using TiledMmaQKPermute = decltype(mute::make_tiled_mma(
+      mute::MP31::SQMMA::ss_op_selector<Element, Element, ElementAccum, TileShapeQKD, TCE::Major::K, TCE::Major::K>(),
+      Layout<Shape<_1,_1,_1>>{},
+      make_tile(Underscore{},
+        get_mp31_permute_tile_for_operand_B<
+                    decltype(get<1>(TileShapeQKD{})), decltype(get<2>(typename TiledMmaPV::AtomShape_MNK{}))>(),
+                Underscore{})));
+
+
+  static constexpr uint32_t TmeTransactionBytesQ = static_cast<uint32_t>(size(SmemLayoutQ{}) * sizeof_bits_v<Element> / 8);
+  static constexpr uint32_t TmeTransactionBytesV = static_cast<uint32_t>(size(take<0,2>(SmemLayoutV{})) * sizeof_bits_v<Element> / 8);
+
+  struct SharedStorage {
+    mute::array_aligned<Element, mute::cosize_v<SmemLayoutQ>, 256> smem_q;
+
+    mute::array_aligned<Element, mute::cosize_v<SmemLayoutK>, 256> smem_k;
+
+    union {
+      mute::array_aligned<Element, mute::cosize_v<SmemLayoutV>, 256> smem_v;
+      mute::array_aligned<Element, mute::cosize_v<SmemLayoutVt>, 256> smem_vt;
+    };
+
+    mute::array_aligned<Element, mute::cosize_v<SmemLayoutP>, 256> smem_p;
+    mute::array_aligned<ElementAccum, mute::cosize_v<SmemLayoutScore>, 256> smem_score;
+    mute::array_aligned<ElementAccum, mute::cosize_v<SmemLayoutLse>, 256> smem_lse;
+  };
+
+  using MainloopPipelineK = mutlass::Mp31PipelineAsync<KStages>;
+  using MainloopPipelineV = mutlass::Mp31PipelineTmeAsync<VStages>;
+  using MainloopPipelineP = mutlass::Mp31PipelineAsync<1>;
+  using MainloopPipelineAlpha = mutlass::Mp31PipelineAsync<KStages>;
+
+  using BarrierQ = mutlass::arch::AsyncTransactionBarrier;
+
+  #pragma pack(push, 1)
+  struct BarrierStorage {
+    uint8_t BarrierQ;
+    uint8_t PipelineK[MainloopPipelineK::NumBarriers];
+    uint8_t PipelineV[MainloopPipelineV::NumBarriers];
+    uint8_t PipelineP[MainloopPipelineP::NumBarriers];
+    uint8_t PipelineAlpha[MainloopPipelineAlpha::NumBarriers];
+  };
+  #pragma pack(pop)
+
+  using IndexType = int32_t;
+
+  // Q, K, D_QK, D_VO, (H_R, H_KV), B
+  using ProblemShape = Shape<int, int, int, int, Shape<int, int>, int>;
+
+  // (0, d, h, b)
+  using StrideQ = Stride<int, _1, int, int>;
+  // (page size, d, h, num_blocks)
+  using StrideK = Stride<int, _1, int, int>;
+  // (page size, d, h, num_blocks)
+  using StrideV = StrideK;
+
+  using StrideO = StrideQ;
+
+  using ShapeQKVO = Shape <int, int, int, int>;
+  using LayoutK = Layout<ShapeQKVO, StrideK>;
+
+  using TME_Q = decltype(make_tme_copy(
+      MP31_TME_LOAD{},
+      make_tensor(
+        make_gmem_ptr(static_cast<Element const*>(nullptr)),
+        repeat_like(StrideQ{}, int32_t(0)),
+        StrideQ{}
+      ),
+      SmemLayoutQ{},
+      select<0, 2>(TileShapeQKD{})));
+
+
+  using TME_V = decltype(make_tme_copy(
+      MP31_TME_LOAD{},
+      make_tensor(
+        make_gmem_ptr(static_cast<Element const*>(nullptr)),
+        repeat_like(StrideV{}, int32_t(0)),
+        StrideV{}
+      ),
+      take<0, 2>(SmemLayoutV{}),
+      select<2, 1>(TileShapePDV{})));
+
+  struct Arguments {
+    // (1, d, h_q, b)
+    Element const* ptr_Q;
+    // (blocks, page_size, h_kv, d)
+    Element const* ptr_K_buffer;
+    // (blocks, page_size, h_kv, d)
+    Element const* ptr_V_buffer;
+
+    // (b, blocks/b)
+    int32_t const* ptr_block_table;
+    // (b,)
+    int32_t const* ptr_seqlen;
+
+    // (1, d, h_q, b)
+    ElementOutput* ptr_O;
+
+    float sm_scale;
+    int num_splits;
+    int num_blocks;
+
+    IndexType stride_block_table_b;
+    IndexType stride_qb;
+    IndexType stride_qh;
+    IndexType stride_buf_kblock;
+    IndexType stride_buf_kpage;
+    IndexType stride_buf_kh;
+    IndexType stride_buf_vblock;
+    IndexType stride_buf_vpage;
+    IndexType stride_buf_vh;
+    IndexType stride_mid_ob;
+    IndexType stride_mid_oh;
+    IndexType stride_mid_os;
+  };
+
+  struct Params {
+    ProblemShape problem_shape;
+
+    TME_Q tme_load_Q;
+    ShapeQKVO shape_q;
+
+    Element const* ptr_K;
+    RobustDescriptor desc_K;
+    LayoutK layout_K;
+
+    TME_V tme_load_V;
+    ShapeQKVO shape_v;
+
+    ElementOutput* ptr_O;
+
+    // (b, blocks/b)
+    int32_t const* ptr_block_table;
+    // (b,)
+    int32_t const* ptr_seqlen;
+
+    float const sm_scale;
+    float const rescale;
+
+    int const num_splits;
+
+    IndexType stride_block_table_b;
+    IndexType stride_qb;
+    IndexType stride_qh;
+    IndexType stride_buf_kblock;
+    IndexType stride_buf_kpage;
+    IndexType stride_buf_kh;
+    IndexType stride_buf_vblock;
+    IndexType stride_buf_vpage;
+    IndexType stride_buf_vh;
+    IndexType stride_mid_ob;
+    IndexType stride_mid_oh;
+    IndexType stride_mid_os;
+  };
+
+
+  static Params
+  to_underlying_arguments(ProblemShape problem_shape, Arguments const& args) {
+    int DimQK = size<2>(problem_shape);
+    int DimVO = size<3>(problem_shape);
+    int H = size<4>(problem_shape);
+    int H_KV = size<4, 1>(problem_shape);
+    int B = size<5>(problem_shape);
+
+    auto layout_Q = make_layout(make_shape(1, DimQK, H, B),
+                                make_stride(0, _1{}, args.stride_qh, args.stride_qb));
+    Tensor mQ = make_tensor(make_gmem_ptr(args.ptr_Q), layout_Q);
+
+    TME_Q tme_load_Q = make_tme_copy(
+        MP31_TME_LOAD{},
+        mQ,
+        SmemLayoutQ{},
+        select<0, 2>(TileShapeQKD{}));
+
+    RobustDescriptor desc_K = make_robust_desc(args.ptr_K_buffer, DimQK * H_KV * PageSize * args.num_blocks);
+
+    auto layout_K = make_layout(make_shape(PageSize, DimQK, H_KV, args.num_blocks),
+                                make_stride(args.stride_buf_kpage, _1{}, args.stride_buf_kh, args.stride_buf_kblock));
+
+    auto layout_V = make_layout(make_shape(PageSize, DimVO, H_KV, args.num_blocks),
+                                make_stride(args.stride_buf_vpage, _1{}, args.stride_buf_vh, args.stride_buf_vblock));
+    Tensor mV = make_tensor(make_gmem_ptr(args.ptr_V_buffer), layout_V);
+
+    TME_V tme_load_V = make_tme_copy(
+        MP31_TME_LOAD{},
+        mV,
+        take<0, 2>(SmemLayoutV{}),
+        select<2, 1>(TileShapePDV{}));
+
+    constexpr float kLog2e = float(1.4426950408889634074);
+    constexpr float rescale = 1.0f / kLog2e;
+
+    float sm_scale = args.sm_scale * kLog2e;
+
+    return {
+      problem_shape,
+      tme_load_Q,
+      shape(layout_Q),
+      args.ptr_K_buffer,
+      desc_K,
+      layout_K,
+      tme_load_V,
+      shape(layout_V),
+      args.ptr_O,
+      args.ptr_block_table,
+      args.ptr_seqlen,
+      sm_scale,
+      rescale,
+      args.num_splits,
+      args.stride_block_table_b,
+      args.stride_qb,
+      args.stride_qh,
+      args.stride_buf_kblock,
+      args.stride_buf_kpage,
+      args.stride_buf_kh,
+      args.stride_buf_vblock,
+      args.stride_buf_vpage,
+      args.stride_buf_vh,
+      args.stride_mid_ob,
+      args.stride_mid_oh,
+      args.stride_mid_os
+    };
+  }
+
+  struct PackGQAArguments {
+    // (q, d, h_q, b)
+    Element const* ptr_Q;
+    // (blocks, page_size, h_kv, d)
+    Element const* ptr_K_buffer;
+    // (blocks, page_size, h_kv, d)
+    Element const* ptr_V_buffer;
+
+    // (b, blocks/b)
+    int32_t const* ptr_block_table;
+    // (b,)
+    int32_t const* ptr_seqlen;
+
+    // (1, d, h_q, b)
+    ElementOutput* ptr_O;
+
+    float sm_scale;
+    int num_splits;
+    int num_blocks;
+
+    IndexType stride_block_table_b;
+
+    // TensorQ (batch_size, seqlen_q * h_r, h_kv, dim_qk)
+    IndexType stride_q_hkv;
+    IndexType stride_q_seq;
+    IndexType stride_q_bs;
+
+    IndexType stride_buf_kblock;
+    IndexType stride_buf_kpage;
+    IndexType stride_buf_kh;
+
+    IndexType stride_buf_vblock;
+    IndexType stride_buf_vpage;
+    IndexType stride_buf_vh;
+
+    IndexType stride_mid_hkv;
+    IndexType stride_mid_seq;
+    IndexType stride_mid_batch;
+    IndexType stride_mid_split;
+  };
+
+  struct PackGQAParams {
+    ProblemShape problem_shape;
+
+    TME_Q tme_load_Q;
+    ShapeQKVO shape_q;
+
+    Element const* ptr_K;
+    RobustDescriptor desc_K;
+    LayoutK layout_K;
+
+    TME_V tme_load_V;
+    ShapeQKVO shape_v;
+
+    ElementOutput* ptr_O;
+
+    // (b, blocks/b)
+    int32_t const* ptr_block_table;
+    // (b,)
+    int32_t const* ptr_seqlen;
+
+    float const sm_scale;
+    float const rescale;
+
+    int const num_splits;
+    int const q_seq_per_hkv;
+    int const h_r;
+
+    IndexType stride_block_table_b;
+
+    // TensorQ (batch_size, seqlen_q * h_r, h_kv, dim_qk)
+    IndexType stride_q_hkv;
+    IndexType stride_q_seq;
+    IndexType stride_q_bs;
+
+    IndexType stride_buf_kblock;
+    IndexType stride_buf_kpage;
+    IndexType stride_buf_kh;
+
+    IndexType stride_buf_vblock;
+    IndexType stride_buf_vpage;
+    IndexType stride_buf_vh;
+
+    IndexType stride_mid_hkv;
+    IndexType stride_mid_seq;
+    IndexType stride_mid_batch;
+    IndexType stride_mid_split;
+  };
+
+  static PackGQAParams
+  to_underlying_arguments(ProblemShape problem_shape, PackGQAArguments const& args) {
+    int SeqQ = size<0>(problem_shape);
+    int DimQK = size<2>(problem_shape);
+    int DimVO = size<3>(problem_shape);
+    int H = size<4>(problem_shape);
+    int H_KV = size<4, 1>(problem_shape);
+    int H_R = size<4, 0>(problem_shape);
+    int B = size<5>(problem_shape);
+
+    auto layout_Q = make_layout(make_shape(SeqQ * H_R, DimQK, H_KV, B),
+                                make_stride(args.stride_q_seq, _1{}, args.stride_q_hkv, args.stride_q_bs));
+
+    Tensor mQ = make_tensor(make_gmem_ptr(args.ptr_Q), layout_Q);
+
+    TME_Q tme_load_Q = make_tme_copy(
+        MP31_TME_LOAD{},
+        mQ,
+        SmemLayoutQ{},
+        select<0, 2>(TileShapeQKD{}));
+
+    RobustDescriptor desc_K = make_robust_desc(args.ptr_K_buffer, DimQK * H_KV * PageSize * args.num_blocks);
+
+    auto layout_K = make_layout(make_shape(PageSize, DimQK, H_KV, args.num_blocks),
+                                make_stride(args.stride_buf_kpage, _1{}, args.stride_buf_kh, args.stride_buf_kblock));
+
+    auto layout_V = make_layout(make_shape(PageSize, DimVO, H_KV, args.num_blocks),
+                                make_stride(args.stride_buf_vpage, _1{}, args.stride_buf_vh, args.stride_buf_vblock));
+    Tensor mV = make_tensor(make_gmem_ptr(args.ptr_V_buffer), layout_V);
+
+    TME_V tme_load_V = make_tme_copy(
+        MP31_TME_LOAD{},
+        mV,
+        take<0, 2>(SmemLayoutV{}),
+        select<2, 1>(TileShapePDV{}));
+
+    constexpr float kLog2e = float(1.4426950408889634074);
+    constexpr float rescale = 1.0f / kLog2e;
+
+    float sm_scale = args.sm_scale * kLog2e;
+
+    return {
+      problem_shape,
+      tme_load_Q,
+      shape(layout_Q),
+      args.ptr_K_buffer,
+      desc_K,
+      layout_K,
+      tme_load_V,
+      shape(layout_V),
+      args.ptr_O,
+      args.ptr_block_table,
+      args.ptr_seqlen,
+      sm_scale,
+      rescale,
+      args.num_splits,
+      SeqQ * H_R,
+      H_R,
+      args.stride_block_table_b,
+      args.stride_q_hkv,
+      args.stride_q_seq,
+      args.stride_q_bs,
+
+      args.stride_buf_kblock,
+      args.stride_buf_kpage,
+      args.stride_buf_kh,
+
+      args.stride_buf_vblock,
+      args.stride_buf_vpage,
+      args.stride_buf_vh,
+
+      args.stride_mid_hkv,
+      args.stride_mid_seq,
+      args.stride_mid_batch,
+      args.stride_mid_split
+    };
+  }
+};
diff --git a/vllm_musa/csrc/decode_mla/mp31_impl/online_softmax.hpp b/vllm_musa/csrc/decode_mla/mp31_impl/online_softmax.hpp
new file mode 100644
index 000000000..73efffdb8
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/mp31_impl/online_softmax.hpp
@@ -0,0 +1,160 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+**************************************************************************************************/
+
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mutlass/fast_math.h>
+#include "../online_softmax_ops.hpp"
+
+using namespace mute;
+
+template <class LayoutAccum>
+MUTE_DEVICE
+auto
+mp31_convert_layout_acc_rowcol(LayoutAccum acc_layout) {
+  // from ((2, VN, VM), MMA_M, MMA_N) to ((VM, MMA_M),(2, VN, MMA_N));
+  static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+  static_assert(decltype(rank(acc_layout))::value == 3);
+  auto l = acc_layout;
+  return make_layout(make_layout(get<0,2>(l), get<1>(l)), make_layout(get<0,0>(l), get<0,1>(l), get<2>(l)));
+}
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class Fragment0, class Fragment1, class AccumType>
+MUTE_HOST_DEVICE
+void
+mp31_online_softmax_scale(Fragment0& accum, Fragment1& accum_softmax, AccumType rln2_scale) {
+  MUTE_UNROLL
+  for (int i = 0; i < size(accum); ++i) {
+    accum_softmax(i) = accum(i) * rln2_scale;
+  }
+}
+
+
+template <class ElementAccum, class Fragment0, class Fragmnet1>
+MUTE_HOST_DEVICE
+void
+mp31_online_softmax_row_max(Fragment0& accum_softmax,
+                       Fragmnet1& row_max,
+                       Fragmnet1& old_row_max,
+                       Fragmnet1& alpha) {
+  Tensor accum_s = make_tensor(accum_softmax.data(), mp31_convert_layout_acc_rowcol(accum_softmax.layout()));
+
+  MUTE_STATIC_ASSERT_V(size(row_max) == size<0>(accum_s));
+  MUTE_STATIC_ASSERT_V(size(old_row_max) == size<0>(accum_s));
+  MUTE_STATIC_ASSERT_V(size(alpha) == size<0>(accum_s));
+  //static_assert(size<1>(accum_s) == 16);
+
+  MaxOp<ElementAccum> op;
+  ThreadReduce<size<1>(accum_s)> thread_reduce;
+  thread_reduce.template run<true>(accum_s, row_max, op);
+  octave_allreduce_(row_max, row_max, op);
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    row_max(mi) = mutlass::fast_max(row_max(mi), old_row_max(mi));
+
+    alpha(mi) = old_row_max(mi) - row_max(mi);
+    alpha(mi) = exp2f(alpha(mi));
+    old_row_max(mi) = row_max(mi);
+  }
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_s); ++ni) {
+      if (row_max(mi) == -mutlass::platform::numeric_limits<ElementAccum>::infinity()) {
+        alpha(mi) = 0.0f;
+        accum_s(mi, ni) = -mutlass::platform::numeric_limits<ElementAccum>::infinity();
+      } else {
+        accum_s(mi, ni) = accum_s(mi, ni) - row_max(mi);
+      }
+    }
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+mp31_update_accum(Fragment0& accum, Fragment1& alpha) {
+  Tensor accum_o = make_tensor(accum.data(), mp31_convert_layout_acc_rowcol(accum.layout()));
+  MUTE_STATIC_ASSERT_V(size(alpha) == size<0>(accum_o));
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_o); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_o); ++ni) {
+      accum_o(mi, ni) *= alpha(mi);
+    }
+  }
+}
+
+
+template <class ElementAccum, class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+mp31_online_softmax_row_sum(Fragment0& accum_softmax, Fragment1& row_sum,
+                       Fragment1& lse, Fragment1& alpha) {
+  Tensor accum_s = make_tensor(accum_softmax.data(), mp31_convert_layout_acc_rowcol(accum_softmax.layout()));
+  //static_assert(size<1>(accum_s) == 16);
+  MUTE_STATIC_ASSERT_V(size<0>(accum_s) == size(row_sum));
+  SumOp<ElementAccum> op;
+  ThreadReduce<size<1>(accum_s)> thread_reduce;
+  thread_reduce.template run<true>(accum_s, row_sum, op);
+  octave_allreduce_(row_sum, row_sum, op);
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    lse(mi) = lse(mi) * alpha(mi) + row_sum(mi);
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+mp31_apply_softmax_normalizer(Fragment0& accum, Fragment1& lse) {
+  Tensor accum_o = make_tensor(accum.data(), mp31_convert_layout_acc_rowcol(accum.layout()));
+  MUTE_STATIC_ASSERT_V(size<0>(accum_o) == size(lse));
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(lse); ++i) {
+    lse(i) = 1.0f / lse(i);
+  }
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_o); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_o); ++ni) {
+      accum_o(mi, ni) *= lse(mi);
+    }
+  }
+}
diff --git a/vllm_musa/csrc/decode_mla/online_softmax_ops.hpp b/vllm_musa/csrc/decode_mla/online_softmax_ops.hpp
new file mode 100644
index 000000000..e14fe0f00
--- /dev/null
+++ b/vllm_musa/csrc/decode_mla/online_softmax_ops.hpp
@@ -0,0 +1,180 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+**************************************************************************************************/
+
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mutlass/fast_math.h>
+
+using namespace mute;
+
+template <typename T>
+struct MaxOp {
+  MUTE_DEVICE T operator()(T const& x, T const& y) { return x > y ? x : y; }
+};
+
+template <>
+struct MaxOp<float> {
+  MUTE_DEVICE float operator()(float const& x, float const& y) { return max(x, y); };
+};
+
+template <typename T>
+struct SumOp {
+  MUTE_DEVICE T operator()(T const& x, T const& y) { return x + y; }
+};
+
+template <int THREADS>
+struct Allreduce {
+  static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
+  template <typename T, typename Operator>
+  static
+  MUTE_DEVICE T run(T x, Operator& op) {
+    constexpr int OFFSET = THREADS / 2;
+    x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
+    return Allreduce<OFFSET>::run(x, op);
+  }
+};
+
+template <>
+struct Allreduce<2> {
+  template <typename T, typename Operator>
+  static
+  MUTE_DEVICE T run(T x, Operator& op) {
+    x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
+    return x;
+  }
+};
+
+template <typename Tensor0, typename Tensor1, typename Operator>
+MUTE_DEVICE void octave_allreduce_(Tensor0& dst, Tensor1& src, Operator& op) {
+  MUTE_STATIC_ASSERT_V(size(dst) == size(src));
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(dst); ++i) {
+    dst(i) = Allreduce<8>::run(src(i), op);
+  }
+}
+
+template <int COLS>
+struct ThreadReduce {
+  template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
+            typename Operator>
+  MUTE_DEVICE void run(Tensor<Engine0, Layout0> const& tensor,
+                       Tensor<Engine1, Layout1>& summary, Operator& op) {
+    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+    static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+    MUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
+
+    MUTE_UNROLL
+    for (int mi = 0; mi < size<0>(tensor); ++mi) {
+      summary(mi) = init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0));
+      MUTE_UNROLL
+      for (int ni = 1; ni < size<1>(tensor); ++ni) {
+        summary(mi) = op(summary(mi), tensor(mi, ni));
+      }
+    }
+  }
+};
+
+
+template <>
+struct ThreadReduce<16> {
+  template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
+            typename Operator>
+  MUTE_DEVICE void run(Tensor<Engine0, Layout0> const& tensor,
+                       Tensor<Engine1, Layout1>& summary, Operator& op) {
+    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+    static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+    MUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
+
+    MUTE_UNROLL
+    for (int mi = 0; mi < size<0>(tensor); ++mi) {
+      // Manually unroll to trigger some nice optimizations
+      Tensor reduce_group = make_tensor<typename Engine1::value_type>(Shape<_4>{});
+
+      reduce_group(0) = op(init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0)), tensor(mi, 0x4));
+      reduce_group(1) = op(tensor(mi, 0x1), tensor(mi, 0x5));
+      reduce_group(2) = op(tensor(mi, 0x2), tensor(mi, 0x6));
+      reduce_group(3) = op(tensor(mi, 0x3), tensor(mi, 0x7));
+
+      reduce_group(0) = op(reduce_group(0), tensor(mi, 0x8));
+      reduce_group(1) = op(reduce_group(1), tensor(mi, 0x9));
+      reduce_group(2) = op(reduce_group(2), tensor(mi, 0xa));
+      reduce_group(3) = op(reduce_group(3), tensor(mi, 0xb));
+
+      reduce_group(0) = op(reduce_group(0), tensor(mi, 0xc));
+      reduce_group(1) = op(reduce_group(1), tensor(mi, 0xd));
+      reduce_group(2) = op(reduce_group(2), tensor(mi, 0xe));
+      reduce_group(3) = op(reduce_group(3), tensor(mi, 0xf));
+
+      reduce_group(0) = op(reduce_group(0), reduce_group(2));
+      reduce_group(1) = op(reduce_group(1), reduce_group(3));
+
+      summary(mi) = op(reduce_group(0), reduce_group(1));
+    }
+  }
+};
+
+template <class ElementAccum, class Fragment0, class Fragment1, class AccumType, class Fragment2>
+MUTE_HOST_DEVICE
+void
+online_softmax_scale_last(Fragment0& accum, Fragment1& accum_softmax, AccumType rln2_scale, Fragment2 tCcS, int oob_n) {
+
+ MUTE_UNROLL
+  for (int i = 0; i < size(accum_softmax); ++i) {
+   if (get<1>(tCcS(i)) < oob_n) {
+     accum_softmax(i) = accum(i) * rln2_scale;
+   } else {
+     accum_softmax(i) = -999999;
+   }
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+add_mask(Fragment0& accum, Fragment1& tCcMask, int m_idx, int n_idx, float margin) {
+  #pragma unroll
+  for (int i = 0; i < size(accum); ++i) {
+    if ((m_idx + get<0>(tCcMask(i)))/(n_idx + get<1>(tCcMask(i))) < margin)
+     accum(i) += -std::numeric_limits<float>::infinity();
+  }
+}
+
+template <class Fragmnet>
+MUTE_HOST_DEVICE
+void
+online_softmax_exp2(Fragmnet& accum) {
+  MUTE_UNROLL
+  for (int i = 0; i < size(accum); ++i) {
+    accum(i) = exp2f(accum(i));
+  }
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/dispatch_utils.h b/vllm_musa/csrc/dispatch_utils.h
new file mode 100644
index 000000000..aa5c8dbba
--- /dev/null
+++ b/vllm_musa/csrc/dispatch_utils.h
@@ -0,0 +1,42 @@
+/*
+ * Adapted from
+ * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h
+ */
+#pragma once
+
+#include <torch/all.h>
+
+#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...)         \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)  \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+
+#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+
+#define VLLM_DISPATCH_CASE_QUANT_TYPES(...)                    \
+  AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
+
+#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
+
+#define VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(...)   \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)    \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)     \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__)
+
+#define VLLM_DISPATCH_FLOATING_AND_BYTE_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME,                               \
+                     VLLM_DISPATCH_CASE_FLOATING_AND_BYTE_TYPES(__VA_ARGS__))
+
+#define VLLM_DISPATCH_CASE_INTEGRAL_TYPES(...)         \
+  AT_DISPATCH_CASE(at::ScalarType::Byte, __VA_ARGS__)  \
+  AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)  \
+  AT_DISPATCH_CASE(at::ScalarType::Short, __VA_ARGS__) \
+  AT_DISPATCH_CASE(at::ScalarType::Int, __VA_ARGS__)   \
+  AT_DISPATCH_CASE(at::ScalarType::Long, __VA_ARGS__)
+
+#define VLLM_DISPATCH_INTEGRAL_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_INTEGRAL_TYPES(__VA_ARGS__))
diff --git a/vllm_musa/csrc/element_wise/ternary.mu b/vllm_musa/csrc/element_wise/ternary.mu
new file mode 100644
index 000000000..f398e8028
--- /dev/null
+++ b/vllm_musa/csrc/element_wise/ternary.mu
@@ -0,0 +1,106 @@
+#include "../cuda_compat.h"
+#include "../dispatch_utils.h"
+#include "../musa.h"
+#include "torch_musa/csrc/aten/musa/MUSADtype.muh"
+#include <torch/all.h>
+
+typedef __half float16_t;
+typedef __mt_bfloat16 bfloat16_t;
+
+namespace vllm_musa {
+__device__ __host__ __forceinline__ constexpr int64_t ceil_div(int64_t a,
+                                                               int64_t b) {
+  return (a + b - 1) / b;
+}
+
+template <typename scalar_t, int64_t vlen, int iobit = 128>
+__global__ void FusedMulAdd(scalar_t *out, scalar_t *self, scalar_t *bias,
+                            const scalar_t scale, const int64_t elem) {
+  constexpr int bits_of_byte = 8;
+  using Vec =
+      at::musa::VecType<scalar_t, vlen * sizeof(scalar_t) * bits_of_byte>;
+
+  int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+  int64_t grid_stride = (int64_t)gridDim.x * blockDim.x;
+  int64_t grid_stride_vec = grid_stride * vlen;
+
+  for (int64_t offset = tid * vlen; offset < elem; offset += grid_stride_vec) {
+    Vec t = Vec::load(self, offset);
+    Vec b = Vec::load(bias, offset);
+    Vec o;
+#pragma unroll
+    for (int k = 0; k < vlen; ++k) {
+      o.val_.elem[k] = b.val_.elem[k] + t.val_.elem[k] * scale;
+    }
+    Vec::store(out, offset, o);
+  }
+}
+
+} // namespace vllm_musa
+
+void _fused_mul_add(torch::Tensor &output, torch::Tensor &self,
+                    torch::Tensor &bias, const float scale) {
+  TORCH_CHECK(self.sizes() == output.sizes(),
+              "self and output shape don't match");
+  TORCH_CHECK(self.sizes() == bias.sizes(), "self and bias shape don't match");
+  TORCH_CHECK(self.scalar_type() == bias.scalar_type(),
+              "self and bias should be same type");
+  TORCH_CHECK(self.scalar_type() == output.scalar_type(),
+              "self and output should be same type");
+  TORCH_CHECK(self.scalar_type() == at::ScalarType::Float ||
+                  self.scalar_type() == at::ScalarType::BFloat16 ||
+                  self.scalar_type() == at::ScalarType::Half,
+              "self's dtype should be in float/half/bfloat16");
+
+  // device guard
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(self));
+
+  // Suppose the uncontiguous elementwise is much slower
+  if C10_UNLIKELY (!self.is_contiguous()) {
+    self = self.contiguous();
+  }
+  if C10_UNLIKELY (!bias.is_contiguous()) {
+    bias = bias.contiguous();
+  }
+
+  // follow mudnn config for arch==22
+  const int64_t max_load_vec =
+      self.scalar_type() == at::ScalarType::Float ? 4 : 8;
+  const int64_t numel = self.numel();
+  size_t thread_per_block = 512;
+  if (vllm_musa::ceil_div(numel, max_load_vec) <= 128) {
+    thread_per_block = 128;
+  } else if (vllm_musa::ceil_div(numel, max_load_vec) <= 256) {
+    thread_per_block = 256;
+  }
+  size_t nr_block = vllm_musa::ceil_div(numel, max_load_vec * thread_per_block);
+
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  switch (self.scalar_type()) {
+  case at::ScalarType::Float:
+    vllm_musa::FusedMulAdd<float, 4><<<nr_block, thread_per_block, 0, stream>>>(
+        static_cast<float *>(output.data_ptr()),
+        static_cast<float *>(self.data_ptr()),
+        static_cast<float *>(bias.data_ptr()), scale, numel);
+    break;
+  case at::ScalarType::Half:
+    vllm_musa::FusedMulAdd<float16_t, 8>
+        <<<nr_block, thread_per_block, 0, stream>>>(
+            static_cast<float16_t *>(output.data_ptr()),
+            static_cast<float16_t *>(self.data_ptr()),
+            static_cast<float16_t *>(bias.data_ptr()),
+            static_cast<float16_t>(scale), numel);
+    break;
+  case at::ScalarType::BFloat16:
+    vllm_musa::FusedMulAdd<bfloat16_t, 8>
+        <<<nr_block, thread_per_block, 0, stream>>>(
+            static_cast<bfloat16_t *>(output.data_ptr()),
+            static_cast<bfloat16_t *>(self.data_ptr()),
+            static_cast<bfloat16_t *>(bias.data_ptr()),
+            static_cast<bfloat16_t>(scale), numel);
+    break;
+  default:
+    break;
+  }
+}
diff --git a/vllm_musa/csrc/flash_attn/flash_attn_fwd.mu b/vllm_musa/csrc/flash_attn/flash_attn_fwd.mu
new file mode 100644
index 000000000..8e087ad14
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/flash_attn_fwd.mu
@@ -0,0 +1,441 @@
+#include <torch/torch.h>
+#include <iostream>
+#include "impl/attention_kernel.hpp"
+#include "impl/kernel_traits.hpp"
+
+#include "../type_convert.muh"
+#include "../dispatch_utils.h"
+
+#define DISPATCH_TYPE(TYPE, NAME, ...)                       \
+    if (TYPE == at::ScalarType::BFloat16)                    \
+    {                                                        \
+        using scalar_t = mutlass::bfloat16_t;                \
+        __VA_ARGS__                                          \
+    }                                                        \
+    else if (TYPE == at::ScalarType::Half)                   \
+    {                                                        \
+        using scalar_t = mutlass::half_t;                    \
+        __VA_ARGS__                                          \
+    }                                                        \
+    else                                                     \
+    {                                                        \
+        TORCH_CHECK(false, "Unsupported data type: ", TYPE); \
+    }
+
+template <int HeadDimQK, int HeadDimVO, typename T, bool ReturnL = false>
+void flash_attn_fwd_kernel_launcher(
+    const T *Q, const T *K, const T *V,
+    T *O, T *P, float *L,
+    int q, int kv, int d_qk, int d_vo,
+    int h, int h_kv, int b, bool causal, mute::Shape<int, int, int> qkv_elem,
+    int q_seq_stride, int q_head_stride, int q_batch_stride,
+    int k_seq_stride, int k_head_stride, int k_batch_stride,
+    int v_seq_stride, int v_head_stride, int v_batch_stride,
+    int o_seq_stride, int o_head_stride, int o_batch_stride,
+    int l_seq_stride, int l_head_stride, int l_batch_stride)
+{
+    TORCH_CHECK(d_qk == HeadDimQK && d_vo == HeadDimVO,
+                "Head dimensions must match kernel specialization");
+
+    constexpr int TileM = 64;
+    constexpr int TileN = 64;
+    const musaStream_t stream = at::musa::getCurrentMUSAStream();
+    using KernelTraits = AttentionForwardKernelTraits<HeadDimQK, HeadDimVO, TileM, TileN,
+                                                      false, // EnableVarlen
+                                                      true,
+                                                      true, // robustQ
+                                                      true, // robustK
+                                                      true, // robustV
+                                                      T,
+                                                      ReturnL>;
+
+    int h_r = h / h_kv;
+    auto problem_shape = mute::make_shape(q, kv, d_qk, d_vo, mute::make_shape(h_r, h_kv), b);
+
+    using AttentionParams = AttentionParams<KernelTraits>;
+    using Arguments = typename AttentionParams::Arguments;
+    Arguments arguments = {Q, K, V, O, P, L};
+
+    auto params = AttentionParams::to_underlying_arguments(problem_shape, arguments, qkv_elem,
+                                                           q_seq_stride, q_head_stride, q_batch_stride,
+                                                           k_seq_stride, k_head_stride, k_batch_stride,
+                                                           v_seq_stride, v_head_stride, v_batch_stride,
+                                                           o_seq_stride, o_head_stride, o_batch_stride,
+                                                           l_head_stride, l_batch_stride);
+
+    dim3 block_dims(KernelTraits::NumThreads);
+    dim3 grid_dims(mute::ceil_div(mute::size<0>(problem_shape), KernelTraits::CTA_Q),
+                   mute::size<4>(problem_shape),
+                   mute::size<5>(problem_shape));
+
+    if (causal)
+    {
+        flash_atten_forward<KernelTraits, true, true><<<grid_dims, block_dims, 0, stream>>>(params);
+    }
+    else
+    {
+        flash_atten_forward<KernelTraits, true, false><<<grid_dims, block_dims, 0, stream>>>(params);
+    }
+}
+
+template <int HeadDimQK, int HeadDimVO, typename T, bool ReturnL = false>
+void varlen_flash_attn_fwd_kernel_launcher(
+    const T *Q, const T *K, const T *V,
+    T *O, T *P, float *L,
+    int q, int kv, int d_qk, int d_vo,
+    int h, int h_kv, int b, bool causal, mute::Shape<int, int, int> qkv_elem,
+    int q_seq_stride, int q_head_stride, int q_batch_stride,
+    int k_seq_stride, int k_head_stride, int k_batch_stride,
+    int v_seq_stride, int v_head_stride, int v_batch_stride,
+    int o_seq_stride, int o_head_stride, int o_batch_stride,
+    int l_seq_stride, int l_head_stride, int l_batch_stride,
+    int *cu_q_seqlen, int *cu_kv_seqlen, int max_seqlen_q, int max_seqlen_kv)
+{
+    TORCH_CHECK(d_qk == HeadDimQK && d_vo == HeadDimVO,
+                "Head dimensions must match kernel specialization");
+
+    constexpr int TileM = 64;
+    constexpr int TileN = 64;
+    const musaStream_t stream = at::musa::getCurrentMUSAStream();
+    VariableLength varlen_q{max_seqlen_q, cu_q_seqlen};
+    VariableLength varlen_kv{max_seqlen_kv, cu_kv_seqlen};
+
+    int h_r = h / h_kv;
+    auto problem_shape = mute::make_shape(varlen_q, varlen_kv, d_qk, d_vo, mute::make_shape(h_r, h_kv), b);
+
+    using KernelTraits = AttentionForwardKernelTraits<HeadDimQK, HeadDimVO, TileM, TileN,
+                                                      true, // EnableVarlen
+                                                      true, // IsQInReg
+                                                      true, // robustQ
+                                                      true, // robustK
+                                                      true, // robustV
+                                                      T,
+                                                      ReturnL>;
+    using AttentionParams = AttentionParams<KernelTraits>;
+    using Arguments = typename AttentionParams::Arguments;
+    Arguments arguments = {Q, K, V, O, P, L};
+
+    auto params = AttentionParams::to_underlying_arguments(problem_shape, arguments, qkv_elem,
+                                                           q_seq_stride, q_head_stride, q_batch_stride,
+                                                           k_seq_stride, k_head_stride, k_batch_stride,
+                                                           v_seq_stride, v_head_stride, v_batch_stride,
+                                                           o_seq_stride, o_head_stride, o_batch_stride,
+                                                           l_head_stride, l_batch_stride);
+
+    dim3 block_dims(KernelTraits::NumThreads);
+    dim3 grid_dims(mute::ceil_div(mute::size<0>(problem_shape), KernelTraits::CTA_Q),
+                   mute::size<4>(problem_shape),
+                   mute::size<5>(problem_shape));
+
+    if (causal)
+    {
+        flash_atten_forward<KernelTraits, true, true><<<grid_dims, block_dims, 0, stream>>>(params);
+    }
+    else
+    {
+        flash_atten_forward<KernelTraits, true, false><<<grid_dims, block_dims, 0, stream>>>(params);
+    }
+}
+
+template <typename scalar_t, bool ReturnL>
+void dispatch_head_dims_fixed(
+    int d_qk, int d_vo,
+    const scalar_t *q_ptr, const scalar_t *k_ptr, const scalar_t *v_ptr,
+    scalar_t *o_ptr, scalar_t *P, float *l_ptr,
+    int seqlen_q, int seqlen_kv, int num_head_q, int num_head_kv, int batch, bool causal,
+    mute::Shape<int, int, int> qkv_elem,
+    int q_seq_stride, int q_head_stride, int q_batch_stride,
+    int k_seq_stride, int k_head_stride, int k_batch_stride,
+    int v_seq_stride, int v_head_stride, int v_batch_stride,
+    int o_seq_stride, int o_head_stride, int o_batch_stride,
+    int l_seq_stride, int l_head_stride, int l_batch_stride)
+{
+    if (d_qk == 192 && d_vo == 128) {
+        flash_attn_fwd_kernel_launcher<192, 128, scalar_t, ReturnL>(
+            q_ptr, k_ptr, v_ptr, o_ptr, P, l_ptr,
+            seqlen_q, seqlen_kv, d_qk, d_vo,
+            num_head_q, num_head_kv, batch, causal, qkv_elem,
+            q_seq_stride, q_head_stride, q_batch_stride,
+            k_seq_stride, k_head_stride, k_batch_stride,
+            v_seq_stride, v_head_stride, v_batch_stride,
+            o_seq_stride, o_head_stride, o_batch_stride,
+            l_seq_stride, l_head_stride, l_batch_stride);
+    } else if (d_qk == 128 && d_vo == 128) {
+        flash_attn_fwd_kernel_launcher<128, 128, scalar_t, ReturnL>(
+            q_ptr, k_ptr, v_ptr, o_ptr, P, l_ptr,
+            seqlen_q, seqlen_kv, d_qk, d_vo,
+            num_head_q, num_head_kv, batch, causal, qkv_elem,
+            q_seq_stride, q_head_stride, q_batch_stride,
+            k_seq_stride, k_head_stride, k_batch_stride,
+            v_seq_stride, v_head_stride, v_batch_stride,
+            o_seq_stride, o_head_stride, o_batch_stride,
+            l_seq_stride, l_head_stride, l_batch_stride);
+    } else if (d_qk == 64 && d_vo == 64) {
+        flash_attn_fwd_kernel_launcher<64, 64, scalar_t, ReturnL>(
+            q_ptr, k_ptr, v_ptr, o_ptr, P, l_ptr,
+            seqlen_q, seqlen_kv, d_qk, d_vo,
+            num_head_q, num_head_kv, batch, causal, qkv_elem,
+            q_seq_stride, q_head_stride, q_batch_stride,
+            k_seq_stride, k_head_stride, k_batch_stride,
+            v_seq_stride, v_head_stride, v_batch_stride,
+            o_seq_stride, o_head_stride, o_batch_stride,
+            l_seq_stride, l_head_stride, l_batch_stride);
+    } else {
+        TORCH_CHECK(false, "Unsupported head dimensions: d_qk=", d_qk, ", d_vo=", d_vo);
+    }
+}
+
+template <typename scalar_t, bool ReturnL>
+void dispatch_head_dims_varlen(
+    int d_qk, int d_vo,
+    const scalar_t *q_ptr, const scalar_t *k_ptr, const scalar_t *v_ptr,
+    scalar_t *o_ptr, scalar_t *P, float *l_ptr,
+    int max_seqlen_q, int max_seqlen_kv, int num_head, int num_head_kv, int batch, bool causal,
+    mute::Shape<int, int, int> qkv_elem,
+    int q_seq_stride, int q_head_stride, int q_batch_stride,
+    int k_seq_stride, int k_head_stride, int k_batch_stride,
+    int v_seq_stride, int v_head_stride, int v_batch_stride,
+    int o_seq_stride, int o_head_stride, int o_batch_stride,
+    int l_seq_stride, int l_head_stride, int l_batch_stride,
+    int *cu_q_ptr, int *cu_kv_ptr)
+{
+    if (d_qk == 192 && d_vo == 128) {
+        varlen_flash_attn_fwd_kernel_launcher<192, 128, scalar_t, ReturnL>(
+            q_ptr, k_ptr, v_ptr, o_ptr, P, l_ptr,
+            max_seqlen_q, max_seqlen_kv, d_qk, d_vo,
+            num_head, num_head_kv, batch, causal, qkv_elem,
+            q_seq_stride, q_head_stride, q_batch_stride,
+            k_seq_stride, k_head_stride, k_batch_stride,
+            v_seq_stride, v_head_stride, v_batch_stride,
+            o_seq_stride, o_head_stride, o_batch_stride,
+            l_seq_stride, l_head_stride, l_batch_stride,
+            cu_q_ptr, cu_kv_ptr, max_seqlen_q, max_seqlen_kv);
+    } else if (d_qk == 128 && d_vo == 128) {
+        varlen_flash_attn_fwd_kernel_launcher<128, 128, scalar_t, ReturnL>(
+            q_ptr, k_ptr, v_ptr, o_ptr, P, l_ptr,
+            max_seqlen_q, max_seqlen_kv, d_qk, d_vo,
+            num_head, num_head_kv, batch, causal, qkv_elem,
+            q_seq_stride, q_head_stride, q_batch_stride,
+            k_seq_stride, k_head_stride, k_batch_stride,
+            v_seq_stride, v_head_stride, v_batch_stride,
+            o_seq_stride, o_head_stride, o_batch_stride,
+            l_seq_stride, l_head_stride, l_batch_stride,
+            cu_q_ptr, cu_kv_ptr, max_seqlen_q, max_seqlen_kv);
+    } else if (d_qk == 64 && d_vo == 64) {
+        varlen_flash_attn_fwd_kernel_launcher<64, 64, scalar_t, ReturnL>(
+            q_ptr, k_ptr, v_ptr, o_ptr, P, l_ptr,
+            max_seqlen_q, max_seqlen_kv, d_qk, d_vo,
+            num_head, num_head_kv, batch, causal, qkv_elem,
+            q_seq_stride, q_head_stride, q_batch_stride,
+            k_seq_stride, k_head_stride, k_batch_stride,
+            v_seq_stride, v_head_stride, v_batch_stride,
+            o_seq_stride, o_head_stride, o_batch_stride,
+            l_seq_stride, l_head_stride, l_batch_stride,
+            cu_q_ptr, cu_kv_ptr, max_seqlen_q, max_seqlen_kv);
+    } else {
+        TORCH_CHECK(false, "Unsupported head dimensions: d_qk=", d_qk, ", d_vo=", d_vo);
+    }
+}
+
+void flash_attn_fwd(
+    const torch::Tensor &Q,            // verlen ? [total_seqlen, num_head, qk_head_dim] : [batch, num_head, seqlen, qk_head_dim]
+    const torch::Tensor &K,            // verlen ? [total_seqlen, kv_num_head, qk_head_dim] : [batch, kv_num_head, seqlen, qk_head_dim]
+    const torch::Tensor &V,            // verlen ? [total_seqlen, kv_num_head, vo_head_dim] : [batch, kv_num_head, seqlen, vo_head_dim]
+    torch::Tensor &O,                  // [total_seqlen, num_head, d_vo]
+    c10::optional<torch::Tensor> &lse, // [num_head, total_seqlen]
+    const torch::Tensor &cu_q_seqlen,
+    const torch::Tensor &cu_kv_seqlen,
+    int max_seqlen_q,
+    int max_seqlen_kv,
+    bool causal,
+    bool varlen,
+    bool return_softmax_lse)
+{
+    TORCH_CHECK(Q.dim() == (varlen ? 3 : 4), "Q must be ", varlen ? 3 : 4, "D");
+    TORCH_CHECK(K.dim() == (varlen ? 3 : 4), "K must be ", varlen ? 3 : 4, "D");
+    TORCH_CHECK(V.dim() == (varlen ? 3 : 4), "V must be ", varlen ? 3 : 4, "D");
+
+    int current_arch = at::musa::getMUSAArch();
+    TORCH_CHECK(current_arch == 220, "flash_attn_varlen_fun not support on MUSA arch ", current_arch);
+
+    const auto dtype = Q.scalar_type();
+    TORCH_CHECK(dtype == K.scalar_type() && dtype == V.scalar_type(),
+                "Q, K, V must have same dtype");
+    TORCH_CHECK(dtype == at::ScalarType::BFloat16 || dtype == at::ScalarType::Half,
+                "Only BFloat16 and Half supported");
+
+    TORCH_CHECK(O.is_contiguous(), "O must be contiguous");
+    TORCH_CHECK(O.scalar_type() == dtype, "O must have same dtype as inputs");
+
+    if (return_softmax_lse)
+    {
+        TORCH_CHECK(lse.has_value(), "lse tensor must be provided when return_softmax_lse=true");
+        auto &lse_tensor = lse.value();
+        TORCH_CHECK(lse_tensor.scalar_type() == torch::kFloat32, "lse must be float32");
+        TORCH_CHECK(lse_tensor.is_contiguous(), "lse must be contiguous");
+    }
+
+    void *P = nullptr;
+    int q_elem = Q.numel();
+    int k_elem = K.numel();
+    int v_elem = V.numel();
+    auto qkv_elem = mute::make_shape(q_elem, k_elem, v_elem);
+    const auto stream = at::musa::getCurrentMUSAStream();
+    const int d_qk = Q.size(-1);
+    const int d_vo = V.size(-1);
+
+    if (varlen)
+    {
+        TORCH_CHECK(cu_q_seqlen.defined() && cu_kv_seqlen.defined(),
+                    "cu_seqlens must be defined for varlen mode");
+        TORCH_CHECK(cu_q_seqlen.dim() == 1, "cu_q_seqlen must be 1D");
+        TORCH_CHECK(cu_kv_seqlen.dim() == 1, "cu_kv_seqlen must be 1D");
+
+        const int batch = cu_q_seqlen.size(0) - 1;
+        const int num_head = Q.size(1);
+        const int num_head_kv = K.size(1);
+        const int total_seqlen = Q.size(0);
+
+        TORCH_CHECK(O.sizes() == torch::IntArrayRef({total_seqlen, num_head, d_vo}),
+                    "O has wrong shape");
+
+        if (return_softmax_lse)
+        {
+            auto &lse_tensor = lse.value();
+            TORCH_CHECK(lse_tensor.sizes() == torch::IntArrayRef({num_head, total_seqlen}),
+                        "lse has wrong shape");
+        }
+
+        auto q_ptr = Q.data_ptr();
+        auto k_ptr = K.data_ptr();
+        auto v_ptr = V.data_ptr();
+        auto o_ptr = O.data_ptr();
+        auto l_ptr = return_softmax_lse ? lse.value().data_ptr() : nullptr;
+        auto cu_q_ptr = cu_q_seqlen.data_ptr<int>();
+        auto cu_kv_ptr = cu_kv_seqlen.data_ptr<int>();
+
+        DISPATCH_TYPE(dtype, "varlen_fwd", {
+            if (return_softmax_lse)
+            {
+                dispatch_head_dims_varlen<scalar_t, true>(
+                    d_qk, d_vo,
+                    static_cast<scalar_t *>(q_ptr),
+                    static_cast<scalar_t *>(k_ptr),
+                    static_cast<scalar_t *>(v_ptr),
+                    static_cast<scalar_t *>(o_ptr),
+                    static_cast<scalar_t *>(P),
+                    static_cast<float *>(l_ptr),
+                    max_seqlen_q, max_seqlen_kv, num_head, num_head_kv, batch, causal, qkv_elem,
+                    Q.stride(0), Q.stride(1), 0, // q strides
+                    K.stride(0), K.stride(1), 0, // k strides
+                    V.stride(0), V.stride(1), 0, // v strides
+                    O.stride(0), O.stride(1), 0, // o strides
+                    return_softmax_lse ? lse.value().stride(1) : 0,
+                    return_softmax_lse ? lse.value().stride(0) : 0,
+                    0, // l_batch_stride
+                    cu_q_ptr, cu_kv_ptr);
+            }
+            else
+            {
+                dispatch_head_dims_varlen<scalar_t, false>(
+                    d_qk, d_vo,
+                    static_cast<scalar_t *>(q_ptr),
+                    static_cast<scalar_t *>(k_ptr),
+                    static_cast<scalar_t *>(v_ptr),
+                    static_cast<scalar_t *>(o_ptr),
+                    static_cast<scalar_t *>(P),
+                    static_cast<float *>(l_ptr),
+                    max_seqlen_q, max_seqlen_kv, num_head, num_head_kv, batch, causal, qkv_elem,
+                    Q.stride(0), Q.stride(1), 0,
+                    K.stride(0), K.stride(1), 0,
+                    V.stride(0), V.stride(1), 0,
+                    O.stride(0), O.stride(1), 0,
+                    0, 0, 0, // lse strides (未使用)
+                    cu_q_ptr, cu_kv_ptr);
+            }
+        });
+    }
+    else
+    {
+        const int batch = Q.size(0);
+        const int num_head_q = Q.size(1);
+        const int num_head_kv = K.size(1);
+        const int seqlen_q = Q.size(2);
+        const int seqlen_kv = K.size(2);
+
+        TORCH_CHECK(O.sizes() == torch::IntArrayRef({batch, num_head_q, seqlen_q, d_vo}),
+                    "O has wrong shape");
+
+        if (return_softmax_lse)
+        {
+            auto &lse_tensor = lse.value();
+            TORCH_CHECK(lse_tensor.sizes() == torch::IntArrayRef({batch, num_head_q, seqlen_q}),
+                        "lse has wrong shape");
+        }
+
+        auto q_ptr = Q.data_ptr();
+        auto k_ptr = K.data_ptr();
+        auto v_ptr = V.data_ptr();
+        auto o_ptr = O.data_ptr();
+        auto l_ptr = return_softmax_lse ? lse.value().data_ptr() : nullptr;
+
+        const int q_batch_stride = Q.stride(0);
+        const int q_head_stride = Q.stride(1);
+        const int q_seq_stride = Q.stride(2);
+
+        const int k_batch_stride = K.stride(0);
+        const int k_head_stride = K.stride(1);
+        const int k_seq_stride = K.stride(2);
+
+        const int v_batch_stride = V.stride(0);
+        const int v_head_stride = V.stride(1);
+        const int v_seq_stride = V.stride(2);
+
+        const int o_batch_stride = O.stride(0);
+        const int o_head_stride = O.stride(1);
+        const int o_seq_stride = O.stride(2);
+
+        const int l_batch_stride = return_softmax_lse ? lse.value().stride(0) : 0;
+        const int l_head_stride = return_softmax_lse ? lse.value().stride(1) : 0;
+        const int l_seq_stride = return_softmax_lse ? lse.value().stride(2) : 0;
+
+        DISPATCH_TYPE(dtype, "fixed_fwd", {
+            if (return_softmax_lse)
+            {
+                dispatch_head_dims_fixed<scalar_t, true>(
+                    d_qk, d_vo,
+                    static_cast<scalar_t *>(q_ptr),
+                    static_cast<scalar_t *>(k_ptr),
+                    static_cast<scalar_t *>(v_ptr),
+                    static_cast<scalar_t *>(o_ptr),
+                    static_cast<scalar_t *>(P),
+                    static_cast<float *>(l_ptr),
+                    seqlen_q, seqlen_kv, num_head_q, num_head_kv, batch, causal, qkv_elem,
+                    q_seq_stride, q_head_stride, q_batch_stride,
+                    k_seq_stride, k_head_stride, k_batch_stride,
+                    v_seq_stride, v_head_stride, v_batch_stride,
+                    o_seq_stride, o_head_stride, o_batch_stride,
+                    l_seq_stride, l_head_stride, l_batch_stride);
+            }
+            else
+            {
+                dispatch_head_dims_fixed<scalar_t, false>(
+                    d_qk, d_vo,
+                    static_cast<scalar_t *>(q_ptr),
+                    static_cast<scalar_t *>(k_ptr),
+                    static_cast<scalar_t *>(v_ptr),
+                    static_cast<scalar_t *>(o_ptr),
+                    static_cast<scalar_t *>(P),
+                    static_cast<float *>(l_ptr),
+                    seqlen_q, seqlen_kv, num_head_q, num_head_kv, batch, causal, qkv_elem,
+                    q_seq_stride, q_head_stride, q_batch_stride,
+                    k_seq_stride, k_head_stride, k_batch_stride,
+                    v_seq_stride, v_head_stride, v_batch_stride,
+                    o_seq_stride, o_head_stride, o_batch_stride,
+                    l_seq_stride, l_head_stride, l_batch_stride);
+            }
+        });
+    }
+    // musaDeviceSynchronize();
+    // TORCH_CHECK(musaGetLastError() == MUSA_SUCCESS, "MUSA kernel execution failed");
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/flash_attn/impl/attention_kernel.hpp b/vllm_musa/csrc/flash_attn/impl/attention_kernel.hpp
new file mode 100644
index 000000000..6acc00186
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/attention_kernel.hpp
@@ -0,0 +1,465 @@
+#pragma once
+
+#include <mutlass/numeric_conversion.h>
+
+#include "attention_params.hpp"
+#include "utils.hpp"
+#include "online_softmax.hpp"
+
+using namespace mute;
+
+
+template <
+  class KernelTraits,
+  bool EnableSoftmax,
+  bool EnableCausal
+>
+__global__ __launch_bounds__(KernelTraits::NumThreads, 1)
+void flash_atten_forward(const typename AttentionParams<KernelTraits>::Params params) {
+  #if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 220)
+  static_assert(KernelTraits::IsQInReg, "Must keep Q in registers now.");
+
+  using ElementAccum = typename KernelTraits::ElementAccumulator;
+  using ElementL = typename KernelTraits::ElementL;
+  constexpr int FragmentSize = KernelTraits::FragmentSize;
+
+  __shared__ char __attribute__((aligned(128))) smem[sizeof(typename KernelTraits::SharedStorage)];
+  auto& shared_storage = *reinterpret_cast<typename KernelTraits::SharedStorage*>(smem);
+
+  auto [coord_q, coord_h, coord_b] = static_cast<uint3>(blockIdx);
+
+  int thread_idx = threadIdx.x;
+
+  auto tile_shape_Q = select<0, 2>(typename KernelTraits::TileShape_QKD{});
+  auto tile_shape_K = select<1, 2>(typename KernelTraits::TileShape_QKD{});
+  auto tile_shape_V = select<2, 1>(typename KernelTraits::TileShape_PDV{});
+  auto tile_shape_O = select<0, 1>(typename KernelTraits::TileShape_PDV{});
+  auto tile_shape_P = select<0, 1>(typename KernelTraits::TileShape_QKD{});
+
+  int actual_seqlen_q = 0;
+  if constexpr (KernelTraits::IsVarlen) {
+    actual_seqlen_q = get<0>(params.problem_shape).cumulative_length[coord_b + 1] - get<0>(params.problem_shape).cumulative_length[coord_b];
+  } else {
+    actual_seqlen_q = get<0>(params.problem_shape);
+  }
+
+  if (coord_q * KernelTraits::CTA_Q >= actual_seqlen_q) return;
+
+  int actual_seqlen_kv = 0;
+  if constexpr (KernelTraits::IsVarlen) {
+    actual_seqlen_kv = get<1>(params.problem_shape).cumulative_length[coord_b + 1] - get<1>(params.problem_shape).cumulative_length[coord_b];
+  } else {
+    actual_seqlen_kv = get<1>(params.problem_shape);
+  }
+
+  int q_ptr_offset = coord_b * get<3>(params.stride_Q);
+  int k_ptr_offset = coord_b * get<3>(params.stride_K);
+  int v_ptr_offset = coord_b * get<3>(params.stride_V);
+  int o_ptr_offset = coord_b * get<3>(params.stride_O);
+  int l_ptr_offset = coord_b * get<3>(params.stride_L);
+
+  if constexpr (KernelTraits::IsVarlen) {
+    q_ptr_offset = get<0>(params.problem_shape).cumulative_length[coord_b] * get<0>(params.stride_Q);
+    k_ptr_offset = get<1>(params.problem_shape).cumulative_length[coord_b] * get<0>(params.stride_K);
+    v_ptr_offset = get<1>(params.problem_shape).cumulative_length[coord_b] * get<0>(params.stride_V);
+    o_ptr_offset = get<0>(params.problem_shape).cumulative_length[coord_b] * get<0>(params.stride_O);
+    l_ptr_offset = get<0>(params.problem_shape).cumulative_length[coord_b] * get<0>(params.stride_L);
+  }
+
+
+  // (actual_seqlen_q, D_QK, (H_R, H_KV))
+  auto shape_Q = make_shape(actual_seqlen_q, get<2>(params.problem_shape), get<4>(params.problem_shape));
+  // (actual_seqlen_k, D_QK, (H_R, H_KV))
+  auto shape_K = make_shape(actual_seqlen_kv, get<2>(params.problem_shape), get<4>(params.problem_shape));
+  // (actual_seqlen_k, D_VO, (H_R, H_KV))
+  auto shape_V = make_shape(actual_seqlen_kv, get<3>(params.problem_shape), get<4>(params.problem_shape));
+  // (actual_seqlen_q, D_VO, (H_R, H_KV))
+  auto shape_O = make_shape(actual_seqlen_q, get<3>(params.problem_shape), get<4>(params.problem_shape));
+  auto shape_L = make_shape(actual_seqlen_q, mute::_1{}, get<4>(params.problem_shape));
+
+  Tensor mQ = make_tensor(make_gmem_ptr(params.ptr_Q + q_ptr_offset), shape_Q, select<0, 1, 2>(params.stride_Q));
+  Tensor mK = make_tensor(make_gmem_ptr(params.ptr_K + k_ptr_offset), shape_K, select<0, 1, 2>(params.stride_K));
+  Tensor mV = make_tensor(make_gmem_ptr(params.ptr_V + v_ptr_offset), shape_V, select<0, 1, 2>(params.stride_V));
+  Tensor mO = make_tensor(make_gmem_ptr(params.ptr_O + o_ptr_offset), shape_O, select<0, 1, 2>(params.stride_O));
+  Tensor mL = make_tensor(make_gmem_ptr(params.ptr_L + l_ptr_offset), shape_L, select<0, 1, 2>(params.stride_L));
+
+  Tensor gQ = local_tile(mQ, tile_shape_Q, make_coord(coord_q, _0{}, coord_h));   // (CTA_Q, HEAD_DIM_QK)
+  Tensor gK = local_tile(mK, tile_shape_K, make_coord(_, _0{}, coord_h));         // (CTA_KV, HEAD_DIM_QK, kv)
+  Tensor gV = local_tile(mV, tile_shape_V, make_coord(_, _0{}, coord_h));         // (CTA_KV, HEAD_DIM_VO, kv)
+  Tensor gO = local_tile(mO, tile_shape_O, make_coord(coord_q, _0{}, coord_h));   // (CTA_Q, CTA_KV, kv)
+  Tensor gL = local_tile(mL, tile_shape_O, make_coord(coord_q, _0{}, coord_h));
+
+  Tensor sQ  = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), typename KernelTraits::SmemLayoutQ{});
+  Tensor sK  = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), typename KernelTraits::SmemLayoutK{});
+  Tensor sV  = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), typename KernelTraits::SmemLayoutV{});
+  Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), typename KernelTraits::SmemLayoutVt{});
+  Tensor sP  = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), typename KernelTraits::SmemLayoutP{});
+
+  typename KernelTraits::TiledMmaQK tiled_mma_qk;
+  typename KernelTraits::TiledMmaPV tiled_mma_pv;
+
+  typename KernelTraits::GmemTiledCopyQ gmem_tiled_copy_Q;
+  typename KernelTraits::GmemTiledCopyK gmem_tiled_copy_K;
+  typename KernelTraits::GmemTiledCopyV gmem_tiled_copy_V;
+  auto gmem_tiled_copy_L = make_tiled_copy_C_atom(
+    Copy_Atom<UniversalCopy<uint_bit_t<sizeof_bits_v<ElementL>>>, ElementL>{},
+    tiled_mma_qk
+  );
+
+  ThrCopy gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(thread_idx);
+  ThrCopy gmem_thr_copy_K = gmem_tiled_copy_K.get_thread_slice(thread_idx);
+  ThrCopy gmem_thr_copy_V = gmem_tiled_copy_V.get_thread_slice(thread_idx);
+  ThrCopy gmem_thr_copy_L = gmem_tiled_copy_L.get_thread_slice(thread_idx);
+
+  Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); // (QCPY, QCPY_M, QCPY_N)
+  Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); // (QCPY, QCPY_M, QCPY_N)
+  Tensor tQrQ = make_fragment_like(tQsQ);
+
+  Tensor tKgK = gmem_thr_copy_K.partition_S(gK); // (KCPY, KCPY_M, KCPY_N, kv)
+  Tensor tKsK = gmem_thr_copy_K.partition_D(sK); // (KCPY, KCPY_M, KCPY_N)
+  Tensor tKrK = make_fragment_like(tKsK);
+
+  Tensor tVgV = gmem_thr_copy_V.partition_S(gV); // (VCPY, VCPY_M, VCPY_N, kv)
+  Tensor tVsV = gmem_thr_copy_V.partition_D(sV); // (VCPY, VCPY_M, VCPY_N)
+  Tensor tVrV = make_fragment_like(tVsV);
+
+  Tensor tSgL = gmem_thr_copy_L.partition_D(gL);
+  // if (thread(0)) {
+  //   // print("liliang l_ptr_offset = ");print(l_ptr_offset);print("\n");
+  //   // print("liliang params.ptr_L = ");print(params.ptr_L);print("\n");
+  //   // print("liliang mL = ");print(mL);print("\n");
+  //   // print("liliang gL = ");print(gL);print("\n");
+  //   // print("liliang tSgL = ");print(tSgL);print("\n");
+  // }
+
+  auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx);
+  auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx);
+
+  // QK MMA AB
+  Tensor tPrQ = thr_mma_qk.partition_fragment_A(sQ); // (MMA, MMA_M, MMA_K)
+  Tensor tPrK = thr_mma_qk.partition_fragment_B(sK); // (MMA, MMA_M, MMA_K)
+
+  // PV MMA AB
+  Tensor tOrP  = thr_mma_pv.partition_fragment_A(sP);
+  Tensor tOrVt = thr_mma_pv.partition_fragment_B(sVt);
+
+  // Allocate accum
+  Tensor accum_p = partition_fragment_C(tiled_mma_qk, tile_shape_P);
+  Tensor accum_o = partition_fragment_C(tiled_mma_pv, tile_shape_O);
+
+  // Allocate softmax
+  Tensor accum_softmax = make_fragment_like(accum_p);
+  Tensor accum_cvt     = make_fragment_like<typename KernelTraits::ElementKV>(accum_softmax);
+
+  Tensor tAcc_frg = recast<mutlass::Array<ElementAccum, FragmentSize>>(accum_softmax);
+  Tensor tCvt_frg = recast<mutlass::Array<typename KernelTraits::ElementKV, FragmentSize>>(accum_cvt);
+
+  // STS P
+  auto tiled_copy_r2s = make_tiled_copy_C(typename KernelTraits::R2SCopyAtomP{}, tiled_mma_qk);
+  auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(thread_idx);
+  Tensor tPrP = thr_copy_r2s.retile_S(accum_cvt);
+  Tensor tPsP = thr_copy_r2s.partition_D(sP);
+
+  // LDS QKPV
+  auto thr_copy_Q = make_tiled_copy_A(typename KernelTraits::SmemCopyAtomQ{}, tiled_mma_qk).get_thread_slice(thread_idx);
+  auto thr_copy_K = make_tiled_copy_B(typename KernelTraits::SmemCopyAtomK{}, tiled_mma_qk).get_thread_slice(thread_idx);
+  auto thr_copy_P = make_tiled_copy_A(typename KernelTraits::SmemCopyAtomP{}, tiled_mma_pv).get_thread_slice(thread_idx);
+  auto thr_copy_V = make_tiled_copy_B(typename KernelTraits::SmemCopyAtomV{}, tiled_mma_pv).get_thread_slice(thread_idx);
+
+  // QK GEMM
+  Tensor tPsQ           = thr_copy_Q.partition_S(sQ);
+  Tensor tPrQ_copy_view = thr_copy_Q.retile_D(tPrQ);
+  MUTE_STATIC_ASSERT_V(size<1>(tPsQ) == size<1>(tPrQ_copy_view));
+
+  Tensor tPsK           = thr_copy_K.partition_S(sK);
+  Tensor tPrK_copy_view = thr_copy_K.retile_D(tPrK);
+  MUTE_STATIC_ASSERT_V(size<1>(tPsK) == size<1>(tPrK_copy_view));
+
+  // Softmax
+  constexpr int RowsPerThread = size<0>(tile_shape_Q) / tile_size<0>(tiled_mma_qk) * 4;
+  Tensor alpha       = make_tensor<ElementAccum>(Shape<Int<RowsPerThread>>{});
+  Tensor lse         = make_fragment_like(alpha);
+  Tensor row_sum     = make_fragment_like(alpha);
+  Tensor row_max     = make_fragment_like(alpha);
+  Tensor old_row_max = make_fragment_like(alpha);
+
+  if constexpr (EnableSoftmax) {
+    fill(lse,         ElementAccum(0.0f));
+    fill(alpha,       ElementAccum(0.0f));
+    fill(old_row_max, -mutlass::platform::numeric_limits<ElementAccum>::infinity());
+  } else {
+    fill(lse, ElementAccum(1.0f));
+  }
+
+  // PV GEMM
+  Tensor tOsP           = thr_copy_P.partition_S(sP);
+  Tensor tOrP_copy_view = thr_copy_P.retile_D(tOrP);
+
+  Tensor tOsVt           = thr_copy_V.partition_S(sVt);
+  Tensor tOrV_copy_view  = thr_copy_V.retile_D(tOrVt);
+
+  clear(accum_p);
+  clear(accum_o);
+
+  //
+  // PREDICATES
+  //
+
+  // Allocate predicate tensors for Q/K/V
+  auto residue_q = size<0>(shape_Q) - size<0>(gQ) * coord_q;
+  auto residue_kv = (size<0>(shape_K) % size<0>(gK) == 0) ? size<0>(gK) : size<0>(shape_K) % size<0>(gK);
+
+  // Deal with kv oob
+  Tensor cP = make_identity_tensor(tile_shape_P);
+  Tensor tPcP = thr_mma_qk.partition_C(cP);
+  Tensor tPcP_view = make_tensor(tPcP.data(), convert_layout_acc_rowcol(tPcP.layout()));
+
+  // Output Tensor
+  Tensor tOgO = thr_mma_pv.partition_C(gO);
+  Tensor cO   = make_identity_tensor(tile_shape_O);
+  Tensor tCcO = thr_mma_pv.partition_C(cO);
+
+  auto temp_shape = mute::Shape<Int<size<0>(tile_shape_P)>, _1>{};
+  auto temp_layout = make_layout(make_shape(size<0>(tile_shape_P), 1), Stride<_1, _0>{});
+  Tensor cL   = make_identity_tensor(temp_shape);
+  Tensor tCcL = thr_mma_qk.partition_C(cL.compose(temp_layout));
+  // Assume head dim match
+  auto oob_mn = mute::make_tuple(residue_q, size<1>(tile_shape_O));
+
+  // Debug save
+  //Tensor tPgP = thr_copy_r2s.partition_D(gP); // (CPY, CPY_M, CPY_N, kv)
+
+  if (thread0()) {
+    //print("row_sum:"); print(row_sum); print("\n");
+    //print("mQ:"); print(mQ); print("\n");
+    //print("gQ:"); print(gQ); print("\n");
+    //print("mK:"); print(mK); print("\n");
+    //print("gK:"); print(gK); print("\n");
+    //print("mV:"); print(mV); print("\n");
+    //print("gV:"); print(gV); print("\n");
+    //print("gP:"); print(gP); print("\n");
+    //print("tPgP:"); print(tPgP); print("\n");
+    //
+    //print("tPcS"); print(tPcS); print("\n");
+    //
+    //print("shapeK:"); print(shape_K); print("\n");
+    //print("residue_kv:"); print(residue_kv); print("\n");
+    //
+    //print("actual seqlen q:"); print(actual_seqlen_q); print("\n");
+    //print("actual seqlen kv:"); print(actual_seqlen_kv); print("\n");
+
+    //print("tPcP_view"); print(tPcP_view); print("\n");
+    //print("v:"); print(tPcP_view(0, 1)); print("\n");
+    //print("v:"); print(tPcP_view(0, 2)); print("\n");
+    //print("sQ:"); print(sQ); print("\n");
+    //print("sK:"); print(sK); print("\n");
+    //print("sV:"); print(sV); print("\n");
+    //print("sP:"); print(sP); print("\n");
+    //print("tQgQ:"); print(tQgQ); print("\n");
+    //print("tQsQ:"); print(tQsQ); print("\n");
+    //print("tKgK:"); print(tKgK); print("\n");
+    //print("tKsK:"); print(tKsK); print("\n");
+    //print("tKrK:"); print(tKrK); print("\n");
+    //print("accum_p:"); print(accum_p); print("\n");
+    //print("accum_o:"); print(accum_o); print("\n");
+    //print("tPrQ:"); print(tPrQ); print("\n");
+    //print("tPrK:"); print(tPrK); print("\n");
+    //print("tPsQ:"); print(tPsQ); print("\n");
+    //print("tPrQ_copy_view:"); print(tPrQ_copy_view); print("\n");
+  }
+
+  //if (threadIdx.x == 0) {
+  //  print("residue_kv:"); print(residue_kv); print("\n");
+  //}
+
+  //
+  // Prologue
+  //
+
+  // LDG Q
+  if constexpr (KernelTraits::UseRobustLoadQ) {
+    copy(gmem_tiled_copy_Q.with(params.desc_Q), tQgQ, tQrQ);
+  } else {
+    copy(gmem_tiled_copy_Q, tQgQ, tQrQ);
+  }
+  // STS Q
+  copy(tQrQ, tQsQ);
+
+  __syncthreads();
+
+  // LDS Q, keep Q in register
+  copy(tPsQ, tPrQ_copy_view);
+
+  auto k_tile_iter = mute::make_coord_iterator(shape<2>(gK));
+  auto v_tile_iter = k_tile_iter;
+  int orig_kv_tile_count = size<2>(gK);
+  int kv_tile_count = 0;
+
+  if constexpr (EnableCausal) {
+    kv_tile_count = std::min(orig_kv_tile_count,
+                        int(mute::ceil_div((coord_q + 1) * size<0>(tile_shape_Q), size<0>(gK))));
+
+  }
+
+  // LDG K
+  if constexpr (KernelTraits::UseRobustLoadK) {
+    copy(gmem_tiled_copy_K.with(params.desc_K), tKgK(_,_,_,*k_tile_iter), tKrK);
+  } else {
+    copy(gmem_tiled_copy_K, tKgK(_,_,_,*k_tile_iter), tKrK);
+  }
+  ++k_tile_iter;
+
+  if constexpr (KernelTraits::ReuseSmemQ) {
+    __syncthreads();
+  }
+
+  copy(tKrK, tKsK);
+  __syncthreads();
+
+  // P = QK
+  gemm<KernelTraits::IsQInReg>(accum_p, tPrQ, tPrK, tPsQ, tPsK, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+
+
+  // Apply Causal Mask
+  if constexpr (EnableCausal) {
+    apply_causal_mask(accum_p, tPcP, coord_q * size<0>(gQ), 0, false);
+  }
+
+  // Softmax(P)
+  if constexpr (EnableSoftmax) {
+    online_softmax_scale(accum_p, accum_softmax, params.rln2_scale, tPcP, residue_kv, orig_kv_tile_count == 1);
+    //online_softmax_scale(accum_p, accum_softmax, params.rln2_scale);
+    online_softmax_row_max<ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+  } else {
+    online_softmax_scale(accum_p, accum_softmax, float(1.0f));
+  }
+
+  for (int loop = 1; loop < kv_tile_count; ++loop) {
+    if constexpr (KernelTraits::UseRobustLoadK) {
+      copy(gmem_tiled_copy_K.with(params.desc_K), tKgK(_,_,_,*k_tile_iter), tKrK);
+    } else {
+      copy(gmem_tiled_copy_K, tKgK(_,_,_,*k_tile_iter), tKrK);
+    }
+    ++k_tile_iter;
+    copy(tKrK, tKsK);
+    __syncthreads();
+
+    clear(accum_p);
+    gemm<KernelTraits::IsQInReg>(accum_p, tPrQ, tPrK, tPsQ, tPsK, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+
+    if constexpr (EnableSoftmax) {
+      // do softmax on accum_softmax
+      online_softmax_exp2(accum_softmax);
+      update_accum(accum_o, alpha);
+    }
+
+    // tCvt_frg -> accum_cvt, tAcc_frg -> accum_softmax
+    MUTE_UNROLL
+    for (int i = 0; i < size(tCvt_frg); ++i) {
+      tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::ElementKV, ElementAccum, FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+    }
+
+    // STS P (tPrP -> accum_cvt)
+    copy(tiled_copy_r2s, tPrP, tPsP);
+
+    // LDG V
+    if constexpr (KernelTraits::UseRobustLoadV) {
+      copy(gmem_tiled_copy_V.with(params.desc_V), tVgV(_,_,_,*v_tile_iter), tVrV);
+    } else {
+      copy(gmem_tiled_copy_V, tVgV(_,_,_,*v_tile_iter), tVrV);
+    }
+    ++v_tile_iter;
+    // STS V
+    copy(tVrV, tVsV);
+
+    __syncthreads();
+
+    // O = PV
+    gemm(accum_o, tOrP, tOrVt, tOsP, tOsVt, tiled_mma_pv, thr_copy_P, thr_copy_V);
+
+    if constexpr (EnableSoftmax) {
+      // do softmax
+      online_softmax_row_sum<ElementAccum>(accum_softmax, row_sum, lse, alpha);
+
+      // Apply Causal Mask
+      if constexpr (EnableCausal) {
+        apply_causal_mask(accum_p, tPcP, coord_q * size<0>(gQ), loop * size<0>(gK), true);
+      }
+      online_softmax_scale(accum_p, accum_softmax, params.rln2_scale, tPcP, residue_kv, loop == orig_kv_tile_count - 1);
+      online_softmax_row_max<ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+    } else {
+      online_softmax_scale(accum_p, accum_softmax, float(1.0f));
+    }
+  }
+
+  // Last tile
+  if constexpr (EnableSoftmax) {
+    // do softmax
+    online_softmax_exp2(accum_softmax);
+    update_accum(accum_o, alpha);
+  }
+
+  // tCvt_frg -> accum_cvt, tAcc_frg -> accum_softmax
+  MUTE_UNROLL
+  for (int i = 0; i < size(tCvt_frg); ++i) {
+    tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::ElementKV, ElementAccum, FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+  }
+
+  // Temp save
+  //copy(tiled_copy_r2s, tPrP, tPgP(_,_,_,0));
+
+  // LDG V
+  if constexpr (KernelTraits::UseRobustLoadV) {
+    copy(gmem_tiled_copy_V.with(params.desc_V), tVgV(_,_,_,*v_tile_iter), tVrV);
+  } else {
+    copy(gmem_tiled_copy_V, tVgV(_,_,_,*v_tile_iter), tVrV);
+  }
+  // STS V
+  copy(tVrV, tVsV);
+
+  // STS P (tPrP -> accum_cvt)
+  copy(tiled_copy_r2s, tPrP, tPsP);
+
+  __syncthreads();
+
+  // O = PV
+  gemm(accum_o, tOrP, tOrVt, tOsP, tOsVt, tiled_mma_pv, thr_copy_P, thr_copy_V);
+
+  if constexpr (EnableSoftmax) {
+    online_softmax_row_sum<ElementAccum>(accum_softmax, row_sum, lse, alpha);
+  }
+
+  // if (thread(0)) {
+  //   // print("liliang get<0>(oob_mn) = ");print(get<0>(oob_mn));print("\n");
+  //   // print("liliang cL = ");print(cL);print("\n");
+  //   // print("liliang tCcL = ");print(tCcL);print("\n");
+  //   // print("liliang lse = ");print(lse);print("\n");
+  //   // print("liliang row_max = ");print(row_max);print("\n");
+  //   // print("liliang ElementL(log2(lse(i)) + row_max(i)) = ");print(ElementL(log2(lse(0)) + row_max(0)) / float(1.4426950408889634074));print("\n");
+  //   // print("liliang lse(i) = ");print(lse(0));print("\n");
+  //   // print("liliang row_max(i) = ");print(row_max(0));print("\n");
+  //   // print("liliang exp2f = ");print(exp2f(1.0));print("\n");
+  //   // print("liliang exp2f = ");print(exp2f(1.0 * float(1.4426950408889634074)));print("\n");
+  // }
+
+  if constexpr(KernelTraits::ReturnL) {
+    #pragma unroll
+    for(int i = 0; i < size(lse); ++i) {
+      if (elem_less(get<0>(tCcL(i * size<0, 0>(tCcL))), get<0>(oob_mn))) {
+        tSgL(i) = (log2(lse(i)) + row_max(i)) / float(1.4426950408889634074);
+      }
+    }
+  }
+
+  if constexpr (EnableSoftmax) {
+    apply_softmax_normalizer(accum_o, lse);
+  }
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(accum_o); ++i) {
+    if (elem_less(tCcO(i), oob_mn)) {
+      tOgO(i) = typename KernelTraits::ElementQ(accum_o(i));
+    }
+  }
+  #endif
+}
diff --git a/vllm_musa/csrc/flash_attn/impl/attention_params.hpp b/vllm_musa/csrc/flash_attn/impl/attention_params.hpp
new file mode 100644
index 000000000..5f14ecc43
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/attention_params.hpp
@@ -0,0 +1,214 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+#include "mutlass_extension.hpp"
+
+
+struct VariableLength {
+  int max_length;
+  int* cumulative_length = nullptr;
+
+  MUTE_HOST_DEVICE operator int() const {
+    return max_length;
+  }
+};
+
+namespace mute {
+template <>
+struct is_integral<VariableLength> : true_type {};
+
+MUTE_HOST_DEVICE
+void print(VariableLength _) {
+  printf("Varlen<%d, %p>", _.max_length, _.cumulative_length);
+}
+
+} // namespace mute
+
+template <class KernelTraits>
+struct AttentionParams {
+  static constexpr bool IsVarlen = KernelTraits::IsVarlen;
+
+  using ElementQ  = typename KernelTraits::ElementQ;
+  using ElementKV = typename KernelTraits::ElementKV;
+  using ElementAccumulator = typename KernelTraits::ElementAccumulator;
+
+  using ElementO = ElementQ;
+  using ElementL = float;
+  using ElementP = ElementKV;
+
+  using IndexType = int32_t;
+
+  // Q, K, D_QK, D_VO, (H_R, H_KV), B
+  using ProblemShapeFixed  = mute::Shape<int, int, int, int, mute::Shape<int, int>, int>;
+  using ProblemShapeVarlen = mute::Shape<VariableLength, VariableLength, int, int, mute::Shape<int, int>, int>;
+  using ProblemShape = mute::conditional_t<IsVarlen,
+                                           ProblemShapeVarlen,
+                                           ProblemShapeFixed>;
+
+  // Q, D_QK, (H_R, H_KV), B
+  using StrideQ = mute::Stride<int, mute::_1, mute::Stride<int, int>, int>;
+
+  // K, D_QK, (H_R, H_KV), B
+  using StrideK = mute::Stride<int, mute::_1, mute::Stride<mute::_0, int>, int>;
+
+  // V, D_VO, (H_R, H_KV), B
+  using StrideV = StrideK;
+
+  // Q, D_VO, (H_R, H_KV), B
+  using StrideO = StrideQ;
+  using StrideL = mute::Stride<mute::_1, mute::_0, mute::Stride<int, int>, int>;
+
+  // Debug
+  // Q, K, (H_R, H_KV), B
+  using StrideP  = StrideQ;
+
+  struct Arguments {
+    ElementQ const* ptr_Q;
+    ElementKV const* ptr_K;
+    ElementKV const* ptr_V;
+
+    ElementO* ptr_O;
+
+    // Debug
+    ElementP* ptr_P;
+    ElementL* ptr_L;
+  };
+
+
+  struct Params {
+    ProblemShape problem_shape;
+
+    ElementQ const* ptr_Q;
+    StrideQ stride_Q;
+    mute::RobustDescriptor desc_Q;
+    ElementKV const* ptr_K;
+    StrideK stride_K;
+    mute::RobustDescriptor desc_K;
+    ElementKV const* ptr_V;
+    StrideV stride_V;
+    mute::RobustDescriptor desc_V;
+
+    ElementO* ptr_O;
+    StrideO stride_O;
+    ElementL* ptr_L;
+    StrideL stride_L;
+
+    float rln2_scale;
+
+    // Debug
+    ElementP* ptr_P;
+    StrideP stride_P;
+  };
+
+  static Params
+  to_underlying_arguments(ProblemShape problem_shape, Arguments const& args, mute::Shape<int, int, int> qkv_elem) {
+    // ProblemShape: Q, K, D_QK, D_VO, (H_R, H_KV), B
+    int SeqQ  = mute::size<0>(problem_shape);
+    int SeqK  = mute::size<1>(problem_shape);
+    int DimQK = mute::size<2>(problem_shape);
+    int DimVO = mute::size<3>(problem_shape);
+    int H     = mute::size<4>(problem_shape);
+    int H_R   = mute::size<4,0>(problem_shape);
+    int H_KV  = mute::size<4,1>(problem_shape);
+    int B     = mute::size<5>(problem_shape);
+
+    // Q = (SeqQ, DimQK, (H_R, H_KV), B)
+    // K = (SeqK, DimQK, (0, H_KV), B)
+    // V = (SeqK, DimVO, (0, H_KV), B)
+    // O = (SeqQ, DimVO, (H_R, H_KV), B)
+    // L = (SeqQ, 1, (H_R, H_KV), B)
+    StrideQ stride_Q = mute::make_stride(DimQK, mute::_1{}, mute::make_stride(H_KV*DimQK*SeqQ, DimQK*SeqQ), IsVarlen ? 0 : DimQK * SeqQ * H);
+    StrideK stride_K = mute::make_stride(DimQK, mute::_1{}, mute::make_stride(mute::_0{}, DimQK * SeqK), IsVarlen ? 0 : DimQK * SeqK * H_KV);
+    StrideV stride_V = mute::make_stride(DimVO, mute::_1{}, mute::make_stride(mute::_0{}, DimVO * SeqK), IsVarlen ? 0 : DimVO * SeqK * H_KV);
+    StrideO stride_O = mute::make_stride(DimVO, mute::_1{}, mute::make_stride(H_KV*DimVO*SeqQ, DimVO*SeqQ), IsVarlen ? 0 : DimVO * SeqQ * H);
+    StrideL stride_L = mute::make_stride(mute::_1{}, mute::_0{}, mute::make_stride(H_KV*SeqQ, SeqQ), IsVarlen ? 0 : SeqQ * H);
+
+    mute::RobustDescriptor desc_Q = mute::make_robust_desc(args.ptr_Q, mute::get<0>(qkv_elem));
+    mute::RobustDescriptor desc_K = mute::make_robust_desc(args.ptr_K, mute::get<1>(qkv_elem));
+    mute::RobustDescriptor desc_V = mute::make_robust_desc(args.ptr_V, mute::get<2>(qkv_elem));
+
+    constexpr float kLog2e = float(1.4426950408889634074);
+    const float softmax_scale = (1.0f / sqrt(float(DimQK)));
+    float rln2_scale = softmax_scale * kLog2e;
+
+    // Debug
+    StrideP stride_P = mute::make_stride(SeqK, mute::_1{}, mute::make_stride(H_KV*SeqQ*SeqK, SeqQ*SeqK), IsVarlen ? 0 : SeqQ * SeqK * H);
+
+    return {
+      problem_shape,
+      args.ptr_Q,
+      stride_Q,
+      desc_Q,
+      args.ptr_K,
+      stride_K,
+      desc_K,
+      args.ptr_V,
+      stride_V,
+      desc_V,
+      args.ptr_O,
+      stride_O,
+      args.ptr_L,
+      stride_L,
+      rln2_scale,
+      args.ptr_P,
+      stride_P
+    };
+  }
+
+
+  static Params
+  to_underlying_arguments(ProblemShape problem_shape, Arguments const& args,
+                          mute::Shape<int, int, int> qkv_elem,
+                          IndexType q_seq_stride, IndexType q_head_stride, IndexType q_batch_stride,
+                          IndexType k_seq_stride, IndexType k_head_stride, IndexType k_batch_stride,
+                          IndexType v_seq_stride, IndexType v_head_stride, IndexType v_batch_stride,
+                          IndexType o_seq_stride, IndexType o_head_stride, IndexType o_batch_stride,
+                          IndexType l_head_stride, IndexType l_batch_stride
+                          ) {
+
+    int DimQK = mute::size<2>(problem_shape);
+    int H_KV  = mute::size<4,1>(problem_shape);
+
+
+    StrideQ stride_Q = mute::make_stride(q_seq_stride, mute::_1{},
+                      mute::make_stride(H_KV*q_head_stride, q_head_stride), IsVarlen ? 0 : q_batch_stride);
+    StrideK stride_K = mute::make_stride(k_seq_stride, mute::_1{},
+                      mute::make_stride(mute::_0{}, k_head_stride), IsVarlen ? 0 : k_batch_stride);
+    StrideV stride_V = mute::make_stride(v_seq_stride, mute::_1{},
+                      mute::make_stride(mute::_0{}, v_head_stride), IsVarlen ? 0 : v_batch_stride);
+    StrideO stride_O = mute::make_stride(o_seq_stride, mute::_1{},
+                      mute::make_stride(H_KV*o_head_stride, o_head_stride), IsVarlen ? 0 : o_batch_stride);
+    StrideL stride_L = mute::make_stride(mute::_1{}, mute::_0{},
+                      mute::make_stride(H_KV*l_head_stride, l_head_stride), IsVarlen ? 0 : l_batch_stride);
+
+    mute::RobustDescriptor desc_Q = mute::make_robust_desc(args.ptr_Q, mute::get<0>(qkv_elem));
+    mute::RobustDescriptor desc_K = mute::make_robust_desc(args.ptr_K, mute::get<1>(qkv_elem));
+    mute::RobustDescriptor desc_V = mute::make_robust_desc(args.ptr_V, mute::get<2>(qkv_elem));
+
+    constexpr float kLog2e = float(1.4426950408889634074);
+    const float softmax_scale = (1.0f / sqrt(float(DimQK)));
+    float rln2_scale = softmax_scale * kLog2e;
+
+    StrideP stride_P;
+
+    return {
+      problem_shape,
+      args.ptr_Q,
+      stride_Q,
+      desc_Q,
+      args.ptr_K,
+      stride_K,
+      desc_K,
+      args.ptr_V,
+      stride_V,
+      desc_V,
+      args.ptr_O,
+      stride_O,
+      args.ptr_L,
+      stride_L,
+      rln2_scale,
+      args.ptr_P,
+      stride_P
+    };
+  }
+};
diff --git a/vllm_musa/csrc/flash_attn/impl/flash_attention.mu b/vllm_musa/csrc/flash_attn/impl/flash_attention.mu
new file mode 100644
index 000000000..4510f496a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/flash_attention.mu
@@ -0,0 +1,510 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 <mutlass/mutlass.h>
+
+#include <mutlass/util/command_line.h>
+#include <mutlass/util/GPU_Clock.hpp>
+#include <mutlass/util/device_memory.h>
+#include <mutlass/util/distribution.h>
+#include <mutlass/util/reference/device/tensor_fill.h>
+#include <mutlass/util/reference/host/tensor_fill.h>
+#include <mutlass/util/GPU_Clock.hpp>
+
+#include "helper.h"
+
+#include <vector>
+
+#include "kernel_traits.hpp"
+#include "attention_params.hpp"
+#include "attention_kernel.hpp"
+#include "reference_attention.hpp"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+struct Options {
+  bool help = false;
+  bool verbose = false;
+  bool skip_check = false;;
+
+  bool varlen = false;
+
+  int q = 128;
+  int kv = 128;
+  int d_qk = 192;
+  int d_vo = 128;
+  int h = 1;
+  int h_kv = 1;
+  int b = 1;
+
+  int warmup_iter, benchmark_iter;
+  mutlass::Distribution::Kind dist_kind;
+
+  Options():
+    help(false),
+    verbose(false),
+    skip_check(false),
+    varlen(false),
+    dist_kind(mutlass::Distribution::Kind::Uniform),
+    warmup_iter(1), benchmark_iter(1),
+    b(1), d_qk(192),d_vo(128),
+    q(128), kv(128), h(1), h_kv(1)
+  { }
+
+  void parse(int argc, char const** args) {
+    mutlass::CommandLine cmd(argc, args);
+
+    if (cmd.check_cmd_line_flag("help")) {
+      help = true;
+      return;
+    }
+    if (cmd.check_cmd_line_flag("verbose")) {
+      verbose = true;
+    }
+    if (cmd.check_cmd_line_flag("skip-check")) {
+      skip_check = true;
+    }
+
+    if (cmd.check_cmd_line_flag("varlen")) {
+      varlen = true;
+    }
+
+    std::string dist_str;
+    cmd.get_cmd_line_argument("dist", dist_str);
+    if (dist_str == "gaussian") {
+      dist_kind = mutlass::Distribution::Gaussian;
+    }
+
+    Options defaults;
+
+    cmd.get_cmd_line_argument("b", b, defaults.b);
+    cmd.get_cmd_line_argument("d_qk", d_qk, defaults.d_qk);
+    cmd.get_cmd_line_argument("d_vo", d_vo, defaults.d_vo);
+    cmd.get_cmd_line_argument("q", q, defaults.q);
+    cmd.get_cmd_line_argument("kv", kv, defaults.kv);
+    cmd.get_cmd_line_argument("h", h, defaults.h);
+    cmd.get_cmd_line_argument("h_kv", h_kv, defaults.h_kv);
+    cmd.get_cmd_line_argument("warmup", warmup_iter, 1);
+    cmd.get_cmd_line_argument("bench", benchmark_iter, 1);
+  }
+
+  std::ostream & print_usage(std::ostream &out) const {
+    out << "FlashAttention Forward:\n\n"
+        << "Options:\n"
+        << "  --help                      If specified, displays this usage statement.\n"
+        << "  --verbose                   If specified, displays detail information.\n"
+        << "  --skip-check                If specified, skip the correctnes check\n"
+        << "  --varlen                    If specified, use varlen input\n"
+        << "  --dist                      Data distribution of input tensors.(--dist=uniform/gaussian, default: uniform)\n"
+        << "  --warmup=<int>              Warmup iterations(default: --warmup=1)\n"
+        << "  --bench=<int>               Benchmark iterations(default: --bench=1)\n"
+        << "  --b=<int>                   B extent\n"
+        << "  --h=<int>                   H extent\n"
+        << "  --h_kv=<int>                H_K/V extent\n"
+        << "  --q=<int>                   Q extent\n"
+        << "  --kv=<int>                  K/V exntent\n"
+        << "  --d_qk=<int>                D_QK extent\n"
+        << "  --d_vo=<int>                D_VO extent\n"
+        << std::endl;
+    return out;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class Element>
+bool initialize_block(
+    mutlass::DeviceAllocation<Element>& block,
+    mutlass::Distribution::Kind dist_kind = mutlass::Distribution::Kind::Uniform,
+    bool all_one = false,
+    uint64_t seed=2024) {
+
+  if (dist_kind == mutlass::Distribution::Uniform) {
+    Element scope_max, scope_min;
+    int bits_input = mutlass::sizeof_bits<Element>::value;
+    if (bits_input == 1) {
+      scope_max = Element(2);
+      scope_min = Element(0);
+    } else if (bits_input <= 8) {
+      scope_max = Element(2);
+      scope_min = Element(-2);
+    } else if (all_one) {
+      scope_max = Element(1);
+      scope_min = Element(1);
+    } else {
+      scope_max = Element(2);
+      scope_min = Element(-2);
+    }
+    mutlass::reference::device::BlockFillRandomUniform(
+      block.get(), block.size(), seed, scope_max, scope_min, 0);
+  } else if (dist_kind == mutlass::Distribution::Gaussian) {
+    mutlass::reference::device::BlockFillRandomGaussian(
+      block.get(), block.size(), seed, Element(0.0f), Element(1.0f));
+  }
+  return true;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class KernelTraits>
+struct FlashAttentionFwdRunner {
+  static constexpr bool IsVarlen = KernelTraits::IsVarlen;
+  using ElementQ = typename KernelTraits::ElementQ;
+  using ElementKV = typename KernelTraits::ElementKV;
+  using ElementAccumulator = typename KernelTraits::ElementAccumulator;
+  static_assert(mute::is_same_v<ElementQ, ElementKV>);
+
+  using AttentionParams = AttentionParams<KernelTraits>;
+  using Arguments = typename AttentionParams::Arguments;
+  using Params = typename AttentionParams::Params;
+  using ProblemShape = typename AttentionParams::ProblemShape;
+
+  // Data members
+  mutlass::DeviceAllocation<ElementQ>  block_Q;
+  mutlass::DeviceAllocation<ElementKV> block_K;
+  mutlass::DeviceAllocation<ElementKV> block_V;
+  mutlass::DeviceAllocation<ElementQ>  block_O;
+  mutlass::DeviceAllocation<ElementKV> block_P;
+
+
+  // Varlen used
+  std::vector<int> cumulative_seqlen_q;
+  std::vector<int> cumulative_seqlen_kv;
+
+  mutlass::DeviceAllocation<int> device_cumulative_seqlen_q;
+  mutlass::DeviceAllocation<int> device_cumulative_seqlen_kv;
+
+  template <class ProblemShapeIn>
+  void initialize(ProblemShapeIn problem_shape, const Options& options) {
+    auto ShapeQ = select<0,2,4,5>(problem_shape);
+    auto ShapeK = select<1,2,4,5>(problem_shape);
+    auto ShapeV = select<1,3,4,5>(problem_shape);
+    auto ShapeO = select<0,3,4,5>(problem_shape);
+    auto ShapeP = select<0,1,4,5>(problem_shape);
+
+
+    block_Q.reset(size(ShapeQ));
+    block_K.reset(size(ShapeK));
+    block_V.reset(size(ShapeV));
+    block_O.reset(size(ShapeO));
+    block_P.reset(size(ShapeP));
+
+    initialize_block(block_Q, options.dist_kind);
+    initialize_block(block_K, options.dist_kind);
+    initialize_block(block_V, options.dist_kind);
+
+    musaMemset(block_O.get(), 0, block_O.bytes());
+  }
+
+  template <bool EnableSoftmax, bool WithCausal>
+  bool verify(const Options& options) {
+    using Element = ElementQ;
+    using ElementAccum = ElementAccumulator;
+    TestAttention<Element, ElementAccum> testBed(options.h, options.h_kv, options.b, options.d_qk, options.d_vo,
+                                                options.q, options.kv, /*alignment*/1, /*use_mask*/false, WithCausal);
+
+    testBed.initialize();
+
+    std::vector<ElementAccum> MiOut(options.q * options.h * options.b);
+    std::vector<ElementAccum> SprimeOut(options.q * options.h * options.b);
+
+    std::vector<Element> gpu_O(block_O.capacity);
+    std::vector<Element> ref_O = gpu_O;
+
+    block_O.copy_to_host(gpu_O.data());
+
+    constexpr bool use_pre_scaling = EnableSoftmax ? true : false;
+    constexpr bool use_pow2 = true;
+
+    #if DEVICE_REF
+    printf("Device Reference...\n");
+    mutlass::DeviceAllocation<Element> block_S_ref(block_S.capacity);
+    mutlass::DeviceAllocation<Element> block_O_ref(block_O.capacity);
+
+    testBed.compute(block_Q.get(), block_K.get(), block_V.get(),
+                    block_S_ref.get(), block_O_ref.get(),
+                    MiOut.data(), SprimeOut.data(),
+                    use_pow2, use_pre_scaling, EnableSoftmax);
+
+    block_O_ref.copy_to_host(ref_O.data());
+    musaDeviceSynchronize();
+
+    #else
+    printf("Host Reference...\n");
+    std::vector<Element> host_Q(block_Q.capacity);
+    std::vector<Element> host_K(block_K.capacity);
+    std::vector<Element> host_V(block_V.capacity);
+    std::vector<Element> host_P(block_P.capacity);
+
+    std::vector<Element> ref_S(block_P.capacity);
+
+    block_Q.copy_to_host(host_Q.data());
+    block_K.copy_to_host(host_K.data());
+    block_V.copy_to_host(host_V.data());
+    block_P.copy_to_host(host_P.data());
+
+    musaDeviceSynchronize();
+
+    testBed.compute(host_Q.data(), host_K.data(), host_V.data(),
+                    ref_S.data(), ref_O.data(),
+                    MiOut.data(), SprimeOut.data(),
+                    use_pow2, use_pre_scaling, EnableSoftmax);
+    #endif
+
+    //bool pass_p = verify_tensor<Element>(host_P, ref_S, false, true);
+    //printf("Pass P:%d\n", pass_p);
+    bool pass_o = verify_tensor<Element>(gpu_O, ref_O, false, true);
+    printf("Pass O:%d\n", pass_o);
+    return pass_o;
+  }
+
+
+  template <bool WithSoftmax, bool WithCausal>
+  bool verify_varlen(Params params) {
+    using Element = ElementQ;
+    using ElementAccum = ElementAccumulator;
+
+    auto problem_shape = params.problem_shape;
+    // use host cumulative length
+    get<0>(problem_shape).cumulative_length = cumulative_seqlen_q.data();
+    get<1>(problem_shape).cumulative_length = cumulative_seqlen_kv.data();
+
+    std::vector<Element> host_Q(block_Q.capacity);
+    std::vector<Element> host_K(block_K.capacity);
+    std::vector<Element> host_V(block_V.capacity);
+    std::vector<ElementAccum> host_P(block_P.capacity);
+    std::vector<Element> host_O(block_O.capacity);
+
+    block_Q.copy_to_host(host_Q.data());
+    block_K.copy_to_host(host_K.data());
+    block_V.copy_to_host(host_V.data());
+
+
+    // construct tensor
+    Tensor mQ = make_tensor(host_Q.data(), select<0, 2, 4>(problem_shape), select<0, 1, 2>(params.stride_Q));
+    Tensor mK = make_tensor(host_K.data(), select<1, 2, 4>(problem_shape), select<0, 1, 2>(params.stride_K));
+    Tensor mV = make_tensor(host_V.data(), select<1, 3, 4>(problem_shape), select<0, 1, 2>(params.stride_V));
+    Tensor mO = make_tensor(host_O.data(), select<0, 3, 4>(problem_shape), select<0, 1, 2>(params.stride_O));
+
+    TestVarlenAttention<Element, ElementAccum> testbed;
+    testbed.template compute<WithCausal, WithSoftmax, true>(problem_shape, mQ, mK, mV, mO);
+
+    std::vector<Element> gpu_O(block_O.capacity);
+    block_O.copy_to_host(gpu_O.data());
+
+    bool pass_o = verify_tensor<Element>(gpu_O, host_O, false, true);
+    printf("Pass O:%d\n", pass_o);
+    return pass_o;
+  }
+
+  template <class ProblemShapeIn>
+  auto
+  initialize_varlen(ProblemShapeIn problem_shape, const bool same_seqlen_qkv = true, const bool use_fixed_q=true) {
+    int B = get<5>(problem_shape);
+
+    cumulative_seqlen_q = {0};
+    cumulative_seqlen_kv = {0};
+
+    int total_seqlen_q = 0;
+    int total_seqlen_kv = 0;
+    int max_seqlen_q = 0;
+    int max_seqlen_kv = 0;
+
+    for (int b = 0; b < B; ++b) {
+      int seqlen_q = use_fixed_q ? get<0>(problem_shape) : rand() % 1024 + 1;
+      int seqlen_kv = same_seqlen_qkv ? seqlen_q : rand() % 1024 + 1;
+
+      total_seqlen_q += seqlen_q;
+      total_seqlen_kv += seqlen_kv;
+
+      max_seqlen_q = std::max(max_seqlen_q, seqlen_q);
+      max_seqlen_kv = std::max(max_seqlen_kv, seqlen_kv);
+
+      cumulative_seqlen_q.push_back(cumulative_seqlen_q.back() + seqlen_q);
+      cumulative_seqlen_kv.push_back(cumulative_seqlen_kv.back() + seqlen_kv);
+    }
+
+    printf("Q max seqlen:%d total seqlen q:%d\n", max_seqlen_q, total_seqlen_q);
+    printf("KV max seqlen:%d total seqlen kv:%d\n", max_seqlen_kv, total_seqlen_kv);
+
+    ProblemShapeIn problem_size_for_init = problem_shape;
+
+    // used to alloca tensor
+    get<5>(problem_size_for_init) = 1;
+    get<0>(problem_size_for_init) = total_seqlen_q;
+    get<1>(problem_size_for_init) = total_seqlen_kv;
+
+    ProblemShape problem_size_for_launch;
+
+    device_cumulative_seqlen_q.reset(cumulative_seqlen_q.size());
+    device_cumulative_seqlen_q.copy_from_host(cumulative_seqlen_q.data());
+
+    device_cumulative_seqlen_kv.reset(cumulative_seqlen_kv.size());
+    device_cumulative_seqlen_kv.copy_from_host(cumulative_seqlen_kv.data());
+
+    get<0>(problem_size_for_launch) = VariableLength{max_seqlen_q, device_cumulative_seqlen_q.get()};  // Q
+    get<1>(problem_size_for_launch) = VariableLength{max_seqlen_kv, device_cumulative_seqlen_kv.get()}; // KV
+    get<2>(problem_size_for_launch) = get<2>(problem_shape); // D_QK
+    get<3>(problem_size_for_launch) = get<3>(problem_shape); // D_VO
+    get<4>(problem_size_for_launch) = get<4>(problem_shape); // (H_R, H_KV)
+    get<5>(problem_size_for_launch) = get<5>(problem_shape); // B
+
+    return mute::make_tuple(problem_size_for_init, problem_size_for_launch);
+  }
+
+  template <bool EnableSoftmax, bool WithCausal>
+  bool run(const Options& options) {
+    int h_r = options.h / options.h_kv;
+    assert(options.h % options.h_kv == 0);
+
+    auto problem_shape_in = mute::make_shape(options.q, options.kv, options.d_qk, options.d_vo, mute::make_shape(h_r, options.h_kv), options.b);
+
+    ProblemShape problem_shape; // used to run
+    decltype(problem_shape_in) problem_size; // used to alloc
+
+    if constexpr (IsVarlen) {
+      auto [problem_shape_init, problem_shape_launch] = initialize_varlen(problem_shape_in, true, true);
+      problem_size = problem_shape_init;
+      problem_shape = problem_shape_launch;
+
+    } else {
+      problem_shape = problem_shape_in;
+      problem_size = problem_shape_in;
+    }
+
+    // use total seqlen to alloca buffer
+    initialize(problem_size, options);
+
+    Arguments arguments {block_Q.get(), block_K.get(), block_V.get(), block_O.get(), block_P.get()};
+
+    auto qkv_elem = make_shape(static_cast<int>(block_Q.capacity), static_cast<int>(block_K.capacity), static_cast<int>(block_V.capacity));
+
+    // use VariableLength to construct params
+    Params params = AttentionParams::to_underlying_arguments(problem_shape, arguments, qkv_elem);
+
+    print("ProblemShape(Q,K,D_QK,D_VO,(H_R,H_KV),B)"); print(problem_shape); print("\n");
+
+    dim3 block_dims(KernelTraits::NumThreads);
+    dim3 grid_dims(ceil_div(size<0>(problem_shape), KernelTraits::CTA_Q), size<4>(problem_shape), size<5>(problem_shape));
+
+    printf("grid:%d %d %d\n", grid_dims.x, grid_dims.y, grid_dims.z);
+
+
+    musaError_t error = musaDeviceSynchronize();
+    if (error != musaSuccess) {
+      std::cerr << "Somthing failed. Error is " << musaGetErrorString(error) << std::endl;
+      return false;
+    }
+
+    // check correctness
+    flash_atten_forward<KernelTraits, EnableSoftmax, WithCausal><<<grid_dims, block_dims>>>(params);
+
+    error = musaGetLastError();
+    if (error != musaSuccess) {
+      std::cerr << "Launch FA failed. Error is " << musaGetErrorString(error) << std::endl;
+      return false;
+    }
+
+    error = musaDeviceSynchronize();
+    if (error != musaSuccess) {
+      std::cerr << "Running FA failed. Error is " << musaGetErrorString(error) << std::endl;
+      return false;
+    }
+
+    printf("GPU FA finished\n");
+
+    if (!options.skip_check) {
+      bool passed = false;
+      if constexpr (IsVarlen) {
+        passed = verify_varlen<EnableSoftmax, WithCausal>(params);
+      } else {
+        passed = verify<EnableSoftmax, WithCausal>(options);
+      }
+
+      if (passed) {
+        std::cout << "FlasAttn Fwd Passed!\n";
+      } else {
+        std::cerr << "FlasAttn Fwd Failed!\n";
+        return false;
+      }
+    }
+
+    return true;
+  }
+};
+
+
+int main(int argc, char const** argv) {
+  musaDeviceProp props;
+  MUSA_CHECK(musaGetDeviceProperties(&props, 0));
+
+  if (props.major * 10 + props.minor != 22) {
+    std::cout
+      << "This experimental code requires a GPU of MooreThreads's MP22 Architecture.\n";
+    return 0;
+  }
+
+  Options options;
+  options.parse(argc, argv);
+
+  if (options.help) {
+    options.print_usage(std::cout) << std::endl;
+    return 0;
+  }
+
+  constexpr int TileM = 64;
+  constexpr int TileN = 64;
+  constexpr int HeadDimQK = 192;
+  constexpr int HeadDimVO = 128;
+
+
+  //print("TileShapeQKD:"); print(KernelTraits::TileShape_QKD{}); print("\n");
+  //print("TileShapePDV:"); print(KernelTraits::TileShape_PDV{}); print("\n");
+  //print("TileMmaQK:"); print(KernelTraits::TiledMmaQK{}); print("\n");
+  //print("TileMmaPV:"); print(KernelTraits::TiledMmaPV{}); print("\n");
+  //print("SmemQ:"); print(KernelTraits::SmemLayoutQ{}); print("\n");
+  //print("SmemK:"); print(KernelTraits::SmemLayoutK{}); print("\n");
+  //print("SmemV:"); print(KernelTraits::SmemLayoutV{}); print("\n");
+  //print("SmemP:"); print(KernelTraits::SmemLayoutP{}); print("\n");
+  //print("SmemSize:"); print(sizeof(KernelTraits::SharedStorage{})); print("\n");
+  //print("GmemTiledCopyQ:"); print(KernelTraits::GmemTiledCopyQ{}); print("\n");
+  //print("GmemTiledCopyK:"); print(KernelTraits::GmemTiledCopyK{}); print("\n");
+
+  if (options.varlen) {
+    using KernelTraits = AttentionForwardKernelTraits<HeadDimQK, HeadDimVO, TileM, TileN, /*IsVarlen*/true, /*IsQInReg*/true, /*robustQ*/true, /*robustK*/true, /*rosbutV*/true, mutlass::half_t>;
+    FlashAttentionFwdRunner<KernelTraits> runner;
+    runner.run</*Enable Softmax*/true, /*EnaleCausal*/true>(options);
+  } else {
+    using KernelTraits = AttentionForwardKernelTraits<HeadDimQK, HeadDimVO, TileM, TileN, /*IsVarlen*/false, /*IsQInReg*/true, /*robustQ*/true, /*robustK*/true, /*rosbutV*/true, mutlass::half_t>;
+    FlashAttentionFwdRunner<KernelTraits> runner;
+    runner.run</*Enable Softmax*/true, /*EnaleCausal*/true>(options);
+  }
+  return 0;
+}
diff --git a/vllm_musa/csrc/flash_attn/impl/helper.h b/vllm_musa/csrc/flash_attn/impl/helper.h
new file mode 100644
index 000000000..db91ced5b
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/helper.h
@@ -0,0 +1,110 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "musa_runtime.h"
+#include <iostream>
+
+/**
+ * Panic wrapper for unwinding MUTLASS errors
+ */
+#define MUTLASS_CHECK(status)                                                                    \
+  {                                                                                              \
+    mutlass::Status error = status;                                                              \
+    if (error != mutlass::Status::kSuccess) {                                                    \
+      std::cerr << "Got mutlass error: " << mutlassGetStatusString(error) << " at: " << __LINE__ \
+                << std::endl;                                                                    \
+      exit(EXIT_FAILURE);                                                                        \
+    }                                                                                            \
+  }
+
+
+/**
+ * Panic wrapper for unwinding MUSA runtime errors
+ */
+#ifndef MUSA_CHECK
+#define MUSA_CHECK(status)                                              \
+  {                                                                     \
+    musaError_t error = status;                                         \
+    if (error != musaSuccess) {                                         \
+      std::cerr << "Got bad musa status: " << musaGetErrorString(error) \
+                << " at line: " << __LINE__ << std::endl;               \
+      exit(EXIT_FAILURE);                                               \
+    }                                                                   \
+  }
+
+#endif
+/**
+ * GPU timer for recording the elapsed time across kernel(s) launched in GPU stream
+ */
+struct GpuTimer
+{
+    musaStream_t _stream_id;
+    musaEvent_t _start;
+    musaEvent_t _stop;
+
+    /// Constructor
+    GpuTimer() : _stream_id(0)
+    {
+        MUSA_CHECK(musaEventCreate(&_start));
+        MUSA_CHECK(musaEventCreate(&_stop));
+    }
+
+    /// Destructor
+    ~GpuTimer()
+    {
+        MUSA_CHECK(musaEventDestroy(_start));
+        MUSA_CHECK(musaEventDestroy(_stop));
+    }
+
+    /// Start the timer for a given stream (defaults to the default stream)
+    void start(musaStream_t stream_id = 0)
+    {
+        _stream_id = stream_id;
+        MUSA_CHECK(musaEventRecord(_start, _stream_id));
+    }
+
+    /// Stop the timer
+    void stop()
+    {
+        MUSA_CHECK(musaEventRecord(_stop, _stream_id));
+    }
+
+    /// Return the elapsed time (in milliseconds)
+    float elapsed_millis()
+    {
+        float elapsed = 0.0;
+        MUSA_CHECK(musaEventSynchronize(_stop));
+        MUSA_CHECK(musaEventElapsedTime(&elapsed, _start, _stop));
+        return elapsed;
+    }
+};
diff --git a/vllm_musa/csrc/flash_attn/impl/kernel_traits.hpp b/vllm_musa/csrc/flash_attn/impl/kernel_traits.hpp
new file mode 100644
index 000000000..1eb0c211e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/kernel_traits.hpp
@@ -0,0 +1,204 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mute/atom/mma_atom.hpp>
+#include <mute/atom/copy_atom.hpp>
+
+#include <mutlass/mutlass.h>
+#include <mutlass/gemm/gemm.h>
+#include <mutlass/gemm/collective/collective_builder.hpp>
+
+#include "mutlass_extension.hpp"
+
+// using namespace mute;
+
+template <
+  class ElementQ,
+  class ElementKV,
+  class SmemLayoutQ,
+  class SmemLayoutK,
+  class SmemLayoutV,
+  class SmemLayoutP
+>
+struct SharedStorage {
+  mute::array_aligned<ElementQ,  mute::cosize_v<SmemLayoutQ>> smem_q;
+  mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutK>> smem_k;
+  mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutV>> smem_v;
+  mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutP>> smem_p;
+};
+
+
+template <
+  class ElementQ,
+  class ElementKV,
+  class SmemLayoutQ,
+  class SmemLayoutK,
+  class SmemLayoutV,
+  class SmemLayoutP
+>
+struct SharedStorageReuseQ {
+  union {
+    mute::array_aligned<ElementQ,  mute::cosize_v<SmemLayoutQ>> smem_q;
+    struct {
+      mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutK>> smem_k;
+      mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutV>> smem_v;
+      mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutP>> smem_p;
+    };
+  };
+};
+
+
+template <class TileN, class PermuteTileShape = TileN>
+constexpr auto
+get_permute_tile_for_operand_B() {
+  using MmaAtomN = mute::_16;
+  static_assert(TileN{} % PermuteTileShape{} == 0);
+  static_assert(PermuteTileShape{} % MmaAtomN{} == 0);
+
+  using RepeatInsidePermuteTile = decltype(PermuteTileShape{} / MmaAtomN{});
+  using RepeatAcrossPermuteTile = decltype(TileN{} / PermuteTileShape{});
+
+  using PermuteTileN = mute::Layout<mute::Shape <               MmaAtomN, RepeatInsidePermuteTile, RepeatAcrossPermuteTile>,
+                              mute::Stride<RepeatInsidePermuteTile,                      mute::_1,        PermuteTileShape>>;
+
+  return PermuteTileN{};
+}
+
+template <
+  int HEAD_DIM_QK_,
+  int HEAD_DIM_VO_,
+  int CTA_Q_,
+  int CTA_KV_,
+  bool IsVarlen_,
+  bool IsQInReg_,
+  bool UseRobustLoadQ_,
+  bool UseRobustLoadK_,
+  bool UseRobustLoadV_,
+  class ElementQ_,
+  bool ReturnL_,
+  class ElementKV_ = ElementQ_,
+  class IdType_ = int32_t
+>
+struct AttentionForwardKernelTraits {
+  using ElementQ  = ElementQ_;
+  using ElementL  = ElementQ_;
+  using ElementKV = ElementKV_;
+  static constexpr bool ReturnL = ReturnL_;
+  using ElementAccumulator = float;
+
+  static constexpr int CTA_Q       = CTA_Q_;
+  static constexpr int CTA_KV      = CTA_KV_;
+  static constexpr int HEAD_DIM_QK = HEAD_DIM_QK_;
+  static constexpr int HEAD_DIM_VO = HEAD_DIM_VO_;
+  static constexpr bool IsVarlen = IsVarlen_;
+  static constexpr bool IsQInReg = IsQInReg_;
+  static constexpr bool ReuseSmemQ = IsQInReg;
+  static constexpr bool UseRobustLoadQ = UseRobustLoadQ_;
+  static constexpr bool UseRobustLoadK = UseRobustLoadK_;
+  static constexpr bool UseRobustLoadV = UseRobustLoadV_;
+
+  // TODO: now only pack T0V0/TV01, try to pack more regs
+  static constexpr int FragmentSize = 2;
+
+  using TileShape_QKD = mute::Shape<mute::Int<CTA_Q>, mute::Int<CTA_KV>, mute::Int<HEAD_DIM_QK>>;
+  using TileShape_PDV = mute::Shape<mute::Int<CTA_Q>, mute::Int<HEAD_DIM_VO>, mute::Int<CTA_KV>>;
+
+  // We extend in M-dim to avoid smem reduce
+  using AtomLayoutMNK = mute::Layout<mute::Shape<mute::Int<CTA_Q/32>, mute::_1, mute::_1>>;
+
+  // Internal Mma Stride type
+  using MmaStrideMNMajor = mute::Stride<mute::_1, int>;
+  using MmaStrideKMajor  = mute::Stride<int, mute::_1>;
+
+  using PermutedTileN = decltype(get_permute_tile_for_operand_B<mute::Int<CTA_KV>, mute::Int<32>>());
+  using PermuteLayoutMNK = mute::Tile<mute::Underscore, PermutedTileN, mute::Underscore>;
+
+  using TiledMmaQK = decltype(mute::make_tiled_mma(
+    mutlass::gemm::collective::detail::mp22_mma_operation_select<
+                                        ElementQ, MmaStrideKMajor,
+                                        ElementKV,MmaStrideKMajor>(),
+    AtomLayoutMNK{},
+    PermuteLayoutMNK{}));
+
+  using TiledMmaPV =decltype(mute::make_tiled_mma(
+    mutlass::gemm::collective::detail::mp22_mma_operation_select<
+                                        ElementKV, MmaStrideKMajor,
+                                        ElementKV, MmaStrideMNMajor>(),
+    AtomLayoutMNK{}));
+
+  static constexpr int NumThreads = size(TiledMmaQK{});
+  static_assert(size(TiledMmaQK{}) == size(TiledMmaPV{}));
+
+  using SmemLayoutAtomQ = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementQ,  MmaStrideKMajor>());
+  using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, mute::select<0,2>(TileShape_QKD{})));
+
+  using SmemLayoutAtomK = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutK = decltype(tile_to_shape(SmemLayoutAtomK{}, mute::select<1,2>(TileShape_QKD{})));
+
+  using SmemLayoutAtomP = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutP = decltype(tile_to_shape(SmemLayoutAtomP{}, mute::select<0,1>(TileShape_QKD{})));
+
+  using SmemLayoutAtomV = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutV  = decltype(tile_to_shape(SmemLayoutAtomV{}, mute::select<2,1>(TileShape_PDV{})));
+  using SmemLayoutVt = decltype(composition(SmemLayoutV{},
+                        make_layout(mute::select<1,2>(TileShape_PDV{}), mute::LayoutRight{})));
+
+  using SmemCopyAtomQ = mute::Copy_Atom<mute::DefaultCopy, ElementQ>;
+  using SmemCopyAtomK = mute::Copy_Atom<mute::DefaultCopy, ElementKV>;
+  using SmemCopyAtomP = mute::Copy_Atom<mute::DefaultCopy, ElementKV>;
+  using SmemCopyAtomV = mute::Copy_Atom<mute::DefaultCopy, ElementKV>;
+
+  using R2SCopyAtomP  = mute::Copy_Atom<mute::UniversalCopy<mute::uint_byte_t<FragmentSize * sizeof(ElementKV)>>, ElementKV>;
+
+  // TileCopy
+  static constexpr int AlignmentQ = 128 / mute::sizeof_bits_v<ElementQ>;
+  static constexpr int ThreadsMajorQ = mute::gcd(HEAD_DIM_QK / AlignmentQ, NumThreads);
+  static constexpr int ThreadsMinorQ = NumThreads / ThreadsMajorQ;
+  using GmemCopyQThrLayout = mute::Layout<mute::Shape <mute::Int<ThreadsMinorQ>, mute::Int<ThreadsMajorQ>>,
+                                    mute::Stride<mute::Int<ThreadsMajorQ>,                mute::_1>>;
+  using GmemCopyQValLayout = mute::Layout<mute::Shape<mute::_1, mute::Int<AlignmentQ>>>;
+  using GmemCopyOpQ = mute::conditional_t<UseRobustLoadQ,
+            mute::MP22_ROBUST_LOAD<mute::uint_bit_t<AlignmentQ * mute::sizeof_bits_v<ElementQ>>>,
+            mute::UniversalCopy<mute::uint_bit_t<AlignmentQ * mute::sizeof_bits_v<ElementQ>>>>;
+
+  using GmemTiledCopyQ = decltype(make_tiled_copy(
+                            mute::Copy_Atom<GmemCopyOpQ, ElementQ>{},
+                            GmemCopyQThrLayout{}, GmemCopyQValLayout{}));
+
+  static constexpr int AlignmentK = 128 / mute::sizeof_bits_v<ElementKV>;
+  static constexpr int ThreadsMajorK = mute::gcd(HEAD_DIM_QK / AlignmentK, NumThreads);
+  static constexpr int ThreadsMinorK = NumThreads / ThreadsMajorK;
+  using GmemCopyKThrLayout = mute::Layout<mute::Shape <mute::Int<ThreadsMinorK>, mute::Int<ThreadsMajorK>>,
+                                    mute::Stride<mute::Int<ThreadsMajorK>,                mute::_1>>;
+  using GmemCopyKValLayout = mute::Layout<mute::Shape<mute::_1, mute::Int<AlignmentK>>>;
+  using GmemCopyOpK = mute::conditional_t<UseRobustLoadK,
+            mute::MP22_ROBUST_LOAD<mute::uint_bit_t<AlignmentK * mute::sizeof_bits_v<ElementKV>>>,
+            mute::UniversalCopy<mute::uint_bit_t<AlignmentK * mute::sizeof_bits_v<ElementKV>>>>;
+
+
+  using GmemTiledCopyK = decltype(make_tiled_copy(
+                            mute::Copy_Atom<GmemCopyOpK, ElementKV>{},
+                            GmemCopyKThrLayout{}, GmemCopyKValLayout{}));
+
+  static constexpr int AlignmentV = 128 / mute::sizeof_bits_v<ElementKV>;
+  static constexpr int ThreadsMajorV = mute::gcd(HEAD_DIM_VO / AlignmentV, NumThreads);
+  static constexpr int ThreadsMinorV = NumThreads / ThreadsMajorV;
+  using GmemCopyVThrLayout = mute::Layout<mute::Shape <mute::Int<ThreadsMinorV>, mute::Int<ThreadsMajorV>>,
+                                    mute::Stride<mute::Int<ThreadsMajorV>,                mute::_1>>;
+  using GmemCopyVValLayout = mute::Layout<mute::Shape<mute::_1, mute::Int<AlignmentV>>>;
+  using GmemCopyOpV = mute::conditional_t<UseRobustLoadV,
+            mute::MP22_ROBUST_LOAD<mute::uint_bit_t<AlignmentV * mute::sizeof_bits_v<ElementKV>>>,
+            mute::UniversalCopy<mute::uint_bit_t<AlignmentV * mute::sizeof_bits_v<ElementKV>>>>;
+
+  using GmemTiledCopyV = decltype(make_tiled_copy(
+                            mute::Copy_Atom<GmemCopyOpV, ElementKV>{},
+                            GmemCopyVThrLayout{}, GmemCopyVValLayout{}));
+
+  using SharedStorage = mute::conditional_t<ReuseSmemQ,
+      SharedStorageReuseQ<ElementQ, ElementKV, SmemLayoutQ, SmemLayoutK, SmemLayoutV, SmemLayoutP>,
+      SharedStorage<ElementQ, ElementKV, SmemLayoutQ, SmemLayoutK, SmemLayoutV, SmemLayoutP>>;
+};
diff --git a/vllm_musa/csrc/flash_attn/impl/mutlass_extension.hpp b/vllm_musa/csrc/flash_attn/impl/mutlass_extension.hpp
new file mode 100644
index 000000000..cd2a22689
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/mutlass_extension.hpp
@@ -0,0 +1,191 @@
+#pragma once
+
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 220)
+#define MUTE_MP22_ROBUST_EXT
+#endif
+
+#include <mute/config.hpp>
+#include <mute/atom/copy_atom.hpp>
+#include <mute/arch/copy_mp31_desc.hpp> // RobustDescriptor
+
+namespace mute {
+
+//
+// Robust LD
+//
+template <class TS, class TD = TS>
+struct MP22_ROBUST_LOAD
+{
+  using SRegisters = TS[1];
+  using DRegisters = TD[1];
+
+  static_assert(sizeof(TS) == sizeof(TD),
+                "MP22_ROBUST_LOAD requires sizeof(src_value_type) == sizeof(dst_value_type)");
+
+  MUTE_HOST_DEVICE static void
+  copy(TS               const& gmem_src,
+       TD                    & rmem_dst,
+       bool                    pred,
+       RobustDescriptor const& src_desc)
+
+  {
+#if defined(MUTE_MP22_ROBUST_EXT)
+    void const* gmem_ptr = pred ? reinterpret_cast<void const*>(&gmem_src)
+                         : reinterpret_cast<void const*>(src_desc.oob_addr);
+    // meet the intrinsic requirement
+    using RobustRegV3 = uint32_t __attribute__((vector_size(12)));
+    union ExtractRobustReg {
+      RobustReg reg_v4;
+      struct {
+        RobustRegV3 reg;
+        uint32_t dummy;
+      };
+    };
+
+    ExtractRobustReg desc {src_desc.reg};
+
+    if constexpr (sizeof(TS) == 1) {
+      auto val = __musa_ld_robust_i8(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 2) {
+      auto val = __musa_ld_robust_i16(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 4) {
+      auto val = __musa_ld_robust_i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 8) {
+      auto val = __musa_ld_robust_i64(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 16) {
+      auto val = __musa_ld_robust_v4i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 32) {
+      auto val = __musa_ld_robust_v8i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 64) {
+      auto val = __musa_ld_robust_v16i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 128) {
+      auto val = __musa_ld_robust_v32i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else {
+      static_assert(sizeof(TS) <= 128, "Currently we only support up-to 1024-bit load");
+    }
+#else
+    MUTE_INVALID_CONTROL_PATH(
+        "Trying to use MP22_ROBUST_LOAD without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+  }
+};
+
+
+//
+// Robust load traits
+//
+template <class S, class D>
+struct MP22_ROBUST_LOAD_OP : MP22_ROBUST_LOAD<S, D> {};
+
+template <class S, class D>
+struct MP22_ROBUST_PRED_LOAD_OP : MP22_ROBUST_LOAD<S, D> {};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_PRED_LOAD_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  // Predicate value that determines whether to load or return zero
+  bool pred;
+
+  // Overload copy_unpack to pass robust buffer access descriptor
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst)
+  {
+    static_assert(is_gmem<TS>::value, "Expected gmem src for MP22_ROBUST_LOAD");
+    static_assert(is_rmem<TD>::value, "Expected rmem dst for MP22_ROBUST_LOAD");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    MUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+        "This src layout is incompatible with this tiled copy.");
+    MUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+        "This dst layout is incompatible with this tiled copy.");
+
+    MP22_ROBUST_LOAD<S, D>::copy(rS[0], rD[0], traits.pred, traits.src_desc);
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_LOAD_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP22_ROBUST_PRED_LOAD_OP<S, D>>
+  with(bool pred) const {
+    return {src_desc, pred};
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_LOAD<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Construct an executable MP22_ROBUST_LOAD with robust buffer access descriptor
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP22_ROBUST_LOAD_OP<S, D>>
+  with(RobustDescriptor desc) const {
+    return {desc};
+  }
+
+  // Don't try to execute a copy with MP22_ROBUST_LOAD before calling .with()
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst) = delete;
+
+};
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/impl/online_softmax.hpp b/vllm_musa/csrc/flash_attn/impl/online_softmax.hpp
new file mode 100644
index 000000000..d85094285
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/online_softmax.hpp
@@ -0,0 +1,386 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mutlass/fast_math.h>
+
+using namespace mute;
+
+template <class LayoutAccum>
+MUTE_DEVICE
+auto
+convert_layout_acc_rowcol(LayoutAccum acc_layout) {
+  if constexpr (decltype(rank<0>(acc_layout))::value == 3) { // MP31
+    // from ((2, VN, VM), MMA_M, MMA_N) to ((VM, MMA_M),(2, VN, MMA_N));
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(rank(acc_layout))::value == 3);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0,2>(l), get<1>(l)), make_layout(get<0,0>(l), get<0,1>(l), get<2>(l)));
+  } else {  // MP22
+    // from((2,4), MMA_M, MMA_N) to ((4, MMA_M), (2, MMA_N))
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(size<0, 1>(acc_layout))::value == 4);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0,1>(l), get<1>(l)), make_layout(get<0,0>(l), get<2>(l)));
+  }
+}
+
+template <typename T>
+struct MaxOp {
+  MUTE_DEVICE T operator()(T const& x, T const& y) { return x > y ? x : y; }
+};
+
+template <>
+struct MaxOp<float> {
+  MUTE_DEVICE float operator()(float const& x, float const& y) { return max(x, y); };
+};
+
+template <typename T>
+struct SumOp {
+  MUTE_DEVICE T operator()(T const& x, T const& y) { return x + y; }
+};
+
+template <int THREADS>
+struct Allreduce {
+  static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
+  template <typename T, typename Operator>
+  static
+  MUTE_DEVICE T run(T x, Operator& op) {
+    constexpr int OFFSET = THREADS / 2;
+    x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
+    return Allreduce<OFFSET>::run(x, op);
+  }
+};
+
+template <>
+struct Allreduce<2> {
+  template <typename T, typename Operator>
+  static
+  MUTE_DEVICE T run(T x, Operator& op) {
+    x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
+    return x;
+  }
+};
+
+template <typename Tensor0, typename Tensor1, typename Operator>
+MUTE_DEVICE void quad_allreduce_(Tensor0& dst, Tensor1& src, Operator& op) {
+  MUTE_STATIC_ASSERT_V(size(dst) == size(src));
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(dst); ++i) {
+    dst(i) = Allreduce<4>::run(src(i), op);
+  }
+}
+
+template <typename Tensor0, typename Tensor1, typename Operator>
+MUTE_DEVICE void extra_allreduce_(Tensor0& dst, Tensor1& src, Operator& op) {
+  MUTE_STATIC_ASSERT_V(size(dst) == size(src));
+  constexpr int mask = ((~(128 - 1)) & 0x7f) << 7 | (128 - 1);
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(dst); ++i) {
+    int val = *reinterpret_cast<int32_t*>(&src(i));
+    int ret = __musa_shfl_xor_sync_i32(val, 32 & 0x7f, mask);
+    // TODO: assume float here...get type from tensor
+    dst(i) = op(src(i), *reinterpret_cast<float*>(&ret));
+  }
+}
+
+template <typename Tensor0, typename Tensor1, typename Operator>
+MUTE_DEVICE void octave_allreduce_(Tensor0& dst, Tensor1& src, Operator& op) {
+  MUTE_STATIC_ASSERT_V(size(dst) == size(src));
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(dst); ++i) {
+    dst(i) = Allreduce<8>::run(src(i), op);
+  }
+}
+
+template <int COLS>
+struct ThreadReduce {
+  template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
+            typename Operator>
+  MUTE_DEVICE void run(Tensor<Engine0, Layout0> const& tensor,
+                       Tensor<Engine1, Layout1>& summary, Operator& op) {
+    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+    static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+    MUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
+
+    MUTE_UNROLL
+    for (int mi = 0; mi < size<0>(tensor); ++mi) {
+      summary(mi) = init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0));
+      MUTE_UNROLL
+      for (int ni = 1; ni < size<1>(tensor); ++ni) {
+        summary(mi) = op(summary(mi), tensor(mi, ni));
+      }
+    }
+  }
+};
+
+template <>
+struct ThreadReduce<16> {
+  template <bool init, typename Engine0, typename Layout0, typename Engine1, typename Layout1,
+            typename Operator>
+  MUTE_DEVICE void run(Tensor<Engine0, Layout0> const& tensor,
+                       Tensor<Engine1, Layout1>& summary, Operator& op) {
+    static_assert(Layout0::rank == 2, "Only support 2D Tensor");
+    static_assert(Layout1::rank == 1, "Only support 1D Tensor");
+    MUTE_STATIC_ASSERT_V(size<0>(summary) == size<0>(tensor));
+
+    MUTE_UNROLL
+    for (int mi = 0; mi < size<0>(tensor); ++mi) {
+      // Manually unroll to trigger some nice optimizations
+      Tensor reduce_group = make_tensor<typename Engine1::value_type>(Shape<_4>{});
+
+      reduce_group(0) = op(init ? tensor(mi, 0) : op(summary(mi), tensor(mi, 0)), tensor(mi, 0x4));
+      reduce_group(1) = op(tensor(mi, 0x1), tensor(mi, 0x5));
+      reduce_group(2) = op(tensor(mi, 0x2), tensor(mi, 0x6));
+      reduce_group(3) = op(tensor(mi, 0x3), tensor(mi, 0x7));
+
+      reduce_group(0) = op(reduce_group(0), tensor(mi, 0x8));
+      reduce_group(1) = op(reduce_group(1), tensor(mi, 0x9));
+      reduce_group(2) = op(reduce_group(2), tensor(mi, 0xa));
+      reduce_group(3) = op(reduce_group(3), tensor(mi, 0xb));
+
+      reduce_group(0) = op(reduce_group(0), tensor(mi, 0xc));
+      reduce_group(1) = op(reduce_group(1), tensor(mi, 0xd));
+      reduce_group(2) = op(reduce_group(2), tensor(mi, 0xe));
+      reduce_group(3) = op(reduce_group(3), tensor(mi, 0xf));
+
+      reduce_group(0) = op(reduce_group(0), reduce_group(2));
+      reduce_group(1) = op(reduce_group(1), reduce_group(3));
+
+      summary(mi) = op(reduce_group(0), reduce_group(1));
+    }
+  }
+};
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class Fragment0, class Fragment1, class AccumType>
+MUTE_HOST_DEVICE
+void
+online_softmax_scale(Fragment0& accum, Fragment1& accum_softmax, AccumType rln2_scale) {
+  MUTE_UNROLL
+  for (int i = 0; i < size(accum); ++i) {
+    accum_softmax(i) = accum(i) * rln2_scale;
+  }
+}
+
+template <class Fragment0, class Fragment1, class AccumType, class Fragment2>
+MUTE_HOST_DEVICE
+void
+online_softmax_scale(Fragment0& accum, Fragment1& accum_softmax, AccumType rln2_scale, Fragment2& tPcP, int oob_n, bool cond) {
+  Tensor accum_dst = make_tensor(accum_softmax.data(), convert_layout_acc_rowcol(accum_softmax.layout()));
+  Tensor accum_src = make_tensor(accum.data(), convert_layout_acc_rowcol(accum.layout()));
+  Tensor cP_view   = make_tensor(tPcP.data(), convert_layout_acc_rowcol(tPcP.layout()));
+
+  MUTE_STATIC_ASSERT_V(size<0>(accum_dst) == size<0>(accum_src));
+  MUTE_STATIC_ASSERT_V(size<1>(accum_dst) == size<1>(accum_src));
+
+  //if (threadIdx.x == 0) {
+  //  printf("bid:%d cond:%d\n", blockIdx.x, cond);
+  //}
+
+  MUTE_UNROLL
+  for (int m = 0; m < size<0>(accum_dst); ++m) {
+    MUTE_UNROLL
+    for (int n = 0; n < size<1>(accum_dst); ++n) {
+      accum_dst(m, n) = (cond && get<1>(cP_view(0,n)) >= oob_n) ?
+                                -mutlass::platform::numeric_limits<AccumType>::infinity()
+                              : accum_src(m, n) * rln2_scale;
+    }
+  }
+
+}
+
+template <class ElementAccum, class Fragment0, class Fragment1, class AccumType, class Fragment2>
+MUTE_HOST_DEVICE
+void
+online_softmax_scale_last(Fragment0& accum, Fragment1& accum_softmax, AccumType rln2_scale, Fragment2 tCcS, int oob_n) {
+
+ MUTE_UNROLL
+  for (int i = 0; i < size(accum_softmax); ++i) {
+   if (get<1>(tCcS(i)) < oob_n) {
+     accum_softmax(i) = accum(i) * rln2_scale;
+   } else {
+     accum_softmax(i) = -999999;
+   }
+  }
+}
+
+
+template <class ElementAccum, class Fragment0, class Fragmnet1>
+MUTE_HOST_DEVICE
+void
+online_softmax_row_max(Fragment0& accum_softmax,
+                       Fragmnet1& row_max,
+                       Fragmnet1& old_row_max,
+                       Fragmnet1& alpha) {
+  Tensor accum_s = make_tensor(accum_softmax.data(), convert_layout_acc_rowcol(accum_softmax.layout()));
+
+  MUTE_STATIC_ASSERT_V(size(row_max) == size<0>(accum_s));
+  MUTE_STATIC_ASSERT_V(size(old_row_max) == size<0>(accum_s));
+  MUTE_STATIC_ASSERT_V(size(alpha) == size<0>(accum_s));
+  //static_assert(size<1>(accum_s) == 16);
+
+  MaxOp<ElementAccum> op;
+  ThreadReduce<size<1>(accum_s)> thread_reduce;
+
+  thread_reduce.template run<true>(accum_s, row_max, op);
+
+  octave_allreduce_(row_max, row_max, op);
+  // use width=128 shulff xor
+  extra_allreduce_(row_max, row_max, op);
+
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    row_max(mi) = mutlass::fast_max(row_max(mi), old_row_max(mi));
+
+    alpha(mi) = old_row_max(mi) - row_max(mi);
+    alpha(mi) = exp2f(alpha(mi));
+    old_row_max(mi) = row_max(mi);
+  }
+
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_s); ++ni) {
+      accum_s(mi, ni) = accum_s(mi, ni) - row_max(mi);
+    }
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+update_accum(Fragment0& accum, Fragment1& alpha) {
+  Tensor accum_o = make_tensor(accum.data(), convert_layout_acc_rowcol(accum.layout()));
+  MUTE_STATIC_ASSERT_V(size(alpha) == size<0>(accum_o));
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_o); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_o); ++ni) {
+      accum_o(mi, ni) *= alpha(mi);
+    }
+  }
+}
+
+template <class Fragmnet>
+MUTE_HOST_DEVICE
+void
+online_softmax_exp2(Fragmnet& accum) {
+  MUTE_UNROLL
+  for (int i = 0; i < size(accum); ++i) {
+    accum(i) = exp2f(accum(i));
+  }
+}
+
+template <class ElementAccum, class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+online_softmax_row_sum(Fragment0& accum_softmax, Fragment1& row_sum,
+                       Fragment1& lse, Fragment1& alpha) {
+  Tensor accum_s = make_tensor(accum_softmax.data(), convert_layout_acc_rowcol(accum_softmax.layout()));
+  //static_assert(size<1>(accum_s) == 16);
+  MUTE_STATIC_ASSERT_V(size<0>(accum_s) == size(row_sum));
+  SumOp<ElementAccum> op;
+  ThreadReduce<size<1>(accum_s)> thread_reduce;
+  thread_reduce.template run<true>(accum_s, row_sum, op);
+  octave_allreduce_(row_sum, row_sum, op);
+  // use width=128 shulff xor
+  extra_allreduce_(row_sum, row_sum, op);
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_s); ++mi) {
+    lse(mi) = lse(mi) * alpha(mi) + row_sum(mi);
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+apply_softmax_normalizer(Fragment0& accum, Fragment1& lse) {
+  Tensor accum_o = make_tensor(accum.data(), convert_layout_acc_rowcol(accum.layout()));
+  MUTE_STATIC_ASSERT_V(size<0>(accum_o) == size(lse));
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(lse); ++i) {
+    lse(i) = 1.0f / lse(i);
+  }
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum_o); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum_o); ++ni) {
+      accum_o(mi, ni) *= lse(mi);
+    }
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+add_mask(Fragment0& accum, Fragment1& tCcMask, int m_idx, int n_idx, float margin) {
+  #pragma unroll
+  for (int i = 0; i < size(accum); ++i) {
+    if ((m_idx + get<0>(tCcMask(i)))/(n_idx + get<1>(tCcMask(i))) < margin)
+     accum(i) += -std::numeric_limits<float>::infinity();
+  }
+}
+
+template <class Fragment0, class Fragment1>
+MUTE_HOST_DEVICE
+void
+apply_causal_mask(Fragment0& accum_, Fragment1& tPcP_, int m_idx, int n_idx, bool debug) {
+  Tensor accum = make_tensor(accum_.data(), convert_layout_acc_rowcol(accum_.layout()));
+  Tensor tPcP  = make_tensor(tPcP_.data(), convert_layout_acc_rowcol(tPcP_.layout()));
+
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<0>(accum); ++mi) {
+    const int row_idx = m_idx + get<0>(tPcP(mi, 0));
+    const int col_idx_limit = row_idx + 1;
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<1>(accum); ++ni) {
+      const int col_idx = n_idx + get<1>(tPcP(0, ni));
+      //if (debug && blockIdx.x == 1) {
+      //  printf("tid:%d col_idx_limit:%d col_idx:%d\n", threadIdx.x, col_idx_limit, col_idx);
+      //}
+      if (col_idx >= col_idx_limit) {
+        accum(mi, ni) = -std::numeric_limits<float>::infinity();
+      }
+    }
+  }
+}
diff --git a/vllm_musa/csrc/flash_attn/impl/reference_attention.hpp b/vllm_musa/csrc/flash_attn/impl/reference_attention.hpp
new file mode 100644
index 000000000..c25165aec
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/reference_attention.hpp
@@ -0,0 +1,570 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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) 2017 - 2023 COLFAX
+ * Distributed under MIT License
+ **************************************************************************************************/
+
+#pragma once
+
+#include <vector>
+
+#include "mutlass/util/device_memory.h"
+#include "mutlass/util/reference/device/gemm_complex.h"
+#include "mutlass/util/reference/host/gemm_complex.h"
+#include "mutlass/util/tensor_view_io.h"
+
+#include "attention_params.hpp"
+
+#define DEVICE_REF 0
+
+/* https://github.com/ColfaxResearch/cutlass-kernels/blob/master/include/utils/random.hpp */
+template <typename Element>
+bool verify_tensor(std::vector<Element> vector_Input,
+                  std::vector<Element> vector_Input_Ref,
+                  bool printValues = false, bool printDiffs = false,
+                  float errCountExpected = 0, int64_t verify_length = -1) {
+  int64_t size = (vector_Input.size() < vector_Input_Ref.size())
+      ? vector_Input.size()
+      : vector_Input_Ref.size();
+  size = (verify_length == -1) ? size : verify_length;
+
+  float abs_tol = 1e-2f;
+
+  float rel_tol = 1e-2f;
+  int errCount = 0;
+  for (int64_t i = 0; i < size; ++i) {
+    if (printValues)
+      std::cout << vector_Input[i] << " " << vector_Input_Ref[i] << std::endl;
+    float diff = (float)(vector_Input[i] - vector_Input_Ref[i]);
+    float abs_diff = fabs(diff);
+    float abs_ref = fabs((float)vector_Input_Ref[i] + 1e-5f);
+    float relative_diff = abs_diff / abs_ref;
+    if ((isnan(vector_Input_Ref[i]) || isnan(abs_diff) || isinf(abs_diff)) ||
+        (abs_diff > abs_tol && relative_diff > rel_tol)) {
+      if (printDiffs)
+        printf("[%d/%d] diff = %f, rel_diff = %f, {computed=%f, ref=%f}.\n",
+              int(i), int(size), abs_diff, relative_diff,
+              (float)(vector_Input[i]), (float)(vector_Input_Ref[i]));
+      errCount++;
+      return false;
+    }
+  }
+  auto errCountComputed = float(errCount) / float(size) * 100;
+
+  return errCountComputed <= errCountExpected ? true : false;
+}
+
+
+template <typename Element, typename AccumType>
+class TestVarlenAttention {
+public:
+  //
+  // Type definitions
+  //
+  static constexpr float kLog2e = float(1.4426950408889634074); // log_2(e) = M_LOG2E
+  using ElementQ = Element;
+  using ElementK = Element;
+  using ElementS = Element;
+  using ElementP = AccumType;
+  using ElementAccumulator = AccumType;
+  using ElementV = Element;
+  using ElementO = Element;
+
+  using ElementCompute = AccumType;
+
+  using ElementNorm = AccumType;
+  using ElementSum = AccumType;
+  using ElementSoftmaxCompute = AccumType;
+
+  template <
+    bool WithCausal,
+    bool WithSoftmax,
+    bool IsVarlen,
+    class ProblemShapeIn,
+    class TensorQ,
+    class TensorK,
+    class TensorV,
+    class TensorO>
+  void compute(ProblemShapeIn problem_shape,
+              TensorQ mQ,
+              TensorK mK,
+              TensorV mV,
+              TensorO mO) {
+
+    auto seqlen_q = get<0>(problem_shape);
+    auto seqlen_kv = get<1>(problem_shape);
+    int DimQK = get<2>(problem_shape);
+    int DimVO = get<3>(problem_shape);
+
+    int H = size<4>(problem_shape);
+    int H_R = get<4,0>(problem_shape);
+    int H_KV = get<4,1>(problem_shape);
+    int B = size<5>(problem_shape);
+
+    static constexpr float kLog2e = float(1.4426950408889634074); // log_2(e) = M_LOG2E
+
+    float alpha = 1.0f / sqrt(float(DimQK));
+
+    for (int b = 0; b < B; ++b) {
+      for (int h = 0; h < H; ++h) {
+
+        int actual_seqlen_q = 0;
+        if constexpr (IsVarlen) {
+          actual_seqlen_q = seqlen_q.cumulative_length[b + 1] - seqlen_q.cumulative_length[b];
+        } else {
+          actual_seqlen_q = seqlen_q;
+        }
+
+        int actual_seqlen_kv = 0;
+        if constexpr (IsVarlen) {
+          actual_seqlen_kv = seqlen_kv.cumulative_length[b + 1] - seqlen_kv.cumulative_length[b];
+        } else {
+          actual_seqlen_kv = seqlen_kv;
+        }
+
+        int q_idx_offset = seqlen_q.cumulative_length[b];
+        int kv_idx_offset = seqlen_kv.cumulative_length[b];
+
+        std::vector<ElementAccumulator> host_P(actual_seqlen_q * actual_seqlen_kv * H);
+
+        Tensor mP = make_tensor(host_P.data(),
+                make_shape(actual_seqlen_q, actual_seqlen_kv, get<4>(problem_shape)),
+                make_stride(actual_seqlen_kv, _1{}, make_stride(H_KV*actual_seqlen_q*actual_seqlen_kv, actual_seqlen_q * actual_seqlen_kv)));
+
+        for (int q = 0; q < actual_seqlen_q; ++q) {
+          // P = QK
+          for (int k = 0; k < actual_seqlen_kv; ++k) {
+            ElementAccumulator acc = ElementAccumulator(0);
+            for (int d = 0; d < DimQK; ++d) {
+              acc += ElementAccumulator(mQ(q_idx_offset + q, d, h)) * ElementAccumulator(mK(kv_idx_offset + k, d, h));
+            }
+
+            if constexpr (WithCausal) {
+              if (k >= (q + 1)) {
+                acc = -std::numeric_limits<ElementAccumulator>::infinity();
+              }
+            }
+            mP(q, k, h) = acc;
+          }
+
+          if constexpr (WithSoftmax) {
+            for (int k = 0; k < actual_seqlen_kv; ++k) {
+              mP(q, k, h) *= kLog2e * alpha;
+            }
+            ElementAccumulator max_val = mP(q, 0, h);
+            for (int k = 1; k < actual_seqlen_kv; ++k) {
+              max_val = std::max(max_val, mP(q, k, h));
+            }
+
+            ElementAccumulator sum_val = 0;
+            for (int k = 0; k < actual_seqlen_kv; ++k) {
+              sum_val += std::exp2f(mP(q, k, h) - max_val);
+            }
+
+            ElementAccumulator inv_sum = 1.0f / sum_val;
+            for (int k = 0; k < actual_seqlen_kv; ++k) {
+              mP(q, k, h) = std::exp2f(mP(q, k, h) - max_val) * inv_sum;
+            }
+          }
+
+          // O = PV
+          for (int d = 0; d < DimVO; ++d) {
+            ElementAccumulator acc = 0;
+            for (int v = 0; v < actual_seqlen_kv; ++v) {
+              acc += ElementAccumulator(mP(q, v, h)) * ElementAccumulator(mV(kv_idx_offset + v, d, h));
+            }
+            mO(q_idx_offset + q, d, h) = ElementO(acc);
+          }
+        }
+      }
+    }
+  }
+};
+
+/* https://github.com/ColfaxResearch/cutlass-kernels/blob/master/include/utils/fmha_cutlass.hpp */
+template <typename Element, typename AccumType>
+class TestAttention {
+public:
+  //
+  // Type definitions
+  //
+  static constexpr float kLog2e = float(1.4426950408889634074); // log_2(e) = M_LOG2E
+  using ElementQ = Element;
+  using ElementK = Element;
+  using ElementS = Element;
+  using ElementP = AccumType;
+  using ElementAccumulator = AccumType;
+  using ElementV = Element;
+  using ElementO = Element;
+
+  using ElementCompute = AccumType;
+
+  using ElementNorm = AccumType;
+  using ElementSum = AccumType;
+  using ElementSoftmaxCompute = AccumType;
+
+  using LayoutQ = mutlass::layout::RowMajor;
+  using LayoutK = mutlass::layout::ColumnMajor;
+  using LayoutP = mutlass::layout::RowMajor;
+  using LayoutV = mutlass::layout::RowMajor;
+  using LayoutO = mutlass::layout::RowMajor;
+  using LayoutNorm = mutlass::layout::RowMajor;
+
+  using MatrixCoord = typename LayoutP::TensorCoord;
+
+private:
+  bool help;
+  bool error;
+  bool reference_check;
+  bool use_mask;
+  bool causal;
+
+  int alignment;
+  int iterations;
+
+  int QSeqLen;
+  int HeadDimQK;
+  int HeadDimV;
+  int KVSeqLen;
+  int HeadNumQ;
+  int HeadNumKV;
+  int BatchSize;
+
+  // alpha0, alpha1 and beta are fixed
+  // in this multi-head attention example
+  float alpha0;
+  float alpha1;
+  float beta;
+
+  mutlass::gemm::GemmCoord problem0;
+  mutlass::gemm::GemmCoord problem1;
+  //
+
+  int64_t ldq;
+  int64_t ldk;
+  int64_t ldp;
+  int64_t ldv;
+  int64_t ldo;
+
+  int64_t elements_Q;
+  int64_t elements_K;
+  int64_t elements_V;
+  int64_t elements_P;
+  int64_t elements_O;
+  int64_t elements_norm;
+
+public:
+  //
+  // Methods
+  //
+
+  TestAttention(int _head_number_q, int _head_number_kv, int _batch_size, int _head_dim_qk, int _head_dim_v,
+                int _seq_length_q, int _seq_length_kv, int _alignment = 1, bool use_mask = false,
+                bool _causal = false) {
+
+    QSeqLen = _seq_length_q;
+    HeadDimQK = _head_dim_qk;
+    HeadDimV = _head_dim_v;
+    KVSeqLen = _seq_length_kv;
+    BatchSize = _batch_size;
+    HeadNumQ = _head_number_q;
+    HeadNumKV = _head_number_kv;
+    causal = _causal;
+
+    mutlass::gemm::GemmCoord _problem0(QSeqLen, KVSeqLen, HeadDimQK);
+    mutlass::gemm::GemmCoord _problem1(QSeqLen, HeadDimV, KVSeqLen);
+
+    problem0 = _problem0;
+    problem1 = _problem1;
+  }
+
+public:
+  /// Initializes data structures
+  void initialize() {
+
+    //
+    // Set scaling factor(s).
+    //
+    alpha0 = 1.0f / sqrt(float(HeadDimQK));
+    alpha1 = 1.0f;
+    beta = 0;
+
+    elements_Q = problem0.m() * problem0.k();
+    elements_K = problem0.n() * problem0.k();
+    elements_P = problem0.m() * problem0.n();
+    elements_V = problem1.n() * problem1.k();
+    elements_O = problem1.m() * problem1.n();
+    elements_norm = problem0.m();
+
+    ldq = LayoutQ::packed({problem0.m(), problem0.k()}).stride(0);
+    ldk = LayoutK::packed({problem0.k(), problem0.n()}).stride(0);
+    ldp = LayoutP::packed({problem0.m(), problem0.n()}).stride(0);
+    ldv = LayoutV::packed({problem1.k(), problem1.n()}).stride(0);
+    ldo = LayoutO::packed({problem1.m(), problem1.n()}).stride(0);
+  }
+
+
+  void compute(const ElementQ *Q, const ElementK *K, const ElementV *V,
+              ElementS *S, ElementO *O, ElementNorm *norm, ElementSum *sum,
+              bool usePow2 = false, bool usePreScaling = true, bool with_softmax = false) {
+    LayoutQ layout_Q(ldq);
+    LayoutK layout_K(ldk);
+    LayoutP layout_P(ldp);
+    LayoutV layout_V(ldv);
+    LayoutO layout_O(ldo);
+
+    LayoutNorm layout_norm(1);
+
+    MatrixCoord extent_Q{problem0.m(), problem0.k()};
+    MatrixCoord extent_K{problem0.k(), problem0.n()};
+    MatrixCoord extent_P{problem0.m(), problem0.n()};
+    MatrixCoord extent_V{problem1.k(), problem1.n()};
+    MatrixCoord extent_O{problem1.m(), problem1.n()};
+    MatrixCoord extent_norm{problem1.m(), 1};
+
+
+#if DEVICE_REF
+    mutlass::DeviceAllocation<ElementP> softmaxP(layout_P.capacity(extent_P));
+    mutlass::TensorView<ElementP, LayoutP> softmaxViewP(softmaxP.get(),
+                                                        layout_P, extent_P);
+#else
+    std::vector<ElementP> softmaxP(layout_P.capacity(extent_P));
+    mutlass::TensorView<ElementP, LayoutP> softmaxViewP(softmaxP.data(),
+                                                        layout_P, extent_P);
+
+#endif
+
+    for (int64_t b = 0; b < BatchSize; ++b) {
+      for (int64_t h = 0; h < HeadNumQ; ++h) {
+        int hkv = h%HeadNumKV;
+        auto offsetQ = Q + b * HeadNumQ * elements_Q + h * elements_Q;
+        auto offsetK = K + b * HeadNumKV * elements_K + hkv * elements_K;
+        auto offsetS = S + b * HeadNumQ * elements_P + h * elements_P;
+        auto offsetV = V + b * HeadNumKV * elements_V + hkv * elements_V;
+        auto offsetO = O + b * HeadNumQ * elements_O + h * elements_O;
+
+        auto offsetNorm = norm + b * HeadNumQ * elements_norm + h * elements_norm;
+        auto offsetSum  =  sum + b * HeadNumQ * elements_norm + h * elements_norm;
+        mutlass::TensorView<ElementQ, LayoutQ> view_Q(
+            const_cast<ElementQ *>(offsetQ), layout_Q, extent_Q);
+        mutlass::TensorView<ElementK, LayoutK> view_K(
+            const_cast<ElementK *>(offsetK), layout_K, extent_K);
+        mutlass::TensorView<ElementV, LayoutV> view_V(
+            const_cast<ElementK *>(offsetV), layout_V, extent_V);
+
+        mutlass::TensorView<ElementS, LayoutP> view_S(offsetS, layout_P, extent_P);
+        mutlass::TensorView<ElementO, LayoutO> view_O(offsetO, layout_O,
+                                                      extent_O);
+
+        mutlass::TensorView<ElementNorm, LayoutNorm> view_Norm_Ref(
+            offsetNorm, layout_norm, extent_norm);
+        mutlass::TensorView<ElementSum, LayoutNorm> view_Sum_Ref(
+            offsetSum, layout_norm, extent_norm);
+
+
+        // Reference GEMM-I.
+
+        float alphaMma0 = usePreScaling ? alpha0 : 1.0f;
+        float postScaling = usePreScaling ? 1.0f : alpha0;
+
+#if DEVICE_REF
+        mutlass::DeviceAllocation<ElementP> block_Ref_P(
+            layout_P.capacity(extent_P));
+        mutlass::TensorView<ElementP, LayoutP> view_P(block_Ref_P.get(),
+                                                      layout_P, extent_P);
+
+        mutlass::reference::device::GemmComplex<
+            ElementQ, LayoutQ, ElementK, LayoutK, ElementP, LayoutP,
+            ElementAccumulator, ElementCompute>(
+            problem0, ElementAccumulator(alphaMma0), view_Q,
+            mutlass::ComplexTransform::kNone, view_K,
+            mutlass::ComplexTransform::kNone, ElementAccumulator(beta), view_P,
+            view_P, ElementAccumulator(0));
+
+        // Compute softmax for P. We need to explicitly compute softmax
+        // over P because softmax is fused to the second GEMM in the
+        // profiled implementation.
+        std::vector<ElementP> matrix_Ref(layout_P.capacity(extent_P));
+        mutlass::device_memory::copy_to_host(
+            matrix_Ref.data(), block_Ref_P.get(), matrix_Ref.size());
+        mutlass::TensorView<ElementP, LayoutP> view_Ref_host(
+            matrix_Ref.data(), layout_P, extent_P);
+
+        std::vector<ElementS> matrix_Ref_S(layout_P.capacity(extent_P));
+        for (int i = 0; i < matrix_Ref.size(); ++i) {
+          matrix_Ref_S.at(i) = ElementS(matrix_Ref.at(i));
+        }
+        // Copy to S (for debugging).
+        mutlass::device_memory::copy_to_device(offsetS, matrix_Ref_S.data(),
+                                              matrix_Ref_S.size());
+#else
+
+        std::vector<ElementP> block_Ref_P(
+            layout_P.capacity(extent_P));
+        mutlass::TensorView<ElementP, LayoutP> view_P(block_Ref_P.data(),
+                                                      layout_P, extent_P);
+        mutlass::reference::host::GemmComplex<
+            ElementQ, LayoutQ, ElementK, LayoutK, ElementP, LayoutP,
+            ElementAccumulator, ElementCompute>(
+            problem0, ElementAccumulator(alphaMma0), view_Q,
+            mutlass::ComplexTransform::kNone, view_K,
+            mutlass::ComplexTransform::kNone, ElementAccumulator(beta), view_P,
+            view_P, ElementAccumulator(0));
+
+        // Compute softmax for P. We need to explicitly compute softmax
+        // over P because softmax is fused to the second GEMM in the
+        // profiled implementation.
+        std::vector<ElementP> matrix_Ref(layout_P.capacity(extent_P));
+        mutlass::device_memory::copy_host_to_host(
+            matrix_Ref.data(), block_Ref_P.data(), matrix_Ref.size());
+        mutlass::TensorView<ElementP, LayoutP> view_Ref_host(
+            matrix_Ref.data(), layout_P, extent_P);
+
+        std::vector<ElementS> matrix_Ref_S(layout_P.capacity(extent_P));
+        for (int i = 0; i < matrix_Ref.size(); ++i) {
+          matrix_Ref_S.at(i) = ElementS(matrix_Ref.at(i));
+        }
+        // Copy to S (for debugging).
+        mutlass::device_memory::copy_host_to_host(offsetS, matrix_Ref_S.data(),
+                                                  matrix_Ref_S.size());
+
+#endif
+        if (causal) {
+          for (int m = 0; m < problem0.m(); ++m) {
+            int n_limit = m + 1;
+            for (int n = 0; n < problem0.n(); ++n) {
+              if (n >= n_limit) {
+                view_Ref_host.ref().at({m, n}) = -std::numeric_limits<float>::infinity();
+              }
+            }
+          }
+        }
+
+        if (with_softmax) {
+          int n_dim_row = problem0.n();
+
+          // Compute softmax for reference matrix
+          // Run this branch
+          if (usePow2) {
+            for (int m = 0; m < problem0.m(); m++) {
+              for (int n = 0; n < n_dim_row; n++) {
+                view_Ref_host.ref().at({m, n}) =
+                    kLog2e * postScaling * view_Ref_host.ref().at({m, n});
+              }
+
+
+              ElementSoftmaxCompute max =
+                  ElementSoftmaxCompute(view_Ref_host.ref().at({m, uint64_t(0)}));
+              for (int n = 1; n < n_dim_row; n++) {
+                max = std::max(
+                    max, ElementSoftmaxCompute(view_Ref_host.ref().at({m, n})));
+              }
+
+
+              view_Norm_Ref.at({m, uint64_t(0)}) = ElementNorm(max);
+
+              ElementSoftmaxCompute sum = ElementSoftmaxCompute();
+              for (int n = 0; n < n_dim_row; n++) {
+                sum += std::exp2f(
+                    ElementSoftmaxCompute(view_Ref_host.ref().at({m, n})) - max);
+              }
+              ElementSoftmaxCompute inv_sum = ElementSoftmaxCompute(1.0f / sum);
+
+              view_Sum_Ref.at({m, uint64_t(0)}) = ElementSum(inv_sum);
+
+              for (int n = 0; n < n_dim_row; n++) {
+                view_Ref_host.ref().at({m, n}) =
+                    ElementP(std::exp2f(ElementSoftmaxCompute(
+                                            view_Ref_host.ref().at({m, n})) -
+                                        max) *
+                            inv_sum);
+              }
+            }
+          } else {
+            for (int64_t m = 0; m < problem0.m(); m++) {
+              ElementSoftmaxCompute max =
+                  ElementSoftmaxCompute(view_Ref_host.ref().at({m, uint64_t(0)}));
+              for (int64_t n = 1; n < n_dim_row; n++) {
+                max = std::max(
+                    max, ElementSoftmaxCompute(view_Ref_host.ref().at({m, n})));
+              }
+
+              view_Norm_Ref.at({m, uint64_t(0)}) = ElementNorm(max);
+
+              ElementSoftmaxCompute sum = ElementSoftmaxCompute();
+              for (int64_t n = 0; n < n_dim_row; n++) {
+                sum += std::exp(
+                    ElementSoftmaxCompute(view_Ref_host.ref().at({m, n})) - max);
+              }
+              ElementSoftmaxCompute inv_sum = ElementSoftmaxCompute(1.0f / sum);
+
+              view_Sum_Ref.at({m, uint64_t(0)}) = ElementSum(inv_sum);
+
+              for (int64_t n = 0; n < n_dim_row; n++) {
+                view_Ref_host.ref().at({m, n}) =
+                    ElementP(std::exp(ElementSoftmaxCompute(
+                                          view_Ref_host.ref().at({m, n})) -
+                                      max) *
+                            inv_sum);
+              }
+            }
+          }
+        }
+
+#if DEVICE_REF
+        mutlass::device_memory::copy_to_device(
+            block_Ref_P.get(), matrix_Ref.data(), matrix_Ref.size());
+
+        mutlass::reference::device::GemmComplex<
+            ElementP, LayoutP, ElementV, LayoutV, ElementO, LayoutO,
+            ElementAccumulator, ElementCompute>(
+            problem1, ElementAccumulator(alpha1), view_P,
+            mutlass::ComplexTransform::kNone, view_V,
+            mutlass::ComplexTransform::kNone, ElementAccumulator(beta), view_O,
+            view_O, ElementAccumulator(0));
+#else
+
+        mutlass::device_memory::copy_host_to_host(
+            block_Ref_P.data(), matrix_Ref.data(), matrix_Ref.size());
+        mutlass::reference::host::GemmComplex<
+            ElementP, LayoutP, ElementV, LayoutV, ElementO, LayoutO,
+            ElementAccumulator, ElementCompute>(
+            problem1, ElementAccumulator(alpha1), view_P,
+            mutlass::ComplexTransform::kNone, view_V,
+            mutlass::ComplexTransform::kNone, ElementAccumulator(beta), view_O,
+            view_O, ElementAccumulator(0));
+#endif
+      }
+    }
+  }
+};
diff --git a/vllm_musa/csrc/flash_attn/impl/utils.hpp b/vllm_musa/csrc/flash_attn/impl/utils.hpp
new file mode 100644
index 000000000..4f53c9ef7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/impl/utils.hpp
@@ -0,0 +1,41 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+
+template <
+  bool A_in_regs = false,
+  bool B_in_regs = false,
+  typename Tensor0,
+  typename Tensor1,
+  typename Tensor2,
+  typename Tensor3,
+  typename Tensor4,
+  typename TiledMma,
+  typename ThrCopyA,
+  typename ThrCopyB
+>
+MUTE_DEVICE
+void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsA,
+          Tensor4 const& tCsB, TiledMma tiled_mma,
+          ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) {
+  MUTE_STATIC_ASSERT_V(mute::size<1>(tCrA) == mute::size<1>(acc));
+  MUTE_STATIC_ASSERT_V(mute::size<1>(tCrB) == mute::size<2>(acc));
+  MUTE_STATIC_ASSERT_V(mute::size<2>(tCrA) == mute::size<2>(tCrB));
+  mute::Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
+  MUTE_STATIC_ASSERT_V(mute::size<1>(tCsA) == mute::size<1>(tCrA_copy_view));
+  mute::Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+  MUTE_STATIC_ASSERT_V(mute::size<1>(tCsB) == mute::size<1>(tCrB_copy_view));
+
+  if constexpr (!A_in_regs) { copy(tCsA(mute::_,mute::_,mute::_0{}), tCrA_copy_view(mute::_,mute::_,mute::_0{})); }
+  if constexpr (!B_in_regs) { copy(tCsB(mute::_,mute::_,mute::_0{}), tCrB_copy_view(mute::_,mute::_,mute::_0{})); }
+
+  MUTE_UNROLL
+  for (int i = 0; i < mute::size<2>(tCrA); ++i) {
+    if (i < mute::size<2>(tCrA) - 1) {
+      if constexpr (!A_in_regs) { copy(tCsA(mute::_,mute::_,i + 1), tCrA_copy_view(mute::_,mute::_,i + 1)); }
+      if constexpr (!B_in_regs) { copy(tCsB(mute::_,mute::_,i + 1), tCrB_copy_view(mute::_,mute::_,i + 1)); }
+    }
+
+    mute::gemm(tiled_mma, tCrA(mute::_,mute::_,i), tCrB(mute::_,mute::_,i), acc);
+  }
+}
diff --git a/vllm_musa/csrc/flash_attn/include/mute/address_space.hpp b/vllm_musa/csrc/flash_attn/include/mute/address_space.hpp
new file mode 100644
index 000000000..73d6cd5b3
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/address_space.hpp
@@ -0,0 +1,49 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+namespace mute {
+
+enum class AddressSpace {
+  Generic = 0,
+  Global  = 1,
+  Shared  = 3,
+};
+
+template <AddressSpace AS>
+MUTE_HOST_DEVICE constexpr
+void __attribute__((address_space(static_cast<int>(AS))))*
+make_ptr_with_address_space(uint64_t ptr) {
+  return reinterpret_cast<void __attribute__((address_space(static_cast<int>(AS))))*>(ptr);
+}
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/axpby.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/axpby.hpp
new file mode 100644
index 000000000..69869dfb7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/axpby.hpp
@@ -0,0 +1,96 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/tensor.hpp>
+#include <mute/tensor_predicate.hpp>
+
+namespace mute
+{
+
+//
+// Accept mutable temporaries
+//
+template <class Alpha,
+          class XEngine, class XLayout,
+          class Beta,
+          class YEngine, class YLayout,
+          class PrdTensor = TrivialPredTensor>
+MUTE_HOST_DEVICE
+void
+axpby(Alpha                    const& alpha,
+      Tensor<XEngine, XLayout> const& x,
+      Beta                     const& beta,
+      Tensor<YEngine, YLayout>     && y,
+      PrdTensor                const& p = {})
+{
+  return axpby(alpha, x, beta, y, p);
+}
+
+//
+// AXPBY
+//
+template <class Alpha,
+          class XEngine, class XLayout,
+          class Beta,
+          class YEngine, class YLayout,
+          class PrdTensor = TrivialPredTensor>
+MUTE_HOST_DEVICE
+void
+axpby(Alpha                    const& alpha,
+      Tensor<XEngine, XLayout> const& x,
+      Beta                     const& beta,
+      Tensor<YEngine, YLayout>      & y,
+      PrdTensor                const& p = {})
+{
+  auto isBetaZero = [&] () {
+    if constexpr (is_complex<Beta>::value) {
+      return beta.real() == Int<0>{} && beta.imag() == Int<0>{};
+    }
+    else {
+      return beta == Int<0>{};
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  } ();
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(x); ++i) {
+    if (p(i)) {
+      y(i) = (isBetaZero ? alpha * x(i) : alpha * x(i) + beta * y(i));
+    }
+  }
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/clear.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/clear.hpp
new file mode 100644
index 000000000..e5ad87ef4
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/clear.hpp
@@ -0,0 +1,67 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/tensor.hpp>
+
+#include <mute/algorithm/fill.hpp>
+
+namespace mute
+{
+
+//
+// Accept mutable temporaries
+//
+template <class Engine, class Layout>
+MUTE_HOST_DEVICE
+void
+clear(Tensor<Engine, Layout>&& tensor)
+{
+  return clear(tensor);
+}
+
+//
+// Set elements to zero
+//
+template <class Engine, class Layout>
+MUTE_HOST_DEVICE
+void
+clear(Tensor<Engine, Layout>& tensor)
+{
+  using T = typename Tensor<Engine,Layout>::value_type;
+
+  fill(tensor, T{});
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/cooperative_copy.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/cooperative_copy.hpp
new file mode 100644
index 000000000..6a5ef7a49
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/cooperative_copy.hpp
@@ -0,0 +1,340 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+#include <mute/layout.hpp>
+#include <mute/layout_composed.hpp> // mute::logical_divide
+#include <mute/swizzle.hpp>         // mute::Swizzle
+#include <mute/swizzle_layout.hpp>  // mute::get_nonswizzle_portion
+#include <mute/tensor.hpp>          // mute::Tensor
+#include <mute/tensor_predicate.hpp>
+#include <mute/algorithm/copy.hpp>
+#include <mute/atom/copy_atom.hpp>
+
+namespace mute
+{
+
+template <uint32_t NumThreads,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE void
+naive_cooperative_copy(uint32_t                     const& tid,
+                       Tensor<SrcEngine, SrcLayout> const& src,
+                       Tensor<DstEngine, DstLayout>      & dst)
+{
+  auto N = size(dst);
+  auto R = N % Int<NumThreads>{};
+  if (R > 0 && tid < R) {                                             // Likely static condition && Residue in-bounds
+    dst[tid] = src[tid];
+  }
+  MUTE_UNROLL
+  for (uint32_t i = uint32_t(R); i < uint32_t(N); i += NumThreads) {  // All in-bounds
+    dst[tid + i] = src[tid + i];
+  }
+}
+
+// Accept mutable temporaries
+template <uint32_t NumThreads,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE void
+naive_cooperative_copy(uint32_t                     const& tid,
+                       Tensor<SrcEngine, SrcLayout> const& src,
+                       Tensor<DstEngine, DstLayout>     && dst)
+{
+  return naive_cooperative_copy(tid, src, dst);
+}
+
+// A heuristic to determine a "good" permutation of two tensors for later vectorization and thr-assignment
+template <class AEngine, class ALayout,
+          class BEngine, class BLayout>
+MUTE_HOST_DEVICE constexpr
+auto
+heuristic_permutation(Tensor<AEngine, ALayout> const& a,
+                      Tensor<BEngine, BLayout> const& b)
+{
+  constexpr bool swizzleA = get_swizzle_t<AEngine>::num_bits != 0 or
+                            get_swizzle_t<ALayout>::num_bits != 0;
+  constexpr bool swizzleB = get_swizzle_t<BEngine>::num_bits != 0 or
+                            get_swizzle_t<BLayout>::num_bits != 0;
+  auto a_inv = right_inverse(get_nonswizzle_portion(a.layout()));
+  auto b_inv = right_inverse(get_nonswizzle_portion(b.layout()));
+
+  constexpr uint8_t scoreA = (uint8_t(swizzleA)                  << 2) |
+                             (uint8_t(is_smem<AEngine>::value)   << 1) |
+                             (uint8_t(size(a_inv) > size(b_inv)) << 0);
+
+  constexpr uint8_t scoreB = (uint8_t(swizzleB)                  << 2) |
+                             (uint8_t(is_smem<BEngine>::value)   << 1) |
+                             (uint8_t(size(b_inv) > size(a_inv)) << 0);
+
+  if constexpr (scoreA >= scoreB) {
+    return a_inv;
+  } else {
+    return b_inv;
+  }
+}
+
+// cooperative_copy<NumThreads, MaxVecBits>(thr_idx, src, dst)
+// Use NumThreads to copy Tensor src to Tensor dst with element-wise vectorization up to MaxVecBits.
+// @pre 0 <= @a tid < NumThreads
+// @pre Tensors @a src and @a dst are aligned up to MaxVecBits.
+//      That is, pointers and dynamic strides are assumed to be aligned up to MaxVecBits.
+//
+template <uint32_t NumThreads, uint32_t MaxVecBits,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout,
+          class CopyPolicy = DefaultCopy>
+MUTE_HOST_DEVICE
+void
+cooperative_copy(uint32_t                     const& tid,
+                 Tensor<SrcEngine, SrcLayout> const& src,
+                 Tensor<DstEngine, DstLayout>      & dst,
+                 CopyPolicy                   const& cpy = {})
+{
+  // Assumes the shapes are static, can generalize/fallback
+  MUTE_STATIC_ASSERT_V(is_static<decltype(shape(src))>{} && is_static<decltype(shape(dst))>{});
+  MUTE_STATIC_ASSERT_V(size(src) == size(dst));
+  // Assumes the types are the same, can generalize/fallback
+  static_assert(mute::is_same<typename SrcEngine::value_type, typename DstEngine::value_type>::value);
+  static_assert(MaxVecBits == sizeof_bits_v<typename SrcEngine::value_type> ||
+                MaxVecBits == 8 || MaxVecBits == 16 || MaxVecBits == 32 || MaxVecBits == 64 || MaxVecBits == 128,
+                "Expected MaxVecBits to be value size or 8 or 16 or 32 or 64 or 128 for alignment and performance.");
+  // Check that the tensors are likely shared across threads: either gmem or smem
+  static_assert((is_gmem<SrcEngine>::value || is_smem<SrcEngine>::value),
+                "cooperative_copy expects shared gmem or smem source tensor.");
+  static_assert((is_gmem<DstEngine>::value || is_smem<DstEngine>::value),
+                "cooperative_copy expects shared gmem or smem destination tensor.");
+  // Precondition on tid in DEBUG
+  assert(tid < NumThreads);
+  // Precondition on pointer alignment in DEBUG
+  assert(is_byte_aligned<ceil_div(MaxVecBits,8u)>(raw_pointer_cast(src.data())));
+  assert(is_byte_aligned<ceil_div(MaxVecBits,8u)>(raw_pointer_cast(dst.data())));
+
+#if 0
+      if (thread0()) {
+        print("   "); print("cooperative_copy\n");
+        print("   "); print("NumThreads: "); print(NumThreads); print("\n");
+        print("   "); print("MaxVecBits: "); print(MaxVecBits); print("\n");
+        print("   "); print("src: "); print(src); print("\n");
+        print("   "); print("dst: "); print(dst); print("\n");
+      }
+#ifdef __MUSA_ARCH__
+      __syncthreads();
+#endif
+#endif
+
+  // The common layout of the two tensors that can be vectorized over elements and threads
+  // vidx -> coord
+  auto common_layout = heuristic_permutation(src, dst);
+
+  // Apply
+  // (V, rest)
+  Tensor src_a = coalesce(logical_divide(src, common_layout), Shape<_1,_1>{});
+  Tensor dst_a = coalesce(logical_divide(dst, common_layout), Shape<_1,_1>{});
+
+  //
+  // Determine vectorization of elems and thrs based on src/dst size and number of threads
+  // NOTE: This heuristic promotes parallelization over vectorization
+  //
+
+  // The number of elements and number of bits
+  constexpr int  elem_bits = sizeof_bits_v<typename SrcEngine::value_type>;
+  constexpr int total_elem = size(SrcLayout{});
+
+  // The number of elements that can be vectorized in values
+  constexpr int common_elem = decltype(max_common_vector(src_a, dst_a))::value;
+
+#if 0
+      if (thread0()) {
+        print("   "); print("common_layout: "); print(common_layout); print("\n");
+        print("   "); print("src_a: "); print(src_a); print("\n");
+        print("   "); print("dst_a: "); print(dst_a); print("\n");
+      }
+#ifdef __MUSA_ARCH__
+      __syncthreads();
+#endif
+#endif
+
+  //
+  if constexpr (total_elem % NumThreads != 0) {
+    // Not attempting to find a partitioning pattern, fallback to dynamically indexed slowpath
+
+    if constexpr (common_elem > 1 && MaxVecBits > elem_bits) {
+      // If the vectorization is non-trivial and divides the maximum vectorizations, then vectorize
+      constexpr auto max_align_src = elem_bits * decltype(max_alignment(src_a.layout()))::value;
+      constexpr auto max_align_dst = elem_bits * decltype(max_alignment(dst_a.layout()))::value;
+      constexpr auto vec_bits = gcd(max_align_src, max_align_dst, MaxVecBits);
+      using VecType = uint_bit_t<vec_bits>;
+
+      static_assert(vec_bits % elem_bits == 0, "Expected divisibility");
+      static_assert((vec_bits >= 8), "No support for subbyte copying");
+
+      Tensor src_v = recast<VecType const>(src_a);
+      Tensor dst_v = recast<VecType      >(dst_a);
+
+#if 0
+      if (thread0()) {
+        print("   "); print("cooperative_copy -- naive\n");
+        print("   "); print("src_v: "); print(src_v); print("\n");
+        print("   "); print("dst_v: "); print(dst_v); print("\n");
+      }
+#ifdef __MUSA_ARCH__
+      __syncthreads();
+#endif
+#endif
+
+      naive_cooperative_copy<NumThreads>(tid, src_v, dst_v);
+    } else {
+      naive_cooperative_copy<NumThreads>(tid, src_a, dst_a);
+    }
+  } else {
+    // If the tensors can be equally partitioned by the threads,
+    // compute vectorization widths in elements and threads.
+
+    // If there are too many threads to allow a full vectorized copy, trunc the vectorization
+    constexpr int total_bits = total_elem * elem_bits;
+    constexpr int max_bits_per_thr = total_bits / NumThreads;
+    // At least elem_bits, at most common_bits
+    constexpr int common_bits = common_elem * elem_bits;
+    constexpr int vec_bits = mute::max(elem_bits, mute::gcd(common_bits, int(MaxVecBits), max_bits_per_thr));
+
+    // Should account for vec_bits < 8 and/or vec_elem <= 1
+    // And also account for subbyte types, which could cause race conditions
+    // Want to ENFORCE sufficient vectorization in those cases
+    static_assert(vec_bits % elem_bits == 0, "Expected divisibility");
+    static_assert(vec_bits >= 8, "No support for subbyte copying");
+
+    using VecType = uint_bit_t<vec_bits>;
+    constexpr int vec_elem = vec_bits / elem_bits;
+
+    constexpr int vec_thrs = mute::min(int(NumThreads), total_elem / vec_elem);
+
+    //
+    // Determine the partitioning patterns for the vec_elems and vec_thrs
+    //
+
+    // Distribute the rest of the V*T to some consistent portion outside of the common_layout, if needed
+    auto common_domain_src = domain_distribute(shape(src_a), Int<vec_elem*vec_thrs>{});
+    auto common_domain_dst = domain_distribute(shape(dst_a), Int<vec_elem*vec_thrs>{});
+
+    // Make sure for now, could fall back here instead
+    MUTE_STATIC_ASSERT_V(size(common_domain_src) == Int<vec_elem*vec_thrs>{});
+    MUTE_STATIC_ASSERT_V(compatible(common_domain_src, common_domain_dst) ||
+                         compatible(common_domain_dst, common_domain_src));
+    // Use the "more specific" domain for the extra elements of V*T
+    auto common_domain = conditional_return(compatible(common_domain_src, common_domain_dst),
+                                            common_domain_dst, common_domain_src);
+
+    // Construct the tiler
+    auto tiler_vt = common_domain.with_shape(Int<vec_elem>{}, Int<vec_thrs>{});
+
+    // Apply and slice
+    Tensor src_v = logical_divide(src_a, tiler_vt)(make_coord(_,tid),_);
+    Tensor dst_v = logical_divide(dst_a, tiler_vt)(make_coord(_,tid),_);
+
+#if 0
+      if (thread0()) {
+        print("   "); print("cooperative_copy -- vec\n");
+        print("   "); print("Used vector: ");  print(vec_elem); print("\n");
+        print("   "); print("Used threads: ");  print(vec_thrs); print("\n");
+        print("   "); print("tiler_vt: "); print(tiler_vt); print("\n");
+        print("   "); print("src_v: "); print(src_v); print("\n");
+        print("   "); print("dst_v: "); print(dst_v); print("\n");
+        print("   "); print("recast<VecType const>(src_v): "); print(recast<VecType const>(src_v)); print("\n");
+        print("   "); print("recast<VecType      >(dst_v): "); print(recast<VecType      >(dst_v)); print("\n");
+      }
+#ifdef __MUSA_ARCH__
+      __syncthreads();
+#endif
+#endif
+
+    // If we're using all threads (static) or the tid is in-range (dynamic)
+    if (vec_thrs == NumThreads or tid < vec_thrs) {
+      auto src_c = recast<VecType const>(src_v);
+      auto dst_c = recast<VecType>(dst_v);
+      return copy(cpy, src_c, dst_c);
+    }
+  }
+}
+
+
+// Default max-vectorization size to value_type size
+template <uint32_t NumThreads,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout,
+          class CopyPolicy = DefaultCopy>
+MUTE_HOST_DEVICE
+void
+cooperative_copy(uint32_t                     const& tid,
+                 Tensor<SrcEngine, SrcLayout> const& src,
+                 Tensor<DstEngine, DstLayout>      & dst,
+                 CopyPolicy                   const& cpy = {})
+{
+  constexpr uint32_t MaxVecBits = sizeof_bits_v<typename SrcEngine::value_type>;
+  return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst, cpy);
+}
+
+//
+// Accept mutable temporaries
+//
+
+template <uint32_t NumThreads,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout,
+          class CopyPolicy = DefaultCopy>
+MUTE_HOST_DEVICE
+void
+cooperative_copy(uint32_t                     const& tid,
+                 Tensor<SrcEngine, SrcLayout> const& src,
+                 Tensor<DstEngine, DstLayout>     && dst,
+                 CopyPolicy                   const& cpy = {})
+{
+  return cooperative_copy<NumThreads>(tid, src, dst, cpy);
+}
+
+template <uint32_t NumThreads, uint32_t MaxVecBits,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout,
+          class CopyPolicy = DefaultCopy>
+MUTE_HOST_DEVICE
+void
+cooperative_copy(uint32_t                     const& tid,
+                 Tensor<SrcEngine, SrcLayout> const& src,
+                 Tensor<DstEngine, DstLayout>     && dst,
+                 CopyPolicy                   const& cpy = {})
+{
+  return cooperative_copy<NumThreads, MaxVecBits>(tid, src, dst, cpy);
+}
+
+} // end namespace mute
\ No newline at end of file
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/cooperative_gemm.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/cooperative_gemm.hpp
new file mode 100644
index 000000000..42c759a35
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/cooperative_gemm.hpp
@@ -0,0 +1,586 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+#include <mute/util/type_traits.hpp>
+
+#include <mute/atom/mma_atom.hpp>
+
+#include <mute/algorithm/axpby.hpp>
+#include <mute/algorithm/functional.hpp>
+#include <mute/algorithm/gemm.hpp>
+
+#include <mute/tensor.hpp>
+
+namespace mute
+{
+
+//
+// Cooperative Shared-Memory GEMMs
+//
+
+namespace detail {
+
+// Slow fallback path:
+template<typename ... Args,
+         typename Alpha, typename TRC, typename RCLayout,
+         typename Beta, class TSC, typename CLayout, typename SCLayout,
+         typename CLoadTransformOp, typename CStoreTransformOp>
+MUTE_HOST_DEVICE
+void
+epilogue_predication(ThrMMA<Args...>    const& thr_mma,
+                     Alpha              const& alpha,
+                     Tensor<TRC, RCLayout>   & tCrC,
+                     Beta               const& beta,
+                     Tensor<TSC, CLayout>    & sC,
+                     Tensor<TSC, SCLayout>   & tCsC,
+                     CLoadTransformOp   const& sC_load_op,  // transforms C values before use in GEMM
+                     CStoreTransformOp  const& sC_store_op) // transforms results before they are stored to C
+{
+  using InputTypeC   = typename TSC::value_type;
+  using ComputeTypeC = typename ThrMMA<Args...>::ValTypeC;
+  MUTE_STATIC_ASSERT(MUTE_STL_NAMESPACE::is_same_v<ComputeTypeC, typename TRC::value_type>);
+
+  // Create coordinate tensors for the problem
+  Tensor cC   = make_identity_tensor(shape(sC));                     // (M,N) -> (m,n)
+  // Repeat partitioning with thr_mma
+  Tensor tCcC = thr_mma.partition_C(cC);                             // (MMA,MMA_M,MMA_N) -> (m,n)
+
+  const bool isBetaZero = [&] () {
+    if constexpr (is_complex<Beta>::value) {
+      return beta.real() == Int<0>{} && beta.imag() == Int<0>{};
+    }
+    else {
+      return beta == Int<0>{};
+    }
+    MUTE_GCC_UNREACHABLE;
+  } ();
+
+  // Custom axpby_if for now
+  MUTE_UNROLL
+  for (int i = 0; i < size(tCrC); ++i)
+  {
+    if (elem_less(tCcC(i), shape(sC)))
+    {
+      tCsC(i) = sC_store_op(isBetaZero ? alpha * tCrC(i)
+                                       : alpha * tCrC(i) +
+                                          beta * static_cast<ComputeTypeC>(sC_load_op(tCsC(i))));
+    }
+  }
+}
+
+template<class Alpha, class TRC, class RCLayout,
+         class Beta, class TSC, class SCLayout,
+         class CLoadTransformOp, class CStoreTransformOp,
+         class SmemCopyOpC>
+MUTE_HOST_DEVICE
+void
+epilogue_no_predication(Alpha              const& alpha,
+                        Tensor<TRC, RCLayout>   & tCrC,
+                        Beta               const& beta,
+                        Tensor<TSC, SCLayout>   & tCsC,
+                        CLoadTransformOp   const& sC_load_op,  // transforms C values before use in GEMM
+                        CStoreTransformOp  const& sC_store_op, // transforms results before they are stored to C
+                        SmemCopyOpC        const& sC_copy_op)
+{
+  using InputTypeC   = typename TSC::value_type;
+  using ComputeTypeC = typename TRC::value_type;
+
+  const bool isBetaZero = [&] () {
+    if constexpr (is_complex<Beta>::value) {
+      return beta.real() == Int<0>{} && beta.imag() == Int<0>{};
+    }
+    else {
+      return beta == Int<0>{};
+    }
+    MUTE_GCC_UNREACHABLE;
+  } ();
+
+  Tensor tCrDi = make_fragment_like(tCsC);
+  Tensor tCrD = make_fragment_like(tCrC);
+  if(!isBetaZero) {
+    copy(sC_copy_op, tCsC, tCrDi);
+    // Transform C on/after load
+    mute::transform(tCrDi, tCrD, sC_load_op);
+  }
+  // C = alpha * (A * B) + beta * C
+  axpby(alpha, tCrC, beta, tCrD);
+  // Transform C before/on store
+  mute::transform(tCrD, tCrDi, sC_store_op);
+  copy(sC_copy_op, tCrDi, tCsC);
+}
+
+// Predicated Cooperative GEMM
+template <class... Args,
+          class TA, class ALayout, class TB, class BLayout,
+          class TC, class RCLayout,
+          class ALoadTransformOp, class BLoadTransformOp>
+MUTE_HOST_DEVICE
+void
+cooperative_gemm_predication(ThrMMA<Args...>     const& thr_mma,
+                             Tensor<TA, ALayout> const& sA,
+                             Tensor<TB, BLayout> const& sB,
+                             Tensor<TC, RCLayout>     & tCrC,
+                             ALoadTransformOp    const& sA_load_op,  // transforms A values before use in GEMM
+                             BLoadTransformOp    const& sB_load_op)  // transforms B values before use in GEMM
+{
+  using InputTypeA        = typename TA::value_type;
+  using InputTypeB        = typename TB::value_type;
+  using InputTypeC        = typename TC::value_type;
+  using ComputeTypeA = typename ThrMMA<Args...>::ValTypeA;
+  using ComputeTypeB = typename ThrMMA<Args...>::ValTypeB;
+  using ComputeTypeC = typename ThrMMA<Args...>::ValTypeC;
+
+  //
+  // MMA Partitioning
+  //
+
+  // Partition the sA, sB, and sC tiles across the threads for the MMA
+  Tensor tCsA = thr_mma.partition_A(sA);                            // (MMA,MMA_M,MMA_K)
+  Tensor tCsB = thr_mma.partition_B(sB);                            // (MMA,MMA_N,MMA_K)
+
+  // Create register tensors for the MMA to operate on
+  Tensor tCrA = thr_mma.make_fragment_A(tCsA);                      // (MMA,MMA_M,MMA_K)
+  Tensor tCrB = thr_mma.make_fragment_B(tCsB);                      // (MMA,MMA_N,MMA_K)
+
+#if 0
+  if (thread0()) {
+    print("  sA: "); print(  sA); print("\n");
+    print("  sB: "); print(  sB); print("\n");
+    print(thr_mma);
+    print("tCsA: "); print(tCsA); print("\n");
+    print("tCsB: "); print(tCsB); print("\n");
+    print("tCrA: "); print(tCrA); print("\n");
+    print("tCrB: "); print(tCrB); print("\n");
+    print("tCrC: "); print(tCrC); print("\n");
+  }
+#endif
+
+  //
+  // PREDICATION
+  //
+
+  // Create coordinate tensors for the problem
+  Tensor cA = make_identity_tensor(shape(sA));                      // (M,K) -> (m,k)
+  Tensor cB = make_identity_tensor(shape(sB));                      // (N,K) -> (n,k)
+
+  // Repeat partitioning with thr_mma
+  Tensor tCcA = thr_mma.partition_A(cA);                            // (MMA,MMA_M,MMA_K) -> (m,k)
+  Tensor tCcB = thr_mma.partition_B(cB);                            // (MMA,MMA_N,MMA_K) -> (n,k)
+
+  // Allocate the preds for MMA- and MMA_MN-modes
+  Tensor tCpA = make_tensor<bool>(make_shape(size<0>(tCsA), size<1>(tCsA)));
+  Tensor tCpB = make_tensor<bool>(make_shape(size<0>(tCsB), size<1>(tCsB)));
+
+  // Populate the predicates on M and N
+  MUTE_UNROLL
+  for (int i = 0; i < size(tCpA); ++i) {
+    tCpA(i) = elem_less(get<0>(tCcA(_,_,Int<0>{})(i)), shape<0>(sA));
+  }
+  MUTE_UNROLL
+  for (int i = 0; i < size(tCpB); ++i) {
+    tCpB(i) = elem_less(get<0>(tCcB(_,_,Int<0>{})(i)), shape<0>(sB));
+  }
+
+#if 0
+  if (thread0()) {
+    print("  cA: "); print(  cA); print("\n");
+    print("  cB: "); print(  cB); print("\n");
+    print("tCcA: "); print(tCcA); print("\n");
+    print("tCcB: "); print(tCcB); print("\n");
+    print_tensor(tCpA);
+    print_tensor(tCpB);
+  }
+#endif
+
+  //
+  // PREFETCH k_block = 0
+  //   Condition the k-predication on (static) k_block == K_BLOCK_MAX-1, the last k_block
+  //   Assumes the MMA-tiling in K is trivial
+  //
+
+  constexpr int K_BLOCK_MAX = size<2>(tCrA);
+
+  MUTE_UNROLL
+  for (int m = 0; m < size<1>(tCrA); ++m) {     // Copy MMA_M
+    MUTE_UNROLL
+    for (int i = 0; i < size<0>(tCrA); ++i) {   // Copy MMA_I
+      tCrA(i,m,0) = (tCpA(i,m) && (0 < K_BLOCK_MAX-1 || elem_less(get<1>(tCcA(i,m,0)), shape<1>(sA)))) ? static_cast<ComputeTypeA>(sA_load_op(tCsA(i,m,0))) : ComputeTypeA{};
+    }
+  }
+  MUTE_UNROLL
+  for (int n = 0; n < size<1>(tCrB); ++n) {     // Copy MMA_N
+    MUTE_UNROLL
+    for (int i = 0; i < size<0>(tCrB); ++i) {   // Copy MMA_I
+      tCrB(i,n,0) = (tCpB(i,n) && (0 < K_BLOCK_MAX-1 || elem_less(get<1>(tCcB(i,n,0)), shape<1>(sB)))) ? static_cast<ComputeTypeB>(sB_load_op(tCsB(i,n,0))) : ComputeTypeB{};
+    }
+  }
+  //
+  // MAINLOOP
+  //
+
+  MUTE_UNROLL
+  for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block)
+  {
+    if (k_block < K_BLOCK_MAX-1)   // static-if not the last k_block
+    {
+      int k_next = k_block + 1;    // Load k_next block
+
+      //   Condition the k-predication on (static) k_block == K_BLOCK_MAX-1, the last k_block
+      //   Assumes the MMA-tiling in K is trivial
+
+      MUTE_UNROLL
+      for (int m = 0; m < size<1>(tCrA); ++m) {       // Copy MMA_M
+        MUTE_UNROLL
+        for (int i = 0; i < size<0>(tCrA); ++i) {     // Copy MMA_I
+          tCrA(i,m,k_next) = (tCpA(i,m) && (k_next < K_BLOCK_MAX-1 || elem_less(get<1>(tCcA(i,m,k_next)), shape<1>(sA)))) ? static_cast<ComputeTypeA>(sA_load_op(tCsA(i,m,k_next))) : ComputeTypeA{};
+        }
+      }
+      MUTE_UNROLL
+      for (int n = 0; n < size<1>(tCrB); ++n) {       // Copy MMA_N
+        MUTE_UNROLL
+        for (int i = 0; i < size<0>(tCrB); ++i) {     // Copy MMA_I
+          tCrB(i,n,k_next) = (tCpB(i,n) && (k_next < K_BLOCK_MAX-1 || elem_less(get<1>(tCcB(i,n,k_next)), shape<1>(sB)))) ? static_cast<ComputeTypeB>(sB_load_op(tCsB(i,n,k_next))) : ComputeTypeB{};
+        }
+      }
+    }
+    // GEMM on k_block in registers
+    gemm(thr_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
+  }
+}
+
+// Unpredicated Cooperative GEMM
+template <class... Args,
+          class TA, class ALayout, class TB, class BLayout,
+          class TC, class CLayout,
+          class ALoadTransformOp, class BLoadTransformOp,
+          class SmemCopyOpA, class SmemCopyOpB>
+MUTE_HOST_DEVICE
+void
+cooperative_gemm_no_predication(uint32_t                   thread_idx,
+                                ThrMMA<Args...>     const& thr_mma,
+                                Tensor<TA, ALayout> const& sA,
+                                Tensor<TB, BLayout> const& sB,
+                                Tensor<TC, CLayout>      & tCrC,
+                                ALoadTransformOp    const& sA_load_op,  // transforms A values before use in GEMM
+                                BLoadTransformOp    const& sB_load_op,  // transforms B values before use in GEMM
+                                SmemCopyOpA         const& sA_copy_op,
+                                SmemCopyOpB         const& sB_copy_op)
+{
+  using InputTypeA   = typename TA::value_type;
+  using InputTypeB   = typename TB::value_type;
+  using InputTypeC   = typename TC::value_type;
+  using ComputeTypeA = typename ThrMMA<Args...>::ValTypeA;
+  using ComputeTypeB = typename ThrMMA<Args...>::ValTypeB;
+  using ComputeTypeC = typename ThrMMA<Args...>::ValTypeC;
+
+
+  //
+  // MMA Partitioning
+  //
+
+  // Create register tensors for the MMA to operate on
+  Tensor tCrA  = thr_mma.partition_fragment_A(sA);                    // (MMA,MMA_M,MMA_K)
+  Tensor tCrB  = thr_mma.partition_fragment_B(sB);                    // (MMA,MMA_N,MMA_K)
+
+  using CopyOpAType = SmemCopyOpA;
+  using CopyOpBType = SmemCopyOpB;
+
+  auto smem_tiled_copy_A = make_tiled_copy_A(Copy_Atom<CopyOpAType, InputTypeA>{}, thr_mma);
+  auto smem_thr_copy_A   = smem_tiled_copy_A.get_thread_slice(thread_idx);
+  Tensor tCsA            = smem_thr_copy_A.partition_S(sA);
+  Tensor tCrAi           = make_fragment_like(tCsA);
+  Tensor tCrAi_copy_view = smem_thr_copy_A.retile_D(tCrAi);
+  MUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrAi_copy_view));             // CPY_M
+  MUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCrAi_copy_view));             // CPY_K
+
+  auto smem_tiled_copy_B = make_tiled_copy_B(Copy_Atom<CopyOpBType, InputTypeB>{}, thr_mma);
+  auto smem_thr_copy_B   = smem_tiled_copy_B.get_thread_slice(thread_idx);
+  Tensor tCsB            = smem_thr_copy_B.partition_S(sB);
+  Tensor tCrBi           = make_fragment_like(tCsB);
+  Tensor tCrBi_copy_view = smem_thr_copy_B.retile_D(tCrBi);
+  MUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrBi_copy_view));            // CPY_N
+  MUTE_STATIC_ASSERT_V(size<2>(tCsB) == size<2>(tCrBi_copy_view));            // CPY_K
+
+#if 0
+  if (thread0()) {
+    print("  sA: "); print(sA); print("\n");
+    print("  sB: "); print(sB); print("\n");
+    print(thr_mma); print("\n");
+    print("tCrA: "); print(tCrA); print("\n");
+    print("tCrB: "); print(tCrB); print("\n");
+    print("tCrC: "); print(tCrC); print("\n");
+    print(smem_thr_copy_A); print("\n");
+    print("tCsA: "); print(tCsA); print("\n");
+    print("tCrA_copy_view: "); print(tCrA_copy_view); print("\n");
+    print(smem_thr_copy_B); print("\n");
+    print("tCsB: "); print(tCsB); print("\n");
+    print("tCrB_copy_view: "); print(tCrB_copy_view); print("\n");
+  }
+#endif
+
+  //
+  // PREFETCH
+  //
+
+  copy(smem_tiled_copy_A, tCsA(_,_,Int<0>{}), tCrAi_copy_view(_,_,Int<0>{}));
+  copy(smem_tiled_copy_B, tCsB(_,_,Int<0>{}), tCrBi_copy_view(_,_,Int<0>{}));
+  //
+  // MAINLOOP
+  //
+
+  constexpr int K_BLOCK_MAX = size<2>(tCrA);
+
+  MUTE_UNROLL
+  for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block)
+  {
+    // static-if load the next k_block. No k-predication required on these loads.
+    if (k_block < K_BLOCK_MAX-1)
+    {
+      // Load the next k_block
+      int k_next = k_block + 1;       // statically unrolled
+      copy(smem_tiled_copy_A, tCsA(_,_,k_next), tCrAi_copy_view(_,_,k_next));
+      copy(smem_tiled_copy_B, tCsB(_,_,k_next), tCrBi_copy_view(_,_,k_next));
+    }
+
+    // Transform A and B, relying on the compiler to remove in case of identity ops
+    mute::transform(tCrAi(_,_,k_block), tCrA(_,_,k_block), sA_load_op);
+    mute::transform(tCrBi(_,_,k_block), tCrB(_,_,k_block), sB_load_op);
+
+    // GEMM on k_block in registers
+    gemm(thr_mma, tCrA(_,_,k_block), tCrB(_,_,k_block), tCrC);
+  }
+}
+
+} // end namespace detail
+
+// C passed as a shared memory tensor
+// Epilogue included
+template <class... Args,
+          class Alpha, class TA, class ALayout, class TB, class BLayout,
+          class Beta,  class TC, class CLayout,
+          class ALoadTransformOp = mute::identity, class BLoadTransformOp  = mute::identity,
+          class CLoadTransformOp = mute::identity, class CStoreTransformOp = mute::identity,
+          class SmemCopyOpA = DefaultCopy, class SmemCopyOpB = DefaultCopy,
+          class SmemCopyOpC = DefaultCopy>
+MUTE_HOST_DEVICE
+void
+cooperative_gemm(uint32_t                   thread_idx,
+                 TiledMMA<Args...>   const& tiled_mma,
+                 Alpha               const& alpha,
+                 Tensor<TA, ALayout> const& sA,
+                 Tensor<TB, BLayout> const& sB,
+                 Beta                const& beta,
+                 Tensor<TC, CLayout>      & sC,
+                 ALoadTransformOp    const& sA_load_op  = {}, // transforms A values before use in GEMM
+                 BLoadTransformOp    const& sB_load_op  = {}, // transforms B values before use in GEMM
+                 CLoadTransformOp    const& sC_load_op  = {}, // transforms C values before use in GEMM
+                 CStoreTransformOp   const& sC_store_op = {}, // transforms results before they are stored to C
+                 SmemCopyOpA         const& sA_copy_op  = {},
+                 SmemCopyOpB         const& sB_copy_op  = {},
+                 SmemCopyOpC         const& sC_copy_op  = {})
+{
+  MUTE_STATIC_ASSERT_V(rank(sA) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(sB) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(sC) == Int<2>{});
+
+  MUTE_STATIC_ASSERT_V(size<0>(sA) == size<0>(sC));  // AM == CM
+  MUTE_STATIC_ASSERT_V(size<0>(sB) == size<1>(sC));  // BN == CN
+  MUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB));  // AK == BK
+
+  using InputTypeA        = typename TA::value_type;
+  using InputTypeB        = typename TB::value_type;
+  using InputTypeC        = typename TC::value_type;
+  using ComputeTypeA = typename TiledMMA<Args...>::ValTypeA;
+  using ComputeTypeB = typename TiledMMA<Args...>::ValTypeB;
+  using ComputeTypeC = typename TiledMMA<Args...>::ValTypeC;
+
+  auto compat = evenly_divides(make_shape(size<0>(sA), size<0>(sB), size<1>(sA)),
+                               tile_shape(TiledMMA<Args...>{}));
+
+  // ThrMMA
+  auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
+  Tensor tCsC  = thr_mma.partition_C(sC);                             // (MMA,MMA_M,MMA_N) :: InputTypeC
+  Tensor tCrC  = thr_mma.make_fragment_C(tCsC);                       // (MMA,MMA_M,MMA_N) :: ComputeTypeC
+
+  // Clear accumulators
+  clear(tCrC);
+
+#if 0
+  if (thread0()) {
+    print("  sC: "); print(sC); print("\n");
+    print("  tCsC: "); print(tCsC); print("\n");
+  }
+#endif
+
+  if constexpr (is_constant<true, decltype(compat)>::value) {
+    detail::cooperative_gemm_no_predication(
+        thread_idx, thr_mma, sA, sB, tCrC, sA_load_op, sB_load_op, sA_copy_op, sB_copy_op
+    );
+    detail::epilogue_no_predication(
+        alpha, tCrC, beta, tCsC, sC_load_op, sC_store_op, sC_copy_op
+    );
+  } else {
+    detail::cooperative_gemm_predication(
+        thr_mma, sA, sB, tCrC, sA_load_op, sB_load_op
+    );
+    detail::epilogue_predication(
+        thr_mma, alpha, tCrC, beta, sC, tCsC, sC_load_op, sC_store_op
+    );
+  }
+}
+
+// C already partitioned into registers on input
+// It can be passed non-empty
+// Epilogue not included
+template <class... Args,
+          class TA, class ALayout, class TB, class BLayout,
+          class TC, class CLayout,
+          class ALoadTransformOp = mute::identity, class BLoadTransformOp  = mute::identity,
+          class SmemCopyOpA = DefaultCopy, class SmemCopyOpB = DefaultCopy>
+MUTE_HOST_DEVICE
+void
+cooperative_gemm(uint32_t                   thread_idx,
+                 TiledMMA<Args...>   const& tiled_mma,
+                 Tensor<TA, ALayout> const& sA,
+                 Tensor<TB, BLayout> const& sB,
+                 Tensor<TC, CLayout>      & tCrC,
+                 ALoadTransformOp    const& sA_load_op  = {}, // transforms A values before use in GEMM
+                 BLoadTransformOp    const& sB_load_op  = {}, // transforms B values before use in GEMM
+                 SmemCopyOpA         const& sA_copy_op  = {},
+                 SmemCopyOpB         const& sB_copy_op  = {})
+{
+  MUTE_STATIC_ASSERT_V(rank(sA) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(sB) == Int<2>{});
+
+  MUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB));  // AK == BK
+
+  using InputTypeA        = typename TA::value_type;
+  using InputTypeB        = typename TB::value_type;
+  using InputTypeC        = typename TC::value_type;
+  using ComputeTypeA = typename TiledMMA<Args...>::ValTypeA;
+  using ComputeTypeB = typename TiledMMA<Args...>::ValTypeB;
+  using ComputeTypeC = typename TiledMMA<Args...>::ValTypeC;
+
+  // Check if input C fragment is compatible with thr_mma and problem size
+  using ref_c_frag = decltype(partition_shape_C(tiled_mma, make_shape(size<0>(sA), size<0>(sB))));
+  MUTE_STATIC_ASSERT_V(compatible(shape(ref_c_frag{}), shape(tCrC)));
+
+  auto compat = evenly_divides(make_shape(size<0>(sA), size<0>(sB), size<1>(sA)),
+                               tile_shape(TiledMMA<Args...>{}));
+
+  // ThrMMA
+  auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
+
+  if constexpr (is_constant<true, decltype(compat)>::value) {
+    detail::cooperative_gemm_no_predication(
+        thread_idx, thr_mma, sA, sB, tCrC, sA_load_op, sB_load_op, sA_copy_op, sB_copy_op
+    );
+  } else {
+    detail::cooperative_gemm_predication(
+        thr_mma, sA, sB, tCrC, sA_load_op, sB_load_op
+    );
+  }
+}
+
+// Accept mutable temporaries
+template <class... Args,
+          class Alpha, class TA, class ALayout, class TB, class BLayout,
+          class Beta,  class TC, class CLayout,
+          class ALoadTransformOp = mute::identity, class BLoadTransformOp  = mute::identity,
+          class CLoadTransformOp = mute::identity, class CStoreTransformOp = mute::identity,
+          class SmemCopyOpA = DefaultCopy, class SmemCopyOpB = DefaultCopy,
+          class SmemCopyOpC = DefaultCopy>
+MUTE_HOST_DEVICE
+void
+cooperative_gemm(uint32_t thread_idx,
+                 TiledMMA<Args...>   const& tiled_mma,
+                 Alpha               const& alpha,
+                 Tensor<TA, ALayout> const& sA,
+                 Tensor<TB, BLayout> const& sB,
+                 Beta                const& beta,
+                 Tensor<TC, CLayout>     && sC,
+                 ALoadTransformOp    const& sA_load_op  = {}, // transforms A values before use in GEMM
+                 BLoadTransformOp    const& sB_load_op  = {}, // transforms B values before use in GEMM
+                 CLoadTransformOp    const& sC_load_op  = {}, // transforms C values before use in GEMM
+                 CStoreTransformOp   const& sC_store_op = {}, // transforms results before they are stored to C
+                 SmemCopyOpA         const& sA_copy_op  = {},
+                 SmemCopyOpB         const& sB_copy_op  = {},
+                 SmemCopyOpC         const& sC_copy_op  = {})
+{
+  cooperative_gemm(thread_idx, tiled_mma, alpha, sA, sB, beta, sC,
+                   sA_load_op, sB_load_op, sC_load_op, sC_store_op,
+                   sA_copy_op, sB_copy_op, sC_copy_op);
+}
+
+// Legacy overload of mute::gemm for backwards-compatibility
+template <class... Args,
+          class Alpha, class TA, class ALayout, class TB, class BLayout,
+          class Beta,  class TC, class CLayout,
+          class ALoadTransformOp = mute::identity, class BLoadTransformOp  = mute::identity,
+          class CLoadTransformOp = mute::identity, class CStoreTransformOp = mute::identity>
+MUTE_HOST_DEVICE
+void
+gemm(ThrMMA<Args...>     const& thr_mma,
+     Alpha               const& alpha,
+     Tensor<TA, ALayout> const& sA,
+     Tensor<TB, BLayout> const& sB,
+     Beta                const& beta,
+     Tensor<TC, CLayout>      & sC,
+     ALoadTransformOp    const& sA_load_op  = {}, // transforms A values before use in GEMM
+     BLoadTransformOp    const& sB_load_op  = {}, // transforms B values before use in GEMM
+     CLoadTransformOp    const& sC_load_op  = {}, // transforms C values before use in GEMM
+     CStoreTransformOp   const& sC_store_op = {}) // transforms results before they are stored to C
+{
+  MUTE_STATIC_ASSERT_V(rank(sA) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(sB) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(sC) == Int<2>{});
+
+  MUTE_STATIC_ASSERT_V(size<0>(sA) == size<0>(sC));  // AM == CM
+  MUTE_STATIC_ASSERT_V(size<0>(sB) == size<1>(sC));  // BN == CN
+  MUTE_STATIC_ASSERT_V(size<1>(sA) == size<1>(sB));  // AK == BK
+
+  Tensor tCsC  = thr_mma.partition_C(sC);                           // (MMA,MMA_M,MMA_N)
+  Tensor tCrC  = thr_mma.make_fragment_C(tCsC);                     // (MMA,MMA_M,MMA_N)
+
+  // Goes directly to the slow path to avoid getting thread_idx from thr_mma
+  detail::cooperative_gemm_predication(
+    thr_mma, sA, sB, sC, sA_load_op, sB_load_op
+  );
+
+  detail::epilogue_predication(
+      thr_mma, alpha, tCrC, beta, sC, tCsC, sC_load_op, sC_store_op
+  );
+}
+
+} // end namespace mute
\ No newline at end of file
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/copy.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/copy.hpp
new file mode 100644
index 000000000..55d782625
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/copy.hpp
@@ -0,0 +1,345 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/container/alignment.hpp>
+
+#include <mute/tensor.hpp>
+#include <mute/tensor_predicate.hpp>
+
+#include <mute/atom/copy_atom.hpp>
+
+namespace mute
+{
+
+//
+// Accept mutable temporaries
+//
+
+template <class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy(Tensor<SrcEngine, SrcLayout> const& src,
+     Tensor<DstEngine, DstLayout>     && dst)
+{
+  return copy(src, dst);
+}
+
+template <class VecType,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
+         Tensor<DstEngine, DstLayout>     && dst)
+{
+  return copy_vec<VecType>(src, dst);
+}
+
+template <class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
+             Tensor<DstEngine, DstLayout>     && dst)
+{
+  return copy_aligned(src, dst);
+}
+
+template <class PrdTensor,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy_if(PrdTensor                    const& pred,
+        Tensor<SrcEngine, SrcLayout> const& src,
+        Tensor<DstEngine, DstLayout>     && dst)
+{
+  return copy_if(pred, src, dst);
+}
+
+template <class CopyPolicy,
+          class PrdTensor,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy_if(CopyPolicy                   const& copy_policy,
+        PrdTensor                    const& pred,
+        Tensor<SrcEngine, SrcLayout> const& src,
+        Tensor<DstEngine, DstLayout>     && dst)
+{
+  return copy_if(copy_policy, pred, src, dst);
+}
+
+template <class CopyPolicy,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy(CopyPolicy                   const& copy_policy,
+     Tensor<SrcEngine, SrcLayout> const& src,
+     Tensor<DstEngine, DstLayout>     && dst)
+{
+  return copy(copy_policy, src, dst);
+}
+
+//
+// copy_if -- Predicated Copy
+//
+
+template <class PrdTensor,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy_if(PrdTensor                    const& pred,
+        Tensor<SrcEngine, SrcLayout> const& src,
+        Tensor<DstEngine, DstLayout>      & dst)
+{
+  auto copy_op = select_elementwise_copy(src, dst);
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(src); ++i) {
+    if (pred(i)) {
+      copy_op.copy(src(i), dst(i));
+    }
+  }
+}
+
+//
+// copy_if -- Predicated CopyAtom
+//
+
+namespace detail {
+
+// Trait that detects if atom's traits has a member function with(bool)
+template <class, class Enable = void>
+constexpr bool has_with_bool = false;
+
+template <class T>
+constexpr bool has_with_bool<T, mute::void_t<decltype(declval<typename T::Traits>().with(declval<bool>()))>> = true;
+
+} // end namespace detail
+
+template <class... CopyArgs,
+          class PredTensor,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy_if(Copy_Atom<CopyArgs...>       const& copy_atom,
+        PredTensor                   const& pred,      // (Rest...)
+        Tensor<SrcEngine, SrcLayout> const& src,       // (V,Rest...)
+        Tensor<DstEngine, DstLayout>      & dst)       // (V,Rest...)
+{
+  static_assert(SrcLayout::rank == DstLayout::rank, "CopyAtom rank-mismatch.");
+  if constexpr (SrcLayout::rank == 1) {   // Dispatch the copy
+    copy_atom.call(src, dst);
+
+    // Unable to compile moe kernel
+    // if constexpr (detail::has_with_bool<Copy_Atom<CopyArgs...>>) {
+	  //   copy_atom.with(pred(0)).call(src,dst);
+    // } else {
+    //   if (pred(0)) copy_atom.call(src, dst);
+    // }
+  } else {                                // Loop over all but the first mode
+    constexpr int R = SrcLayout::rank;
+    Tensor src_v = group_modes<1,R>(src);
+    Tensor dst_v = group_modes<1,R>(dst);
+    MUTE_UNROLL
+    for (int i = 0; i < size<1>(src_v); ++i) {
+      // If copy traits can be transformed with a predicate value, do it, otherwise branch here
+      if constexpr (detail::has_with_bool<Copy_Atom<CopyArgs...>>) {
+        copy_atom.with(pred(i)).call(src_v(_,i), dst_v(_,i));
+      } else {
+        if (pred(i)) {
+          copy_atom.call(src_v(_,i), dst_v(_,i));
+        }
+      }
+    }
+  }
+}
+
+//
+// copy_vec -- attempt vectorized copy with VecType
+//
+
+template <class VecType,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy_vec(Tensor<SrcEngine, SrcLayout> const& src,
+         Tensor<DstEngine, DstLayout>      & dst)
+{
+  static_assert(sizeof_bits_v<VecType> >= 8 && sizeof_bits_v<VecType> % 8 == 0,
+                "Expected a vectorization type of at least a byte.");
+  using SrcType = typename SrcEngine::element_type;
+  using DstType = typename DstEngine::element_type;
+  if constexpr (sizeof_bits_v<SrcType> == sizeof_bits_v<DstType> &&
+                sizeof_bits_v<VecType>  > sizeof_bits_v<DstType>)
+  {
+    // Preserve volatility of Src/Dst types.
+    using SrcVecType = conditional_t<is_volatile_v<SrcType>, VecType const volatile, VecType const>;
+    using DstVecType = conditional_t<is_volatile_v<DstType>, VecType       volatile, VecType      >;
+    Tensor src_v = recast<SrcVecType>(src);
+    Tensor dst_v = recast<DstVecType>(dst);
+
+#if 0
+    if (thread0()) {
+      print("copy_vec<%db> -- vectorizing copy:\n", int(sizeof_bits_v<VecType>));
+      print("   "); print(src); print(" => "); print(src_v); print("\n");
+      print("   "); print(dst); print(" => "); print(dst_v); print("\n");
+    }
+#endif
+
+    return copy_if(TrivialPredTensor{}, src_v, dst_v);
+  } else {
+#if 0
+  if (thread0()) {
+    print("copy_vec<%db> -- NOT vectorizing copy:\n", int(sizeof_bits_v<VecType>));
+    print("   "); print(src); print("\n");
+    print("   "); print(dst); print("\n");
+  }
+#endif
+
+    return copy_if(TrivialPredTensor{}, src, dst);
+  }
+}
+
+//
+// copy -- CopyAtom
+//
+
+template <class... CopyArgs,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy(Copy_Atom<CopyArgs...>       const& copy_atom,
+     Tensor<SrcEngine, SrcLayout> const& src,
+     Tensor<DstEngine, DstLayout>      & dst)
+{
+  return copy_if(copy_atom, TrivialPredTensor{}, src, dst);
+}
+
+//////////////////////////////////////////
+// Special Auto-Vectorizing Overloads
+//////////////////////////////////////////
+
+// Specialization for AutoVectorizingCopyAssumedAlignment<MaxVecBits>
+template <int MaxVecBits, class... Args,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits> const&,
+     Tensor<SrcEngine, SrcLayout>                        const& src,
+     Tensor<DstEngine, DstLayout>                             & dst)
+{
+  constexpr int vec_elem = decltype(max_common_vector(src, dst))::value;
+
+  constexpr int src_bits = sizeof_bits<typename SrcEngine::value_type>::value;
+  // When layouts are static,  accept vec_bits up to 128
+  // When layouts are dynamic, accept vec_bits up to MaxVecBits
+  constexpr int vec_bits = (is_static<SrcLayout>::value && is_static<DstLayout>::value) ?
+                            mute::min(vec_elem * src_bits, 128) :
+                            mute::min(vec_elem * src_bits, MaxVecBits);
+
+#if 0
+  if (thread0()) {
+    print("copy -- found max_common_vector of %d elems and vectorization to %d bits\n", vec_elem, vec_bits);
+    print("   "); print(src); print("\n");
+    print("   "); print(dst); print("\n");
+  }
+#endif
+
+  if constexpr (vec_elem > 1 && vec_bits >= 8) {
+    return copy_vec<uint_bit_t<vec_bits>>(src, dst);
+  } else {
+    return copy_if(TrivialPredTensor{}, src, dst);
+  }
+}
+
+// Auto-vectorizing copy for static layouts
+template <class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy(Tensor<SrcEngine, SrcLayout> const& src,
+     Tensor<DstEngine, DstLayout>      & dst)
+{
+  return copy(AutoVectorizingCopy{}, src, dst);
+}
+
+// Auto-vectorizing copy with assumed alignment of dynamic layout strides up to 128bit.
+template <class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy_aligned(Tensor<SrcEngine, SrcLayout> const& src,
+             Tensor<DstEngine, DstLayout>      & dst)
+{
+  return copy(AutoVectorizingCopyWithAssumedAlignment<128>{}, src, dst);
+}
+
+// Specializaton for Atom AutoVectorizingCopy
+template <class... Args,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy(Copy_Atom<AutoVectorizingCopy, Args...> const&,
+     Tensor<SrcEngine, SrcLayout>            const& src,
+     Tensor<DstEngine, DstLayout>                 & dst)
+{
+  return copy(AutoVectorizingCopy{}, src, dst);
+}
+
+// Specializaton for Atom AutoVectorizingCopyAssumedAlignment
+template <int MaxVecBits, class... Args,
+          class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE
+void
+copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>, Args...> const&,
+     Tensor<SrcEngine, SrcLayout>                                            const& src,
+     Tensor<DstEngine, DstLayout>                                                 & dst)
+{
+  return copy(AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>{}, src, dst);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/fill.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/fill.hpp
new file mode 100644
index 000000000..1ed2192bf
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/fill.hpp
@@ -0,0 +1,88 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/tensor.hpp>
+#include <mute/algorithm/prefer.hpp>
+
+namespace mute
+{
+
+//
+// Accept mutable temporaries
+//
+template <class Engine, class Layout, class T>
+MUTE_HOST_DEVICE
+void
+fill(Tensor<Engine, Layout>&& tensor, T const& value)
+{
+  return fill(tensor, value);
+}
+
+namespace detail
+{
+
+// Prefer fill(tensor.data(), value), if possible
+template <class Engine, class Layout, class T>
+MUTE_HOST_DEVICE
+auto
+fill(Tensor<Engine, Layout>& tensor, T const& value, prefer<1>)
+    -> decltype(fill(tensor.data(), value))
+{
+  fill(tensor.data(), value);
+}
+
+// Default implementation
+template <class Engine, class Layout, class T>
+MUTE_HOST_DEVICE
+void
+fill(Tensor<Engine, Layout>& tensor, T const& value, prefer<0>)
+{
+  MUTE_UNROLL
+  for (int i = 0; i < size(tensor); ++i) {
+    tensor(i) = value;
+  }
+}
+
+} // end namespace detail
+
+template <class Engine, class Layout, class T>
+MUTE_HOST_DEVICE
+void
+fill(Tensor<Engine, Layout>& tensor, T const& value)
+{
+  return detail::fill(tensor, value, prefer<1>{});
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/functional.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/functional.hpp
new file mode 100644
index 000000000..1ae04047a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/functional.hpp
@@ -0,0 +1,291 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>          // MUTE_HOST_DEVICE
+#include <mute/numeric/math.hpp>    // mute::max, mute::min
+#include <mute/numeric/complex.hpp> // mute::conj
+
+/** C++14 <functional> extensions */
+
+namespace mute {
+
+/**************/
+/** Identity **/
+/**************/
+
+struct identity {
+  template <class T>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto) operator()(T&& arg) const {
+    return static_cast<T&&>(arg);
+  }
+};
+
+template <class R>
+struct constant_fn {
+  template <class... T>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto) operator()(T&&...) const {
+    return r_;
+  }
+  R r_;
+};
+
+/***********/
+/** Unary **/
+/***********/
+
+#define MUTE_LEFT_UNARY_OP(NAME,OP)                                  \
+  struct NAME {                                                      \
+    template <class T>                                               \
+    MUTE_HOST_DEVICE constexpr                                       \
+    decltype(auto) operator()(T&& arg) const {                       \
+      return OP static_cast<T&&>(arg);                                \
+    }                                                                \
+  }
+#define MUTE_RIGHT_UNARY_OP(NAME,OP)                                 \
+  struct NAME {                                                      \
+    template <class T>                                               \
+    MUTE_HOST_DEVICE constexpr                                       \
+    decltype(auto) operator()(T&& arg) const {                       \
+      return static_cast<T&&>(arg) OP ;                               \
+    }                                                                \
+  }
+#define MUTE_NAMED_UNARY_OP(NAME,OP)                                 \
+  struct NAME {                                                      \
+    template <class T>                                               \
+    MUTE_HOST_DEVICE constexpr                                       \
+    decltype(auto) operator()(T&& arg) const {                       \
+      return OP (static_cast<T&&>(arg));                              \
+    }                                                                \
+  }
+
+MUTE_LEFT_UNARY_OP(unary_plus,       +);
+MUTE_LEFT_UNARY_OP(negate,           -);
+MUTE_LEFT_UNARY_OP(bit_not,          ~);
+MUTE_LEFT_UNARY_OP(logical_not,      !);
+MUTE_LEFT_UNARY_OP(dereference,      *);
+MUTE_LEFT_UNARY_OP(address_of,       &);
+MUTE_LEFT_UNARY_OP(pre_increment,   ++);
+MUTE_LEFT_UNARY_OP(pre_decrement,   --);
+
+MUTE_RIGHT_UNARY_OP(post_increment, ++);
+MUTE_RIGHT_UNARY_OP(post_decrement, --);
+
+MUTE_NAMED_UNARY_OP(abs_fn,           abs);
+MUTE_NAMED_UNARY_OP(conjugate, mute::conj);
+
+#undef MUTE_LEFT_UNARY_OP
+#undef MUTE_RIGHT_UNARY_OP
+#undef MUTE_NAMED_UNARY_OP
+
+template <int Shift_>
+struct shift_right_const {
+  static constexpr int Shift = Shift_;
+
+  template <class T>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto) operator()(T&& arg) const {
+    return static_cast<T&&>(arg) >> Shift;
+  }
+};
+
+template <int Shift_>
+struct shift_left_const {
+  static constexpr int Shift = Shift_;
+
+  template <class T>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto) operator()(T&& arg) const {
+    return static_cast<T&&>(arg) << Shift;
+  }
+};
+
+/************/
+/** Binary **/
+/************/
+
+#define MUTE_BINARY_OP(NAME,OP)                                      \
+  struct NAME {                                                      \
+    template <class T, class U>                                      \
+    MUTE_HOST_DEVICE constexpr                                       \
+    decltype(auto) operator()(T&& lhs, U&& rhs) const {              \
+      return static_cast<T&&>(lhs) OP static_cast<U&&>(rhs);           \
+    }                                                                \
+  }
+#define MUTE_NAMED_BINARY_OP(NAME,OP)                                \
+  struct NAME {                                                      \
+    template <class T, class U>                                      \
+    MUTE_HOST_DEVICE constexpr                                       \
+    decltype(auto) operator()(T&& lhs, U&& rhs) const {              \
+      return OP (static_cast<T&&>(lhs), static_cast<U&&>(rhs));        \
+    }                                                                \
+  }
+
+
+MUTE_BINARY_OP(plus,                 +);
+MUTE_BINARY_OP(minus,                -);
+MUTE_BINARY_OP(multiplies,           *);
+MUTE_BINARY_OP(divides,              /);
+MUTE_BINARY_OP(modulus,              %);
+
+MUTE_BINARY_OP(plus_assign,         +=);
+MUTE_BINARY_OP(minus_assign,        -=);
+MUTE_BINARY_OP(multiplies_assign,   *=);
+MUTE_BINARY_OP(divides_assign,      /=);
+MUTE_BINARY_OP(modulus_assign,      %=);
+
+MUTE_BINARY_OP(bit_and,              &);
+MUTE_BINARY_OP(bit_or,               |);
+MUTE_BINARY_OP(bit_xor,              ^);
+MUTE_BINARY_OP(left_shift,          <<);
+MUTE_BINARY_OP(right_shift,         >>);
+
+MUTE_BINARY_OP(bit_and_assign,      &=);
+MUTE_BINARY_OP(bit_or_assign,       |=);
+MUTE_BINARY_OP(bit_xor_assign,      ^=);
+MUTE_BINARY_OP(left_shift_assign,  <<=);
+MUTE_BINARY_OP(right_shift_assign, >>=);
+
+MUTE_BINARY_OP(logical_and,         &&);
+MUTE_BINARY_OP(logical_or,          ||);
+
+MUTE_BINARY_OP(equal_to,            ==);
+MUTE_BINARY_OP(not_equal_to,        !=);
+MUTE_BINARY_OP(greater,              >);
+MUTE_BINARY_OP(less,                 <);
+MUTE_BINARY_OP(greater_equal,       >=);
+MUTE_BINARY_OP(less_equal,          <=);
+
+MUTE_NAMED_BINARY_OP(max_fn, mute::max);
+MUTE_NAMED_BINARY_OP(min_fn, mute::min);
+
+#undef MUTE_BINARY_OP
+#undef MUTE_NAMED_BINARY_OP
+
+/**********/
+/** Fold **/
+/**********/
+
+#define MUTE_FOLD_OP(NAME,OP)                                        \
+  struct NAME##_unary_rfold {                                        \
+    template <class... T>                                            \
+    MUTE_HOST_DEVICE constexpr                                       \
+    auto operator()(T&&... t) const {                                \
+      return (t OP ...);                                             \
+    }                                                                \
+  };                                                                 \
+  struct NAME##_unary_lfold {                                        \
+    template <class... T>                                            \
+    MUTE_HOST_DEVICE constexpr                                       \
+    auto operator()(T&&... t) const {                                \
+      return (... OP t);                                             \
+    }                                                                \
+  };                                                                 \
+  struct NAME##_binary_rfold {                                       \
+    template <class U, class... T>                                   \
+    MUTE_HOST_DEVICE constexpr                                       \
+    auto operator()(U&& u, T&&... t) const {                         \
+      return (t OP ... OP u);                                        \
+    }                                                                \
+  };                                                                 \
+  struct NAME##_binary_lfold {                                       \
+    template <class U, class... T>                                   \
+    MUTE_HOST_DEVICE constexpr                                       \
+    auto operator()(U&& u, T&&... t) const {                         \
+      return (u OP ... OP t);                                        \
+    }                                                                \
+  }
+
+MUTE_FOLD_OP(plus,                 +);
+MUTE_FOLD_OP(minus,                -);
+MUTE_FOLD_OP(multiplies,           *);
+MUTE_FOLD_OP(divides,              /);
+MUTE_FOLD_OP(modulus,              %);
+
+MUTE_FOLD_OP(plus_assign,         +=);
+MUTE_FOLD_OP(minus_assign,        -=);
+MUTE_FOLD_OP(multiplies_assign,   *=);
+MUTE_FOLD_OP(divides_assign,      /=);
+MUTE_FOLD_OP(modulus_assign,      %=);
+
+MUTE_FOLD_OP(bit_and,              &);
+MUTE_FOLD_OP(bit_or,               |);
+MUTE_FOLD_OP(bit_xor,              ^);
+MUTE_FOLD_OP(left_shift,          <<);
+MUTE_FOLD_OP(right_shift,         >>);
+
+MUTE_FOLD_OP(bit_and_assign,      &=);
+MUTE_FOLD_OP(bit_or_assign,       |=);
+MUTE_FOLD_OP(bit_xor_assign,      ^=);
+MUTE_FOLD_OP(left_shift_assign,  <<=);
+MUTE_FOLD_OP(right_shift_assign, >>=);
+
+MUTE_FOLD_OP(logical_and,         &&);
+MUTE_FOLD_OP(logical_or,          ||);
+
+MUTE_FOLD_OP(equal_to,            ==);
+MUTE_FOLD_OP(not_equal_to,        !=);
+MUTE_FOLD_OP(greater,              >);
+MUTE_FOLD_OP(less,                 <);
+MUTE_FOLD_OP(greater_equal,       >=);
+MUTE_FOLD_OP(less_equal,          <=);
+
+#undef MUTE_FOLD_OP
+
+/**********/
+/** Meta **/
+/**********/
+
+template <class Fn, class Arg>
+struct bound_fn {
+
+  template <class T>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  operator()(T&& arg) {
+    return fn_(arg_, static_cast<T&&>(arg));
+  }
+
+  Fn fn_;
+  Arg arg_;
+};
+
+template <class Fn, class Arg>
+MUTE_HOST_DEVICE constexpr
+auto
+bind(Fn const& fn, Arg const& arg) {
+  return bound_fn<Fn,Arg>{fn, arg};
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/gemm.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/gemm.hpp
new file mode 100644
index 000000000..00528ed9c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/gemm.hpp
@@ -0,0 +1,501 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/util/type_traits.hpp>
+#include <mute/algorithm/functional.hpp>
+
+#include <mute/tensor.hpp>
+
+#include <mute/atom/mma_atom.hpp>
+
+/** The gemm algorithm takes four (or three) tensors and computes
+ *   D = A * B + C
+ * It dispatches based on the number of modes each tensor has:
+ *
+ * 1. `(V) x (V) => (V)`.
+ *      The element-wise product of vectors. Dispatches to FMA or MMA.
+ * 2. `(M) x (N) => (M,N)`.
+ *      The outer product of vectors. Dispatches to [3] with new mode K=(1).
+ * 3. `(M,K) x (N,K) => (M,N)`.
+ *      The product of matrices. Dispatches to [5] with MMA vector-mode V.
+ * 4. `(V,M) x (V,N) => (V,M,N)`.
+ *      The batched outer product of vectors. Accounts for register reuse and dispatches to [1] for each (m,n).
+ * 5. `(V,M,K) x (V,N,K) => (V,M,N)`.
+ *      The batched product of matrices. Dispatches to [4] for each (k).
+ */
+
+namespace mute
+{
+
+//
+// Three arguments to four
+//
+
+template <class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE
+void
+gemm(Tensor<TA, ALayout> const& A,
+     Tensor<TB, BLayout> const& B,
+     Tensor<TC, CLayout>      & C)
+{
+  return gemm(C, A, B, C);
+}
+
+template <class MMA,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TA, ALayout> const& A,
+     Tensor<TB, BLayout> const& B,
+     Tensor<TC, CLayout>      & C)
+{
+  return gemm(mma, C, A, B, C);
+}
+
+//
+// Accept mutable temporaries
+//
+
+template <class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE
+void
+gemm(Tensor<TA, ALayout> const& A,
+     Tensor<TB, BLayout> const& B,
+     Tensor<TC, CLayout>     && C)
+{
+  return gemm(C, A, B, C);
+}
+
+template <class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE
+void
+gemm(Tensor<TD, DLayout>     && D,
+     Tensor<TA, ALayout> const& A,
+     Tensor<TB, BLayout> const& B,
+     Tensor<TC, CLayout> const& C)
+{
+  return gemm(D, A, B, C);
+}
+
+template <class MMA,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TA, ALayout> const& A,
+     Tensor<TB, BLayout> const& B,
+     Tensor<TC, CLayout>     && C)
+{
+  return gemm(mma, C, A, B, C);
+}
+
+template <class MMA,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TD, DLayout>     && D,
+     Tensor<TA, ALayout> const& A,
+     Tensor<TB, BLayout> const& B,
+     Tensor<TC, CLayout> const& C)
+{
+  return gemm(mma, D, A, B, C);
+}
+
+//
+// Default MMA is UniversalFMA
+//
+
+template <class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE
+void
+gemm(Tensor<TD, DLayout>      & D,
+     Tensor<TA, ALayout> const& A,
+     Tensor<TB, BLayout> const& B,
+     Tensor<TC, CLayout> const& C)
+{
+  using MMA = MMA_Atom<UniversalFMA<typename Tensor<TD,DLayout>::value_type,
+                                    typename Tensor<TA,ALayout>::value_type,
+                                    typename Tensor<TB,BLayout>::value_type,
+                                    typename Tensor<TC,CLayout>::value_type>>;
+
+  return gemm(MMA{}, D, A, B, C);
+}
+
+//
+// Thread-Local Register-Memory GEMMs
+//
+
+// Dispatch [1]: (V) x (V) => (V)
+template <class MMA,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout,
+          __MUTE_REQUIRES(DLayout::rank == 1 && is_rmem<TD>::value &&
+                          ALayout::rank == 1 && is_rmem<TA>::value &&
+                          BLayout::rank == 1 && is_rmem<TB>::value &&
+                          CLayout::rank == 1 && is_rmem<TC>::value)>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TD, DLayout>      & D,  // (V) Logical data
+     Tensor<TA, ALayout> const& A,  // (V) Logical data
+     Tensor<TB, BLayout> const& B,  // (V) Logical data
+     Tensor<TC, CLayout> const& C)  // (V) Logical data
+{
+  // No static assertions on (V), MMA checks compatibility
+  mma.call(D, A, B, C);
+}
+
+// Dispatch [2]: (M) x (N) => (M,N)
+template <class MMA,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout,
+          __MUTE_REQUIRES(DLayout::rank == 2 && is_rmem<TD>::value &&
+                          ALayout::rank == 1 && is_rmem<TA>::value &&
+                          BLayout::rank == 1 && is_rmem<TB>::value &&
+                          CLayout::rank == 2 && is_rmem<TC>::value)>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TD, DLayout>      & D,  // (M,N) Logical data
+     Tensor<TA, ALayout> const& A,  // (M)   Logical data
+     Tensor<TB, BLayout> const& B,  // (N)   Logical data
+     Tensor<TC, CLayout> const& C)  // (M,N) Logical data
+{
+  MUTE_STATIC_ASSERT_V(size<0>(A) == size<0>(C));  // AM == CM
+  MUTE_STATIC_ASSERT_V(size<0>(B) == size<1>(C));  // BN == CN
+  MUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D));
+  gemm(mma,
+       D,                                                       // (M,N)
+       make_tensor(A.data(), append<2>(A.layout())),            // (M,1)
+       make_tensor(B.data(), append<2>(B.layout())),            // (N,1)
+       C);                                                      // (M,N)
+}
+
+// Dispatch [3]: (M,K) x (N,K) => (M,N)
+template <class MMA,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout,
+          __MUTE_REQUIRES(DLayout::rank == 2 && is_rmem<TD>::value &&
+                          ALayout::rank == 2 && is_rmem<TA>::value &&
+                          BLayout::rank == 2 && is_rmem<TB>::value &&
+                          CLayout::rank == 2 && is_rmem<TC>::value)>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TD, DLayout>      & D,  // (M,N) Logical data
+     Tensor<TA, ALayout> const& A,  // (M,K) Logical data
+     Tensor<TB, BLayout> const& B,  // (N,K) Logical data
+     Tensor<TC, CLayout> const& C)  // (M,N) Logical data
+{
+  MUTE_STATIC_ASSERT_V(size<0>(A) == size<0>(C));  // AM == CM
+  MUTE_STATIC_ASSERT_V(size<0>(B) == size<1>(C));  // BN == CN
+  MUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(B));  // AK == BK
+  MUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D));
+
+  // Assert this is a 1-value MMA
+  MUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutC_TV{}) == Int<1>{});
+  MUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutA_TV{}) == Int<1>{});
+  MUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutB_TV{}) == Int<1>{});
+
+  gemm(mma,
+       make_tensor(D.data(), prepend<3>(D.layout())),      // (1,M,N)
+       make_tensor(A.data(), prepend<3>(A.layout())),      // (1,M,K)
+       make_tensor(B.data(), prepend<3>(B.layout())),      // (1,N,K)
+       make_tensor(C.data(), prepend<3>(C.layout())));     // (1,M,N)
+}
+
+// Dispatch [4]: (V,M) x (V,N) => (V,M,N)
+template <class MMA,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout,
+          __MUTE_REQUIRES(DLayout::rank == 3 && is_rmem<TD>::value &&
+                          ALayout::rank == 2 && is_rmem<TA>::value &&
+                          BLayout::rank == 2 && is_rmem<TB>::value &&
+                          CLayout::rank == 3 && is_rmem<TC>::value)>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TD, DLayout>      & D,  // (V,M,N) Logical data
+     Tensor<TA, ALayout> const& A,  // (V,M)   Logical data
+     Tensor<TB, BLayout> const& B,  // (V,N)   Logical data
+     Tensor<TC, CLayout> const& C)  // (V,M,N) Logical data
+{
+  MUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(C));  // AM == CM
+  MUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C));  // BN == CN
+  MUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D));
+  auto M = size<1>(A);
+  auto N = size<1>(B);
+  // REGISTER .reuse OPTIMIZATIONS
+  // 64-bit traversal specialization -- serpentine path
+  if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 8 &&
+                decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 8)
+  {
+#if 1 // NOTE: Row- vs Col- major could depend on the C-matrix order... (which we can test)
+    // Row-major serpentine iteration
+    MUTE_UNROLL
+    for (int m = 0; m < M; ++m) {
+      MUTE_UNROLL
+      for (int n = 0; n < N; ++n) {
+        int ns = (m & 1) ? N-1-n : n;  // Serpentine coordinate
+        gemm(mma, D(_,m,ns), A(_,m), B(_,ns), C(_,m,ns));
+      }
+    }
+#else
+    // Col-major serpentine iteration
+    MUTE_UNROLL
+    for (int n = 0; n < N; ++n) {
+      MUTE_UNROLL
+      for (int m = 0; m < M; ++m) {
+        int ms = (n & 1) ? M-1-m : m;  // Serpentine coordinate
+        gemm(mma, D(_,ms,n), A(_,ms), B(_,n), C(_,ms,n));
+      }
+    }
+#endif
+  } else
+  // 32-bit traversal specialization -- kinked serpentine path
+  if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 4 &&
+                decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 4)
+  {
+#if 1  // NOTE: Row- vs Col- major could depend on the C-matrix order... (which we can test)
+    // Row-major kinked serpentine iteration
+    MUTE_UNROLL
+    for (int m = 0; m < M; m += 2) {
+      MUTE_UNROLL
+      for (int n = 0; n < N; ++n) {
+        int ns = (m & 2) ? N-1-n : n;
+        gemm(mma, D(_,m+0,ns), A(_,m+0), B(_,ns), C(_,m+0,ns));
+
+        if (m+1 < M) {
+          gemm(mma, D(_,m+1,ns), A(_,m+1), B(_,ns), C(_,m+1,ns));
+        }
+      }
+    }
+#else
+    // Col-major kinked serpentine iteration
+    MUTE_UNROLL
+    for (int n = 0; n < N; n += 2) {
+      MUTE_UNROLL
+      for (int m = 0; m < M; ++m) {
+        // Kinked serpentine traversal for maximum register reuse
+        int ms = (n & 2) ? M-1-m : m;
+        gemm(mma, D(_,ms,n+0), A(_,ms), B(_,n+0), C(_,ms,n+0));
+
+        if (n+1 < N) {
+          gemm(mma, D(_,ms,n+1), A(_,ms), B(_,n+1), C(_,ms,n+1));
+        }
+      }
+    }
+#endif
+  } else
+  // 64-bit + 32-bit traversal order -- keep A (64-bit) in the outer loop and serpentine B
+  if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 8 &&
+                decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 4) {
+    // Row-major serpentine iteration
+    MUTE_UNROLL
+    for (int m = 0; m < M; ++m) {
+      MUTE_UNROLL
+      for (int n = 0; n < N; ++n) {
+        int ns = (m & 1) ? N-1-n : n;  // Serpentine coordinate
+        gemm(mma, D(_,m,ns), A(_,m), B(_,ns), C(_,m,ns));
+      }
+    }
+  } else
+  // 32-bit + 64-bit traversal order -- keep B (64-bit) in the outer loop and serpentine A
+  if constexpr (decltype(size<0>(A))::value * sizeof(typename TA::value_type) == 4 &&
+                decltype(size<0>(B))::value * sizeof(typename TB::value_type) == 8) {
+    // Col-major serpentine iteration
+    MUTE_UNROLL
+    for (int n = 0; n < N; ++n) {
+      MUTE_UNROLL
+      for (int m = 0; m < M; ++m) {
+        int ms = (n & 1) ? M-1-m : m;  // Serpentine coordinate
+        gemm(mma, D(_,ms,n), A(_,ms), B(_,n), C(_,ms,n));
+      }
+    }
+  } else
+  // Fallback to serpentine loop
+  {
+    // Col-major serpentine iteration
+    MUTE_UNROLL
+    for (int n = 0; n < N; ++n) {
+      MUTE_UNROLL
+      for (int m = 0; m < M; ++m) {
+        int ms = (n & 1) ? M-1-m : m;  // Serpentine coordinate
+        gemm(mma, D(_,ms,n), A(_,ms), B(_,n), C(_,ms,n));
+      }
+    }
+  }
+}
+
+// Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N)
+template <class MMA,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout,
+          __MUTE_REQUIRES(DLayout::rank == 3 && is_rmem<TD>::value &&
+                          ALayout::rank == 3 && is_rmem<TA>::value &&
+                          BLayout::rank == 3 && is_rmem<TB>::value &&
+                          CLayout::rank == 3 && is_rmem<TC>::value)>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TD, DLayout>      & D,  // (V,M,N) Logical data
+     Tensor<TA, ALayout> const& A,  // (V,M,K) Logical data
+     Tensor<TB, BLayout> const& B,  // (V,N,K) Logical data
+     Tensor<TC, CLayout> const& C)  // (V,M,N) Logical data
+{
+  MUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(C));  // AM == CM
+  MUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C));  // BN == CN
+  MUTE_STATIC_ASSERT_V(size<2>(A) == size<2>(B));  // AK == BK
+  MUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D));
+  auto K = size<2>(A);
+
+  MUTE_UNROLL
+  for (int k = 0; k < K; ++k) {
+    gemm(mma, D, A(_,_,k), B(_,_,k), C);
+  }
+}
+
+//
+// Thread-Local Shared-Memory GEMMs
+//
+
+// Dispatch [1]: (V) x (V) => (V)
+// Dispatch [2]: (M) x (N) => (M,N)
+// Dispatch [3]: (M,K) x (N,K) => (M,N)
+// Dispatch [4]: (V,M) x (V,N) => (V,M,N)
+// Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N)
+// Dispatch [3]: (M,K) x (N,K) => (M,N)
+template <class MMA,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout,
+          __MUTE_REQUIRES(DLayout::rank == 2 && is_rmem<TD>::value &&
+                          ALayout::rank == 2 && is_smem<TA>::value &&
+                          BLayout::rank == 2 && is_smem<TB>::value &&
+                          CLayout::rank == 2 && is_rmem<TC>::value)>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TD, DLayout>      & D,  // (M,N) Logical data
+     Tensor<TA, ALayout> const& A,  // (M,K) Logical data
+     Tensor<TB, BLayout> const& B,  // (N,K) Logical data
+     Tensor<TC, CLayout> const& C)  // (M,N) Logical data
+{
+  MUTE_STATIC_ASSERT_V(size<0>(A) == size<0>(C));  // AM == CM
+  MUTE_STATIC_ASSERT_V(size<0>(B) == size<1>(C));  // BN == CN
+  MUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(B));  // AK == BK
+  MUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D));
+
+  // Assert this is a 1-value MMA
+  MUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutC_TV{}) == Int<1>{});
+  MUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutA_TV{}) == Int<1>{});
+  MUTE_STATIC_ASSERT_V(size<1>(typename MMA_Atom<MMA>::LayoutB_TV{}) == Int<1>{});
+
+  gemm(mma,
+       make_tensor(D.data(), prepend<3>(D.layout())),      // (1,M,N)
+       make_tensor(A.data(), prepend<3>(A.layout())),      // (1,M,K)
+       make_tensor(B.data(), prepend<3>(B.layout())),      // (1,N,K)
+       make_tensor(C.data(), prepend<3>(C.layout())));     // (1,M,N)
+}
+
+// Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N)
+template <class MMA,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout,
+          __MUTE_REQUIRES(DLayout::rank == 3 && is_rmem<TD>::value &&
+                          ALayout::rank == 3 && is_smem<TA>::value &&
+                          BLayout::rank == 3 && is_smem<TB>::value &&
+                          CLayout::rank == 3 && is_rmem<TC>::value)>
+MUTE_HOST_DEVICE
+void
+gemm(MMA_Atom<MMA>       const& mma,
+     Tensor<TD, DLayout>      & D,  // (V,M,N) Logical data
+     Tensor<TA, ALayout> const& A,  // (V,M,K) Logical data
+     Tensor<TB, BLayout> const& B,  // (V,N,K) Logical data
+     Tensor<TC, CLayout> const& C)  // (V,M,N) Logical data
+{
+  MUTE_STATIC_ASSERT_V(size<1>(A) == size<1>(C));  // AM == CM
+  MUTE_STATIC_ASSERT_V(size<1>(B) == size<2>(C));  // BN == CN
+  MUTE_STATIC_ASSERT_V(size<2>(A) == size<2>(B));  // AK == BK
+  MUTE_STATIC_ASSERT_V(size<0>(C) == size<0>(D) && size<1>(C) == size<1>(D) && size<2>(C) == size<2>(D));
+
+  auto rA = MMA_Atom<MMA>::make_fragment_A(A);
+  auto rB = MMA_Atom<MMA>::make_fragment_B(B);
+
+  auto K = size<2>(A);
+
+  MUTE_UNROLL
+  for (int k = 0; k < K; ++k)
+  {
+    copy(A(_,_,k), rA(_,_,k));
+    copy(B(_,_,k), rB(_,_,k));
+    // Thread-level register gemm for k
+    gemm(mma, D, rA(_,_,k), rB(_,_,k), C);
+  }
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/prefer.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/prefer.hpp
new file mode 100644
index 000000000..539590424
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/prefer.hpp
@@ -0,0 +1,47 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+namespace mute
+{
+
+// Infinite types that inherit from each other
+template <size_t N>
+struct prefer : prefer<N-1> {};
+
+template <>
+struct prefer<0> {};
+
+// Can be used to preferencially overload implementations
+// Higher N in prefer<N> have higher priority.
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/prefetch.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/prefetch.hpp
new file mode 100644
index 000000000..5abe01cdd
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/prefetch.hpp
@@ -0,0 +1,117 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>          // MUTE_HOST_DEVICE
+#include <mute/tensor.hpp>          // mute::Tensor
+#include <mute/atom/copy_atom.hpp>  // mute::Copy_Atom
+namespace mute
+{
+
+//
+// Prefetch global tensors into L2
+//
+
+template <uint32_t NumThreads, uint32_t FetchBytes = 64,
+          class GEngine, class GLayout>
+MUTE_HOST_DEVICE
+void
+cooperative_prefetch(uint32_t                 const& tid,
+                     Tensor<GEngine, GLayout> const& src)
+{
+  static_assert(is_gmem<GEngine>::value, "Expected global tensor for prefetch");
+
+  constexpr int V = decltype(max_common_vector(src, src))::value;
+
+  if constexpr (V > 1) {
+    // L2 sector is 32B, default fetch granularity is 64B
+    using VecType = conditional_t<(V * sizeof_bits_v<typename GEngine::value_type>) < (FetchBytes * 8),
+                                  ArrayEngine<typename GEngine::value_type, V>,
+                                  uint8_t[FetchBytes]                        >;
+
+    Tensor src_v = recast<VecType const>(src);
+    MUTE_UNROLL
+    for (int i = tid; i < size(src_v); i += NumThreads) {
+      prefetch(raw_pointer_cast(&src_v(i)));
+    }
+  } else {
+    MUTE_UNROLL
+    for (int i = tid; i < size(src); i += NumThreads) {
+      prefetch(raw_pointer_cast(&src(i)));
+    }
+  }
+}
+
+template <class GEngine, class GLayout>
+MUTE_HOST_DEVICE
+void
+prefetch(Tensor<GEngine, GLayout> const& src)
+{
+  return cooperative_prefetch<1>(0, src);
+}
+
+// Prefetch with copy atom
+namespace detail {
+
+template <class CopyOp, class = void>
+constexpr bool has_prefetch = false;
+
+template <class CopyOp>
+constexpr bool has_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = true;
+
+template <class CopyOp, class = void>
+constexpr bool is_prefetch = false;
+
+template <class CopyOp>
+constexpr bool is_prefetch<CopyOp, void_t<typename CopyOp::PREFETCH>> = is_same_v<CopyOp, typename CopyOp::PREFETCH>;
+
+} // end namespace detail
+
+template <class CopyOp, class... CT_Args, class... CA_Args,
+          class GEngine, class GLayout>
+MUTE_HOST_DEVICE
+void
+prefetch(Copy_Atom<Copy_Traits<CopyOp, CT_Args...>, CA_Args...> const& atom,
+         Tensor<GEngine, GLayout>                               const& src)
+{
+  if constexpr (detail::has_prefetch<CopyOp>) {
+    using Prefetch_Traits = Copy_Traits<typename CopyOp::PREFETCH, CT_Args...>;
+    using Prefetch_Atom = Copy_Atom<Prefetch_Traits, CA_Args...>;
+    Prefetch_Atom prefetch_atom{atom};
+    auto& dst = const_cast<Tensor<GEngine, GLayout>&>(src); // dst is ignored for prefetch atoms
+    return copy(prefetch_atom, src, dst);
+  } else {
+    return prefetch(src);
+  }
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/tensor_algorithms.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/tensor_algorithms.hpp
new file mode 100644
index 000000000..780cbab54
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/tensor_algorithms.hpp
@@ -0,0 +1,168 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+/** Common algorithms on (hierarchical) tensors */
+
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/tensor.hpp>
+
+namespace mute
+{
+
+//
+// for_each
+//
+
+template <class Engine, class Layout, class UnaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+for_each(Tensor<Engine,Layout> const& tensor, UnaryOp&& op)
+{
+  MUTE_UNROLL
+  for (int i = 0; i < size(tensor); ++i) {
+    op(tensor(i));
+  }
+}
+
+template <class Engine, class Layout, class UnaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+for_each(Tensor<Engine,Layout>& tensor, UnaryOp&& op)
+{
+  MUTE_UNROLL
+  for (int i = 0; i < size(tensor); ++i) {
+    op(tensor(i));
+  }
+}
+
+// Accept mutable temporaries
+template <class Engine, class Layout, class UnaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+for_each(Tensor<Engine,Layout>&& tensor, UnaryOp&& op)
+{
+  return for_each(tensor, op);
+}
+
+//
+// transform
+//
+
+// Similar to std::transform but does not return number of elements affected
+template <class Engine, class Layout, class UnaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+transform(Tensor<Engine,Layout>& tensor, UnaryOp&& op)
+{
+  MUTE_UNROLL
+  for (int i = 0; i < size(tensor); ++i) {
+    tensor(i) = op(tensor(i));
+  }
+}
+
+// Accept mutable temporaries
+template <class Engine, class Layout, class UnaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+transform(Tensor<Engine,Layout>&& tensor, UnaryOp&& op)
+{
+  return transform(tensor, op);
+}
+
+// Similar to std::transform transforms one tensors and assigns it to another
+template <class EngineIn, class LayoutIn,
+          class EngineOut, class LayoutOut,
+          class UnaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
+          Tensor<EngineOut,LayoutOut>      & tensor_out,
+          UnaryOp&& op)
+{
+  MUTE_UNROLL
+  for (int i = 0; i < size(tensor_in); ++i) {
+    tensor_out(i) = op(tensor_in(i));
+  }
+}
+
+// Accept mutable temporaries
+template <class EngineIn, class LayoutIn,
+          class EngineOut, class LayoutOut,
+          class UnaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+transform(Tensor<EngineIn, LayoutIn > const& tensor_in,
+          Tensor<EngineOut,LayoutOut>     && tensor_out,
+          UnaryOp&& op)
+{
+  return transform(tensor_in, tensor_out, op);
+}
+
+// Similar to std::transform with a binary operation
+// Takes two tensors as input and one tensor as output.
+// Applies the binary_op to tensor_in1 and tensor_in2 and
+// assigns it to tensor_out
+template <class EngineIn1, class LayoutIn1,
+          class EngineIn2, class LayoutIn2,
+          class EngineOut, class LayoutOut,
+          class BinaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
+          Tensor<EngineIn2,LayoutIn2> const& tensor_in2,
+          Tensor<EngineOut,LayoutOut>      & tensor_out,
+          BinaryOp&& op)
+{
+  MUTE_UNROLL
+  for (int i = 0; i < size(tensor_in1); ++i) {
+    tensor_out(i) = op(tensor_in1(i), tensor_in2(i));
+  }
+}
+
+// Accept mutable temporaries
+template <class EngineIn1, class LayoutIn1,
+          class EngineIn2, class LayoutIn2,
+          class EngineOut, class LayoutOut,
+          class BinaryOp>
+MUTE_HOST_DEVICE constexpr
+void
+transform(Tensor<EngineIn1,LayoutIn1> const& tensor_in1,
+          Tensor<EngineIn2,LayoutIn2> const& tensor_in2,
+          Tensor<EngineOut,LayoutOut>     && tensor_out,
+          BinaryOp&& op)
+{
+  return transform(tensor_in1, tensor_in2, tensor_out, op);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/algorithm/tuple_algorithms.hpp b/vllm_musa/csrc/flash_attn/include/mute/algorithm/tuple_algorithms.hpp
new file mode 100644
index 000000000..759b5733d
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/algorithm/tuple_algorithms.hpp
@@ -0,0 +1,1082 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/util/type_traits.hpp>
+#include <mute/container/tuple.hpp>
+#include <mute/algorithm/functional.hpp>
+#include <mute/numeric/integer_sequence.hpp>
+#include <mute/numeric/integral_constant.hpp>
+
+/// @file tuple_algorithms.hpp
+/// @brief Common algorithms on (hierarchical) tuples
+///
+/// Code guidelines and style preferences:
+///
+/// For perfect forwarding, don't use std::forward, because it may not
+/// be defined in device code when compiling with NVRTC.  Instead, use
+/// `static_cast<ParameterType&&>(parameter_name)`.
+///
+/// MuTe generally does not bother forwarding functions, as
+/// reference-qualified member functions are rare in this code base.
+///
+/// Throughout MUTLASS, mute::make_tuple always needs to be called
+/// namespace-qualified, EVEN If inside the mute namespace and/or in
+/// scope of a "using namespace mute" declaration.  Otherwise, the
+/// compiler may select std::make_tuple instead of mute::make_tuple,
+/// due to argument-dependent lookup.  Two problems may result from
+/// that.
+///
+/// 1. Functions have an unexpected return type (std::tuple instead of
+///    mute::tuple), so functions that take mute::tuple parameters
+///    fail to compile (generally inside functions that have template
+///    parameters expected to be mute::tuple).
+///
+/// 2. std::tuple does not have the required __host__ __device__
+///    markings, so the MUSA compiler complains if you use it in
+///    device code.
+///
+/// mute::make_tuple will occur more often than std::make_tuple would
+/// in modern C++ code, because mute::tuple's design deprioritizes
+/// correct operation of CTAD (constructor template argument
+/// deduction) in favor of implementation simplicity.
+
+namespace mute
+{
+
+//
+// Apply (Unpack)
+// (t, f) => f(t_0,t_1,...,t_n)
+//
+
+namespace detail {
+
+template <class T, class F, int... I>
+MUTE_HOST_DEVICE constexpr
+auto
+apply(T&& t, F&& f, seq<I...>)
+{
+  return f(get<I>(static_cast<T&&>(t))...);
+}
+
+} // end namespace detail
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+apply(T&& t, F&& f)
+{
+  return detail::apply(static_cast<T&&>(t), f, tuple_seq<T>{});
+}
+
+//
+// Transform Apply
+// (t, f, g) => g(f(t_0),f(t_1),...)
+//
+
+namespace detail {
+
+template <class T, class F, class G, int... I>
+MUTE_HOST_DEVICE constexpr
+auto
+tapply(T&& t, F&& f, G&& g, seq<I...>)
+{
+  return g(f(get<I>(static_cast<T&&>(t)))...);
+}
+
+template <class T0, class T1, class F, class G, int... I>
+MUTE_HOST_DEVICE constexpr
+auto
+tapply(T0&& t0, T1&& t1, F&& f, G&& g, seq<I...>)
+{
+  return g(f(get<I>(static_cast<T0&&>(t0)),
+             get<I>(static_cast<T1&&>(t1)))...);
+}
+
+template <class T0, class T1, class T2, class F, class G, int... I>
+MUTE_HOST_DEVICE constexpr
+auto
+tapply(T0&& t0, T1&& t1, T2&& t2, F&& f, G&& g, seq<I...>)
+{
+  return g(f(get<I>(static_cast<T0&&>(t0)),
+             get<I>(static_cast<T1&&>(t1)),
+             get<I>(static_cast<T2&&>(t2)))...);
+}
+
+} // end namespace detail
+
+template <class T, class F, class G>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_apply(T&& t, F&& f, G&& g)
+{
+  if constexpr (is_tuple<remove_cvref_t<T>>::value) {
+    return detail::tapply(static_cast<T&&>(t), f, g, tuple_seq<T>{});
+  } else {
+    return g(f(static_cast<T&&>(t)));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T0, class T1, class F, class G>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_apply(T0&& t0, T1&& t1, F&& f, G&& g)
+{
+  if constexpr (is_tuple<remove_cvref_t<T0>>::value) {
+    return detail::tapply(static_cast<T0&&>(t0), static_cast<T1&&>(t1), f, g, tuple_seq<T0>{});
+  } else {
+    return g(f(static_cast<T0&&>(t0), static_cast<T1&&>(t1)));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T0, class T1, class T2, class F, class G>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_apply(T0&& t0, T1&& t1, T2&& t2, F&& f, G&& g)
+{
+  if constexpr (is_tuple<remove_cvref_t<T0>>::value) {
+    return detail::tapply(static_cast<T0&&>(t0), static_cast<T1&&>(t1), static_cast<T2&&>(t2), f, g, tuple_seq<T0>{});
+  } else {
+    return g(f(static_cast<T0&&>(t0), static_cast<T1&&>(t1), static_cast<T2&&>(t2)));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// For Each
+// (t, f) => f(t_0),f(t_1),...,f(t_n)
+//
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+void
+for_each(T&& t, F&& f)
+{
+  if constexpr (is_tuple<remove_cvref_t<T>>::value) {
+    return detail::apply(t, [&](auto&&... a) { (f(static_cast<decltype(a)&&>(a)), ...); }, tuple_seq<T>{});
+  } else {
+    return f(static_cast<T&&>(t));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+for_each_leaf(T&& t, F&& f)
+{
+  if constexpr (is_tuple<remove_cvref_t<T>>::value) {
+    return detail::apply(static_cast<T&&>(t), [&](auto&&... a){ return (for_each_leaf(static_cast<decltype(a)&&>(a), f), ...); }, tuple_seq<T>{});
+  } else {
+    return f(static_cast<T&&>(t));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Transform
+// (t, f) => (f(t_0),f(t_1),...,f(t_n))
+//
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+transform(T const& t, F&& f)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::tapply(t, f, [](auto const&... a){ return mute::make_tuple(a...); }, tuple_seq<T>{});
+  } else {
+    return f(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T0, class T1, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+transform(T0 const& t0, T1 const& t1, F&& f)
+{
+  if constexpr (is_tuple<T0>::value) {
+    static_assert(tuple_size<T0>::value == tuple_size<T1>::value, "Mismatched tuple_size");
+    return detail::tapply(t0, t1, f, [](auto const&... a){ return mute::make_tuple(a...); }, tuple_seq<T0>{});
+  } else {
+    return f(t0, t1);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T0, class T1, class T2, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+transform(T0 const& t0, T1 const& t1, T2 const& t2, F&& f)
+{
+  if constexpr (is_tuple<T0>::value) {
+    static_assert(tuple_size<T0>::value == tuple_size<T1>::value, "Mismatched tuple_size");
+    static_assert(tuple_size<T0>::value == tuple_size<T2>::value, "Mismatched tuple_size");
+    return detail::tapply(t0, t1, t2, f, [](auto const&... a){ return mute::make_tuple(a...); }, tuple_seq<T0>{});
+  } else {
+    return f(t0, t1, t2);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_leaf(T const& t, F&& f)
+{
+  if constexpr (is_tuple<T>::value) {
+    return transform(t, [&](auto const& a) { return transform_leaf(a, f); });
+  } else {
+    return f(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T0, class T1, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_leaf(T0 const& t0, T1 const& t1, F&& f)
+{
+  if constexpr (is_tuple<T0>::value) {
+    return transform(t0, t1, [&](auto const& a, auto const& b) { return transform_leaf(a, b, f); });
+  } else {
+    return f(t0, t1);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// find and find_if
+//
+
+namespace detail {
+
+template <class T, class F, int I, int... Is>
+MUTE_HOST_DEVICE constexpr
+auto
+find_if(T const& t, F&& f, seq<I,Is...>)
+{
+  if constexpr (decltype(f(get<I>(t)))::value) {
+    return mute::C<I>{};
+  } else
+  if constexpr (sizeof...(Is) == 0) {
+    return mute::C<I+1>{};
+  } else {
+    return find_if(t, f, seq<Is...>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+find_if(T const& t, F&& f)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::find_if(t, f, tuple_seq<T>{});
+  } else {
+    return mute::C<decltype(f(t))::value ? 0 : 1>{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+find(T const& t, X const& x)
+{
+  return find_if(t, [&](auto const& v) { return v == x; });  // This should always return a static true/false
+}
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+any_of(T const& t, F&& f)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::apply(mute::transform(t, f), [&] (auto const&... a) { return (false_type{} || ... || a); }, tuple_seq<T>{});
+  } else {
+    return f(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+all_of(T const& t, F&& f)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::apply(t, [&] (auto const&... a) { return (true_type{} && ... && f(a)); }, tuple_seq<T>{});
+  } else {
+    return f(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+none_of(T const& t, F&& f)
+{
+  return not any_of(t, f);
+}
+
+//
+// Filter
+// (t, f) => <f(t_0),f(t_1),...,f(t_n)>
+//
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+filter_tuple(T const& t, F&& f)
+{
+  return transform_apply(t, f, [](auto const&... a) { return mute::tuple_cat(a...); });
+}
+
+template <class T0, class T1, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+filter_tuple(T0 const& t0, T1 const& t1, F&& f)
+{
+  return transform_apply(t0, t1, f, [](auto const&... a) { return mute::tuple_cat(a...); });
+}
+
+template <class T0, class T1, class T2, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+filter_tuple(T0 const& t0, T1 const& t1, T2 const& t2, F&& f)
+{
+  return transform_apply(t0, t1, t2, f, [](auto const&... a) { return mute::tuple_cat(a...); });
+}
+
+//
+// Fold (Reduce, Accumulate)
+// (t, v, f) => f(...f(f(v,t_0),t_1),...,t_n)
+//
+
+namespace detail {
+
+template <class Fn, class Val>
+struct FoldAdaptor {
+  template <class X>
+  MUTE_HOST_DEVICE constexpr auto operator|(X&& x) {
+    auto r = fn_(val_, static_cast<X&&>(x));
+    return FoldAdaptor<Fn, decltype(r)>{fn_, r};
+  }
+  Fn fn_;
+  Val val_;
+};
+
+template <class T, class V, class F, int... Is>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+fold(T&& t, V const& v, F&& f, seq<Is...>)
+{
+  return (FoldAdaptor<F,V>{f,v} | ... | get<Is>(static_cast<T&&>(t))).val_;
+}
+
+} // end namespace detail
+
+template <class T, class V, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+fold(T&& t, V const& v, F&& f)
+{
+  if constexpr (is_tuple<remove_cvref_t<T>>::value) {
+    return detail::fold(static_cast<T&&>(t), v, f, tuple_seq<T>{});
+  } else {
+    return f(v, static_cast<T&&>(t));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class F>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+fold_first(T&& t, F&& f)
+{
+  if constexpr (is_tuple<remove_cvref_t<T>>::value) {
+    return detail::fold(static_cast<T&&>(t), get<0>(t), f, make_range<1,tuple_size<remove_cvref_t<T>>::value>{});
+  } else {
+    return static_cast<T&&>(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// front, back, take, select, unwrap
+//
+
+// Get the first non-tuple element in a hierarchical tuple
+template <class T>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+front(T&& t)
+{
+  if constexpr (is_tuple<remove_cvref_t<T>>::value) {
+    return front(get<0>(static_cast<T&&>(t)));
+  } else {
+    return static_cast<T&&>(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Get the last non-tuple element in a hierarchical tuple
+template <class T>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+back(T&& t)
+{
+  if constexpr (is_tuple<remove_cvref_t<T>>::value) {
+    constexpr int N = tuple_size<remove_cvref_t<T>>::value;
+
+    // MSVC needs a bit of extra help here deducing return types.
+    // We help it by peeling off the nonrecursive case a level "early."
+    if constexpr (! is_tuple<remove_cvref_t<decltype(get<N - 1>(static_cast<T&&>(t)))>>::value) {
+      return get<N - 1>(static_cast<T&&>(t));
+    } else {
+      return back(get<N - 1>(static_cast<T&&>(t)));
+    }
+  } else {
+    return static_cast<T&&>(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Takes the elements in the range [B,E)
+template <int B, int E, class T>
+MUTE_HOST_DEVICE constexpr
+auto
+take(T const& t)
+{
+  if constexpr (E == -1) {
+    if constexpr (is_tuple<T>::value) {
+      return take<B,tuple_size<T>::value>(t);
+    } else {
+      return take<B,1>(t);
+    }
+  } else if constexpr (B <= E) {
+    return detail::apply(t, [](auto const&... a) { return mute::make_tuple(a...); }, make_range<B,E>{});
+  } else {
+    static_assert(B <= E);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Select tuple elements with given indices.
+
+template <int... I, class T>
+MUTE_HOST_DEVICE constexpr
+auto
+select(T const& t)
+{
+  return mute::make_tuple(get<I>(t)...);
+}
+
+// Wrap non-tuples into rank-1 tuples or forward
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+wrap(T const& t)
+{
+  if constexpr (is_tuple<T>::value) {
+    return t;
+  } else {
+    return mute::make_tuple(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Unwrap rank-1 tuples until we're left with a rank>1 tuple or a non-tuple
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+unwrap(T const& t)
+{
+  if constexpr (is_tuple<T>::value) {
+    if constexpr (tuple_size<T>::value == 1) {
+      return unwrap(get<0>(t));
+    } else {
+      return t;
+    }
+  } else {
+    return t;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Flatten and Unflatten
+//
+
+template <class T>
+struct is_flat : true_type {};
+
+template <class... Ts>
+struct is_flat<tuple<Ts...>> : bool_constant<(true && ... && (not is_tuple<Ts>::value))> {};
+
+// Flatten a hierarchical tuple to a tuple of depth one
+//   and wrap non-tuples into a rank-1 tuple.
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+flatten_to_tuple(T const& t)
+{
+  if constexpr (is_tuple<T>::value) {
+    if constexpr (is_flat<T>::value) {      // Shortcut for perf
+      return t;
+    } else {
+      return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); });
+    }
+  } else {
+    return mute::make_tuple(t);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Flatten a hierarchical tuple to a tuple of depth one
+//   and leave non-tuple untouched.
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+flatten(T const& t)
+{
+  if constexpr (is_tuple<T>::value) {
+    if constexpr (is_flat<T>::value) {      // Shortcut for perf
+      return t;
+    } else {
+      return filter_tuple(t, [](auto const& a) { return flatten_to_tuple(a); });
+    }
+  } else {
+    return t;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+namespace detail {
+
+template <class FlatTuple, class TargetProfile>
+MUTE_HOST_DEVICE constexpr
+auto
+unflatten_impl(FlatTuple const& flat_tuple, TargetProfile const& target_profile)
+{
+  if constexpr (is_tuple<TargetProfile>::value) {
+    return fold(target_profile, mute::make_tuple(mute::make_tuple(), flat_tuple), [](auto const& v, auto const& t) {
+      auto [result, remaining_tuple] = v;
+      auto [sub_result, sub_tuple] = unflatten_impl(remaining_tuple, t);
+      return mute::make_tuple(append(result, sub_result), sub_tuple);
+    });
+  } else {
+    return mute::make_tuple(get<0>(flat_tuple), take<1, decltype(rank(flat_tuple))::value>(flat_tuple));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+}  // end namespace detail
+
+// Unflatten a flat tuple into a hierarchical tuple
+// @pre flatten(@a flat_tuple) == @a flat_tuple
+// @pre rank(flatten(@a target_profile)) == rank(@a flat_tuple)
+// @post congruent(@a result, @a target_profile)
+// @post flatten(@a result) == @a flat_tuple
+template <class FlatTuple, class TargetProfile>
+MUTE_HOST_DEVICE constexpr
+auto
+unflatten(FlatTuple const& flat_tuple, TargetProfile const& target_profile)
+{
+  auto [unflatten_tuple, flat_remainder] = detail::unflatten_impl(flat_tuple, target_profile);
+  MUTE_STATIC_ASSERT_V(rank(flat_remainder) == Int<0>{});
+  return unflatten_tuple;
+}
+
+//
+// insert and remove and replace
+//
+
+namespace detail {
+
+// Shortcut around mute::tuple_cat for common insert/remove/repeat cases
+template <class T, class X, int... I, int... J, int... K>
+MUTE_HOST_DEVICE constexpr
+auto
+construct(T const& t, X const& x, seq<I...>, seq<J...>, seq<K...>)
+{
+  return mute::make_tuple(get<I>(t)..., (void(J),x)..., get<K>(t)...);
+}
+
+} // end namespace detail
+
+// Insert x into the Nth position of the tuple
+template <int N, class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+insert(T const& t, X const& x)
+{
+  return detail::construct(t, x, make_seq<N>{}, seq<0>{}, make_range<N,tuple_size<T>::value>{});
+}
+
+// Remove the Nth element of the tuple
+template <int N, class T>
+MUTE_HOST_DEVICE constexpr
+auto
+remove(T const& t)
+{
+  return detail::construct(t, 0, make_seq<N>{}, seq<>{}, make_range<N+1,tuple_size<T>::value>{});
+}
+
+// Replace the Nth element of the tuple with x
+template <int N, class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+replace(T const& t, X const& x)
+{
+  return detail::construct(t, x, make_seq<N>{}, seq<0>{}, make_range<N+1,tuple_size<T>::value>{});
+}
+
+// Replace the first element of the tuple with x
+template <class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+replace_front(T const& t, X const& x)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::construct(t, x, seq<>{}, seq<0>{}, make_range<1,tuple_size<T>::value>{});
+  } else {
+    return x;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Replace the last element of the tuple with x
+template <class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+replace_back(T const& t, X const& x)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::construct(t, x, make_seq<tuple_size<T>::value-1>{}, seq<0>{}, seq<>{});
+  } else {
+    return x;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Make a tuple of Xs of tuple_size N
+//
+
+template <int N, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_repeat(X const& x)
+{
+  return detail::construct(0, x, seq<>{}, make_seq<N>{}, seq<>{});
+}
+
+//
+// Make repeated Xs of rank N
+//
+
+template <int N, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+repeat(X const& x)
+{
+  if constexpr (N == 1) {
+    return x;
+  } else {
+    return detail::construct(0, x, seq<>{}, make_seq<N>{}, seq<>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Make a tuple of Xs the same profile as tuple T
+//
+
+template <class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+repeat_like(T const& t, X const& x)
+{
+  if constexpr (is_tuple<T>::value) {
+    return transform(t, [&](auto const& a) { return repeat_like(a,x); });
+  } else {
+    return x;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Group the elements [B,E) of a T into a single element
+// e.g. group<2,4>(T<_1,_2,_3,_4,_5,_6>{})
+//              => T<_1,_2,T<_3,_4>,_5,_6>{}
+template <int B, int E, class T>
+MUTE_HOST_DEVICE constexpr
+auto
+group(T const& t)
+{
+  if constexpr (not is_tuple<T>::value) {
+    if constexpr (E == -1) {
+      return group<B,1>(t);
+    } else {
+      return detail::construct(t, take<B,E>(t), make_seq<B>{}, make_seq<(B < E)>{}, make_range<E,1>{});
+    }
+  } else
+  if constexpr (E == -1) {
+    return group<B,tuple_size<T>::value>(t);
+  } else
+  if constexpr (B <= E) {
+    return detail::construct(t, take<B,E>(t), make_seq<B>{}, make_seq<(B < E)>{}, make_range<E,tuple_size<T>::value>{});
+  } else {
+    static_assert(B <= E);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Extend a T to rank N by appending/prepending an element
+//
+
+template <int N, class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+append(T const& a, X const& x)
+{
+  if constexpr (is_tuple<T>::value) {
+    if constexpr (N == tuple_size<T>::value) {
+      return a;
+    } else {
+      static_assert(N > tuple_size<T>::value);
+      return detail::construct(a, x, make_seq<tuple_size<T>::value>{}, make_seq<N-tuple_size<T>::value>{}, seq<>{});
+    }
+  } else {
+    if constexpr (N == 1) {
+      return a;
+    } else {
+      return detail::construct(mute::make_tuple(a), x, seq<0>{}, make_seq<N-1>{}, seq<>{});
+    }
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+append(T const& a, X const& x)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::construct(a, x, make_seq<tuple_size<T>::value>{}, seq<0>{}, seq<>{});
+  } else {
+    return mute::make_tuple(a, x);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <int N, class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+prepend(T const& a, X const& x)
+{
+  if constexpr (is_tuple<T>::value) {
+    if constexpr (N == tuple_size<T>::value) {
+      return a;
+    } else {
+      static_assert(N > tuple_size<T>::value);
+      return detail::construct(a, x, seq<>{}, make_seq<N-tuple_size<T>::value>{}, make_seq<tuple_size<T>::value>{});
+    }
+  } else {
+    if constexpr (N == 1) {
+      return a;
+    } else {
+      static_assert(N > 1);
+      return detail::construct(mute::make_tuple(a), x, seq<>{}, make_seq<N-1>{}, seq<0>{});
+    }
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+prepend(T const& a, X const& x)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::construct(a, x, seq<>{}, seq<0>{}, make_seq<tuple_size<T>::value>{});
+  } else {
+    return mute::make_tuple(x, a);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Inclusive scan (prefix sum)
+//
+
+namespace detail {
+
+template <class T, class V, class F, int I, int... Is>
+MUTE_HOST_DEVICE constexpr
+auto
+iscan(T const& t, V const& v, F&& f, seq<I,Is...>)
+{
+  // Apply the function to v and the element at I
+  auto v_next = f(v, get<I>(t));
+  // Replace I with v_next
+  auto t_next = replace<I>(t, v_next);
+
+#if 0
+  std::cout << "ISCAN i" << I << std::endl;
+  std::cout << "  t      " << t << std::endl;
+  std::cout << "  i      " << v << std::endl;
+  std::cout << "  f(i,t) " << v_next << std::endl;
+  std::cout << "  t_n    " << t_next << std::endl;
+#endif
+
+  if constexpr (sizeof...(Is) == 0) {
+    return t_next;
+  } else {
+    return iscan(t_next, v_next, f, seq<Is...>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+template <class T, class V, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+iscan(T const& t, V const& v, F&& f)
+{
+  return detail::iscan(t, v, f, tuple_seq<T>{});
+}
+
+//
+// Exclusive scan (prefix sum)
+//
+
+namespace detail {
+
+template <class T, class V, class F, int I, int... Is>
+MUTE_HOST_DEVICE constexpr
+auto
+escan(T const& t, V const& v, F&& f, seq<I,Is...>)
+{
+  if constexpr (sizeof...(Is) == 0) {
+    // Replace I with v
+    return replace<I>(t, v);
+  } else {
+    // Apply the function to v and the element at I
+    auto v_next = f(v, get<I>(t));
+    // Replace I with v
+    auto t_next = replace<I>(t, v);
+
+#if 0
+    std::cout << "ESCAN i" << I << std::endl;
+    std::cout << "  t      " << t << std::endl;
+    std::cout << "  i      " << v << std::endl;
+    std::cout << "  f(i,t) " << v_next << std::endl;
+    std::cout << "  t_n    " << t_next << std::endl;
+#endif
+
+    // Recurse
+    return escan(t_next, v_next, f, seq<Is...>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+template <class T, class V, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+escan(T const& t, V const& v, F&& f)
+{
+  return detail::escan(t, v, f, tuple_seq<T>{});
+}
+
+//
+// Zip (Transpose)
+//
+
+// Take       ((a,b,c,...),(x,y,z,...),...)        rank-R0 x rank-R1 input
+// to produce ((a,x,...),(b,y,...),(c,z,...),...)  rank-R1 x rank-R0 output
+
+namespace detail {
+
+template <int J, class... Ts>
+MUTE_HOST_DEVICE constexpr
+auto
+zip_(Ts const&... ts)
+{
+  return mute::make_tuple(get<J>(ts)...);
+}
+
+template <class T, int... Is, int... Js>
+MUTE_HOST_DEVICE constexpr
+auto
+zip(T const& t, seq<Is...>, seq<Js...>)
+{
+  static_assert(conjunction<bool_constant<tuple_size<tuple_element_t<0,T>>::value == tuple_size<tuple_element_t<Is,T>>::value>...>::value, "Mismatched Ranks");
+  return mute::make_tuple(zip_<Js>(get<Is>(t)...)...);
+}
+
+} // end namespace detail
+
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+zip(T const& t)
+{
+  if constexpr (is_tuple<T>::value) {
+    if constexpr (is_tuple<tuple_element_t<0,T>>::value) {
+      return detail::zip(t, tuple_seq<T>{}, tuple_seq<tuple_element_t<0,T>>{});
+    } else {
+      return mute::make_tuple(t);
+    }
+  } else {
+    return t;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Convenient to pass them in separately
+template <class T0, class T1, class... Ts>
+MUTE_HOST_DEVICE constexpr
+auto
+zip(T0 const& t0, T1 const& t1, Ts const&... ts)
+{
+  return zip(mute::make_tuple(t0, t1, ts...));
+}
+
+//
+// zip2_by -- A guided zip for rank-2 tuples
+//   Take a tuple like ((A,a),((B,b),(C,c)),d)
+//   and produce a tuple ((A,(B,C)),(a,(b,c),d))
+//   where the rank-2 modes are selected by the terminals of the guide (X,(X,X))
+//
+
+namespace detail {
+
+template <class T, class TG, int... Is, int... Js>
+MUTE_HOST_DEVICE constexpr
+auto
+zip2_by(T const& t, TG const& guide, seq<Is...>, seq<Js...>)
+{
+  // zip2_by produces the modes like ((A,a),(B,b),...)
+  auto split = mute::make_tuple(zip2_by(get<Is>(t), get<Is>(guide))...);
+
+  // Rearrange and append missing modes from t to make ((A,B,...),(a,b,...,x,y))
+  return mute::make_tuple(mute::make_tuple(get<0>(get<Is>(split))...),
+                          mute::make_tuple(get<1>(get<Is>(split))..., get<Js>(t)...));
+}
+
+} // end namespace detail
+
+template <class T, class TG>
+MUTE_HOST_DEVICE constexpr
+auto
+zip2_by(T const& t, TG const& guide)
+{
+  if constexpr (is_tuple<TG>::value) {
+    constexpr int TR = tuple_size<T>::value;
+    constexpr int GR = tuple_size<TG>::value;
+    static_assert(TR >= GR, "Mismatched ranks");
+    return detail::zip2_by(t, guide,
+                           make_range< 0, GR>{},
+                           make_range<GR, TR>{});
+  } else {
+    static_assert(tuple_size<T>::value == 2, "Mismatched ranks");
+    return t;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/// @return A tuple of the elements of @c t in reverse order.
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+reverse(T const& t)
+{
+  if constexpr (is_tuple<T>::value) {
+    return detail::apply(t, [](auto const&... a){ return mute::make_tuple(a...); }, tuple_rseq<T>{});
+  } else {
+    return t;
+  }
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/copy.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/copy.hpp
new file mode 100644
index 000000000..b328765b8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/copy.hpp
@@ -0,0 +1,108 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/arch/util.hpp>
+#include <mute/numeric/numeric_types.hpp>
+
+namespace mute
+{
+
+//
+// Direct Copy for any specific types
+//
+
+template <class S, class D = S>
+struct UniversalCopy
+{
+  using SRegisters = S[1];
+  using DRegisters = D[1];
+
+  // Sanity
+  static_assert(sizeof_bits_v<S> >= 8);
+  static_assert(sizeof_bits_v<D> >= 8);
+
+  MUTE_HOST_DEVICE static constexpr void
+  copy(S const& src,
+       D      & dst)
+  {
+    dst = src;
+  }
+};
+
+//
+// Placeholder for the copy algorithm's stronger auto-vectorizing behavior
+//   that assumes alignment of pointers and dynamic layouts up to MaxVecBits
+//
+
+template <int MaxVecBits = 128>
+struct AutoVectorizingCopyWithAssumedAlignment
+     : UniversalCopy<uint_bit_t<MaxVecBits>>
+{
+  static_assert(MaxVecBits == 8 || MaxVecBits == 16 || MaxVecBits == 32 || MaxVecBits == 64 || MaxVecBits == 128,
+                "Expected MaxVecBits to be 8 or 16 or 32 or 64 or 128 for alignment and performance.");
+};
+
+//
+// AutoVectorizingCopy alias assumes maximal alignment of pointers and dynamic strides.
+//   If this is not the case then AutoVectorizingCopyWithAssumedAlignment should be used instead
+//
+
+using AutoVectorizingCopy = AutoVectorizingCopyWithAssumedAlignment<128>;
+
+//
+// DefaultCopy alias does not assume alignment of pointers or dynamic strides.
+//
+
+using DefaultCopy = AutoVectorizingCopyWithAssumedAlignment<8>;
+
+//
+// Copy policy automatically selecting between
+// UniversalCopy and cp.async , based on type and memory space.
+//
+struct AutoCopyAsync {};
+
+//
+// Global memory prefetch into L2
+//
+
+MUTE_HOST_DEVICE static void
+prefetch(void const* gmem_ptr)
+{
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 310)
+  __musa_prefetch_only(gmem_ptr);
+#endif
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31.hpp
new file mode 100644
index 000000000..f836c09fc
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31.hpp
@@ -0,0 +1,212 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 310)
+#define MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED
+#define MUTE_ARCH_TME_MP31_ACTIVATED
+#endif
+
+#include <mute/config.hpp>
+
+#include <mute/arch/copy.hpp>
+#include <mute/arch/copy_mp31_desc.hpp>
+
+
+namespace mute {
+
+//
+// Robust Buffer Access Load/Store
+//
+
+template <class TS, class TD = TS>
+struct MP31_ROBUST_LOAD
+{
+  using SRegisters = TS[1];
+  using DRegisters = TD[1];
+
+  static_assert(sizeof(TS) == sizeof(TD),
+                "MP31_ROBUST_LOAD requires sizeof(src_value_type) == sizeof(dst_value_type)");
+
+  MUTE_HOST_DEVICE static void
+  copy(TS               const& gmem_src,
+       TD                    & rmem_dst,
+       bool                    pred,
+       RobustDescriptor const& src_desc)
+  {
+#if defined(MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED)
+    void const* gmem_ptr = pred ? reinterpret_cast<void const*>(&gmem_src)
+                         : reinterpret_cast<void const*>(src_desc.oob_addr);
+    constexpr PrefetchSize prefetch_size = PrefetchSize::B128;
+    if constexpr (sizeof(TS) == 1) {
+      auto val = __musa_ld_v4_robust_i8(gmem_ptr, src_desc.reg, static_cast<int>(prefetch_size));
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 2) {
+      auto val = __musa_ld_v4_robust_i16(gmem_ptr, src_desc.reg, static_cast<int>(prefetch_size));
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 4) {
+      auto val = __musa_ld_v4_robust_i32(gmem_ptr, src_desc.reg, static_cast<int>(prefetch_size));
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 8) {
+      auto val = __musa_ld_v4_robust_i64(gmem_ptr, src_desc.reg, static_cast<int>(prefetch_size));
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 16) {
+      auto val = __musa_ld_v4_robust_v4i32(gmem_ptr, src_desc.reg, static_cast<int>(prefetch_size));
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 32) {
+      auto val = __musa_ld_v4_robust_v8i32(gmem_ptr, src_desc.reg, static_cast<int>(prefetch_size));
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 64) {
+      auto val = __musa_ld_v4_robust_v16i32(gmem_ptr, src_desc.reg, static_cast<int>(prefetch_size));
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 128) {
+      auto val = __musa_ld_v4_robust_v32i32(gmem_ptr, src_desc.reg, static_cast<int>(prefetch_size));
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else {
+      static_assert(sizeof(TS) <= 128, "Currently we only support up-to 1024-bit load");
+    }
+#else
+    MUTE_INVALID_CONTROL_PATH(
+        "Trying to use MP31_ROBUST_LOAD without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+  }
+};
+
+template <class TS, class TD = TS>
+struct MP31_ROBUST_STORE
+{
+  using SRegisters = TS[1];
+  using DRegisters = TD[1];
+
+  static_assert(sizeof(TS) == sizeof(TD),
+      "MP31_ROBUST_STORE requires sizeof(src_value_type) == sizeof(dst_value_type)");
+
+  MUTE_HOST_DEVICE static void
+  copy(TS               const& rmem_src,
+       TD                    & gmem_dst,
+       bool                    pred,
+       RobustDescriptor const& dst_desc)
+  {
+#if defined(MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED)
+    void* gmem_ptr = pred ? reinterpret_cast<void*>(&gmem_dst)
+                          : reinterpret_cast<void*>(dst_desc.oob_addr);
+    if constexpr (sizeof(TS) == 1) {
+      __musa_st_v4_robust_i8(*reinterpret_cast<int8_t const*>(&rmem_src),
+                             gmem_ptr, dst_desc.reg);
+    } else if constexpr (sizeof(TS) == 2) {
+      __musa_st_v4_robust_i16(*reinterpret_cast<mute::int16_t const*>(&rmem_src),
+                              gmem_ptr, dst_desc.reg);
+    } else if constexpr (sizeof(TS) == 4) {
+      __musa_st_v4_robust_i32(*reinterpret_cast<mute::int32_t const*>(&rmem_src),
+                              gmem_ptr, dst_desc.reg);
+    } else if constexpr (sizeof(TS) == 8) {
+      __musa_st_v4_robust_i64(*reinterpret_cast<mute::int64_t const*>(&rmem_src),
+                              gmem_ptr, dst_desc.reg);
+    } else if constexpr (sizeof(TS) == 16) {
+      __musa_st_v4_robust_v4i32(*reinterpret_cast<mute::int128_t const*>(&rmem_src),
+                                gmem_ptr, dst_desc.reg);
+    } else if constexpr (sizeof(TS) == 32) {
+      __musa_st_v4_robust_v8i32(*reinterpret_cast<mute::int256_t const*>(&rmem_src),
+                                gmem_ptr, dst_desc.reg);
+    } else if constexpr (sizeof(TS) == 64) {
+      __musa_st_v4_robust_v16i32(*reinterpret_cast<mute::int512_t const*>(&rmem_src),
+                                 gmem_ptr, dst_desc.reg);
+    } else if constexpr (sizeof(TS) == 128) {
+      __musa_st_v4_robust_v32i32(*reinterpret_cast<mute::int1024_t const*>(&rmem_src),
+                                 gmem_ptr, dst_desc.reg);
+    } else {
+      static_assert(sizeof(TS) <= 128, "Currently we only support up-to 1024-bit store");
+    }
+#else
+    MUTE_INVALID_CONTROL_PATH(
+        "Trying to use MP31_ROBUST_STORE without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+  }
+};
+
+
+//
+// Robust Buffer Access LDGSTS
+//
+
+template <class TS, class TD = TS>
+struct MP31_ROBUST_LDGSTS
+{
+  using SRegisters = TS[1];
+  using DRegisters = TD[1];
+
+  static_assert(sizeof(TS) == sizeof(TD),
+                "MP31_ROBUST_LDGSTS requires sizeof(src_value_type) == sizeof(dst_value_type)");
+  static_assert(sizeof(TS) == 1  || sizeof(TS) == 2  || sizeof(TS) == 4  || sizeof(TS) == 8 ||
+                sizeof(TS) == 16 || sizeof(TS) == 32 || sizeof(TS) == 64 || sizeof(TS) == 128,
+                "MP31_ROBUST_LDGSTS sizeof(TS) is not supported");
+
+  MUTE_HOST_DEVICE static void
+  copy(TS               const& gmem_src,
+       TD                    & smem_dst,
+       bool                    pred,
+       RobustDescriptor const& src_desc)
+  {
+#if defined(MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED)
+    uint64_t gmem_int_ptr = reinterpret_cast<uint64_t>(&gmem_src);
+    uint32_t smem_int_ptr = cast_smem_ptr_to_uint(&smem_dst);
+    auto gmem_ptr = make_ptr_with_address_space<AddressSpace::Global>(
+                            pred ? gmem_int_ptr : src_desc.oob_addr);
+    auto smem_ptr = make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr);
+    constexpr PrefetchSize prefetch_size = PrefetchSize::B128;
+    __musa_memcpy_g2s_robust_v4(smem_ptr, gmem_ptr, sizeof(TS), src_desc.reg, static_cast<int>(prefetch_size));
+#else
+    MUTE_INVALID_CONTROL_PATH(
+        "Trying to use MP31_ROBUST_LDGSTS without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+  }
+};
+
+MUTE_HOST_DEVICE
+void
+ldgsts_wait() {
+#if defined(MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED)
+  __musa_memcpy_g2s_wait();
+#else
+  MUTE_INVALID_CONTROL_PATH("Trying to use MemcpyG2S Wait without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mute
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include <mute/arch/copy_mp31_tme.hpp>
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31_desc.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31_desc.hpp
new file mode 100644
index 000000000..77b12687e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31_desc.hpp
@@ -0,0 +1,293 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/arch/copy.hpp>
+
+namespace mute
+{
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// Data Prefetch enum
+////////////////////////////////////////////////////////////////////////////////////////////////////
+enum class PrefetchSize : uint8_t {
+  NONE = 0,
+  B64  = 64,
+  B128 = 128,
+};
+
+MUTE_HOST_DEVICE char const* to_string(PrefetchSize const& t) {
+  switch (t) {
+    case PrefetchSize::NONE: return "PrefetchNone";
+    case PrefetchSize::B64:  return "Prefetch64B";
+    case PrefetchSize::B128: return "Prefetch128B";
+  }
+  return nullptr;
+}
+
+#if !defined(__MUSACC_RTC__)
+MUTE_HOST std::ostream& operator<<(std::ostream& os, PrefetchSize const& t) {
+  char const* s = to_string(t);
+  if (s) {
+    std::operator<<(os, s);  // Explicit call to avoid ambiguity
+  } else {
+    os.setstate(std::ios_base::failbit);
+  }
+  return os;
+}
+#endif // !defined(__MUSACC_RTC__)
+
+
+//////////////////////////////////////////////////////////////////////////////////////////////////////
+// Barriers are 32-bit of user-managed information used in broadly two types syncronization patterns
+// 1) arrive/wait on warps
+// 2) transaction-based
+//////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Initialize barrier with its barrier id
+MUTE_HOST_DEVICE
+void
+initialize_barrier(uint32_t barrier_id,                // 32 bits user-managed barrier's id
+                   uint32_t warp_count = 1,            // Warp count expected to arrive/wait on this barrier
+                   uint32_t init_phase = 0)            // Init phase on this barrier
+{
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+  __musa_async_init_arrival(barrier_id, warp_count, init_phase);
+#endif
+}
+
+// Set the number of bytes transfered per transaction
+MUTE_HOST_DEVICE
+void
+set_barrier_transaction_bytes(uint32_t barrier_id,     // 32 bits user-managed barrier's id
+                              uint32_t bytes)          // Number of bytes transfered by per TME transaction
+{
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+  __musa_async_add_trans(barrier_id, bytes);
+#endif
+}
+
+// Barrier arrive
+template <bool return_phase = true>
+MUTE_HOST_DEVICE
+auto                                                   // Phase bit of the barrier if return_phase = true
+arrive_barrier(uint32_t barrier_id)                    // 32 bits user-managed barrier's id
+{
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+  if constexpr (return_phase) {
+    return __musa_async_arrive(barrier_id);
+  } else {
+    __musa_async_arrive(barrier_id);
+  }
+#else
+  if constexpr (return_phase) {
+    return -1;
+  } else {
+    return;
+  }
+#endif
+}
+
+// Barrier wait
+MUTE_HOST_DEVICE
+void
+wait_barrier(uint32_t barrier_id,                      // 32 bits user-managed barrier's id
+             uint32_t phase_bit)                       // Current phase bit the barrier waiting to flip
+{
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+  __musa_async_wait(barrier_id, phase_bit);
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// TME Descriptor and utilities
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace TME {
+
+enum class SmemSwizzleGranularity : uint8_t {
+  NONE = 0,
+  B16  = 1,
+  B32  = 2,
+  B64  = 3,
+};
+
+enum class SmemSwizzleStride : uint8_t {
+  B32  = 0,
+  B64  = 1,
+  B128 = 2,
+  B256 = 3,
+};
+
+enum class SmemSwizzleLine : uint8_t {
+  B128 = 0,
+  B256 = 1,
+};
+
+
+constexpr bool operator<(SmemSwizzleGranularity sg, SmemSwizzleStride ss) {
+  return static_cast<uint8_t>(sg) < (static_cast<uint8_t>(ss) + 2);
+}
+
+constexpr bool operator<=(SmemSwizzleStride ss, SmemSwizzleLine sl) {
+  return static_cast<uint8_t>(ss) <= (static_cast<uint8_t>(sl) + 2);
+}
+
+MUTE_HOST_DEVICE char const* to_string(SmemSwizzleGranularity const& t) {
+  switch (t) {
+    case SmemSwizzleGranularity::NONE: return "NONE";
+    case SmemSwizzleGranularity::B16:  return "SG_16B";
+    case SmemSwizzleGranularity::B32:  return "SG_32B";
+    case SmemSwizzleGranularity::B64:  return "SG_64B";
+  }
+  return nullptr;
+}
+
+MUTE_HOST_DEVICE char const* to_string(SmemSwizzleStride const& t) {
+  switch (t) {
+    case SmemSwizzleStride::B32:  return "SS_32B";
+    case SmemSwizzleStride::B64:  return "SS_64B";
+    case SmemSwizzleStride::B128: return "SS_128B";
+    case SmemSwizzleStride::B256: return "SS_256B";
+  }
+  return nullptr;
+}
+
+MUTE_HOST_DEVICE char const* to_string(SmemSwizzleLine const& t) {
+  switch (t) {
+    case SmemSwizzleLine::B128: return "SL_128B";
+    case SmemSwizzleLine::B256: return "SL_256B";
+  }
+  return nullptr;
+}
+
+
+
+#if !defined(__MUSACC_RTC__)
+MUTE_HOST std::ostream& operator<<(std::ostream& os, SmemSwizzleGranularity const& t) {
+  char const* s = to_string(t);
+  if (s) {
+    std::operator<<(os, s);  // Explicit call to avoid ambiguity
+  } else {
+    os.setstate(std::ios_base::failbit);
+  }
+  return os;
+}
+
+MUTE_HOST std::ostream& operator<<(std::ostream& os, SmemSwizzleStride const& t) {
+  char const* s = to_string(t);
+  if (s) {
+    std::operator<<(os, s);  // Explicit call to avoid ambiguity
+  } else {
+    os.setstate(std::ios_base::failbit);
+  }
+  return os;
+}
+
+MUTE_HOST std::ostream& operator<<(std::ostream& os, SmemSwizzleLine const& t) {
+  char const* s = to_string(t);
+  if (s) {
+    std::operator<<(os, s);  // Explicit call to avoid ambiguity
+  } else {
+    os.setstate(std::ios_base::failbit);
+  }
+  return os;
+}
+
+/// @return The TME descriptor datatype enum corresponding to T.
+template <typename T>
+inline MUtensorDescriptorDataType
+to_MUtensorDescriptorDataType() {
+  if constexpr (is_same_v<T,       int8_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_INT8;     } else
+  if constexpr (is_same_v<T,      uint8_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_UINT8;    } else
+  if constexpr (is_same_v<T, float_e4m3_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_UINT8;    } else
+  if constexpr (is_same_v<T, float_e5m2_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_UINT8;    } else
+  if constexpr (is_same_v<T,     uint16_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_UINT16;   } else
+  if constexpr (is_same_v<T,     uint32_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_UINT32;   } else
+  if constexpr (is_same_v<T,     uint64_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_UINT64;   } else
+  if constexpr (is_same_v<T,      int32_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_INT32;    } else
+  if constexpr (is_same_v<T,      int64_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_INT64;    } else
+  if constexpr (is_same_v<T,       half_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_FLOAT16;  } else
+  if constexpr (is_same_v<T,        float>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_FLOAT32;  } else
+  if constexpr (is_same_v<T,       double>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_FLOAT64;  } else
+  if constexpr (is_same_v<T,   bfloat16_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_BFLOAT16; } else
+  if constexpr (is_same_v<T,   tfloat32_t>) { return MU_TENSOR_DESCRIPTOR_DATA_TYPE_TFLOAT32; } else
+  { static_assert(sizeof(T) < 0, "Unknown TME Format!"); }
+}
+#endif // !defined(__MUSACC_RTC__)
+
+} // namespace TME
+
+using TmeDescriptor = MUtensorDescriptor;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// Initiates a TensorDescriptor Prefetch
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+MUTE_HOST_DEVICE
+void
+prefetch_tme_descriptor(TmeDescriptor const* desc_ptr)
+{
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 310)
+#if (__MUSA_ARCH__ == 310)
+  prefetch(desc_ptr);
+#endif
+#else
+  MUTE_INVALID_CONTROL_PATH("Trying to use TME Descriptor Prefetch on invalid arch");
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// Robust Buffer Access
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+using RobustReg = uint64_t __attribute__((vector_size(16)));
+
+struct RobustDescriptor {
+  RobustReg reg;
+  uint64_t oob_addr;
+};
+
+template <class T>
+MUTE_HOST_DEVICE constexpr
+RobustDescriptor
+make_robust_desc(T const* ptr, size_t elements) {
+  uint64_t base_addr = reinterpret_cast<uint64_t>(ptr);
+  uint64_t buff_size = static_cast<uint64_t>(sizeof(T) * elements);
+  uint64_t oob_addr  = base_addr + buff_size;
+
+  RobustReg reg {base_addr, buff_size};
+
+  return {reg, oob_addr};
+}
+
+} //namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31_tme.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31_tme.hpp
new file mode 100644
index 000000000..7b00e087e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/copy_mp31_tme.hpp
@@ -0,0 +1,684 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/arch/copy.hpp>
+#include <mute/arch/copy_mp31.hpp>
+
+namespace mute
+{
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// TME_LOAD : Initiates a TME copy, in tile mode, from global memory to shared memory
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP31_TME_LOAD {};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct MP31_TME_LOAD_1D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& dim0)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    __musa_tme_ld_tile_1d(bar_id, make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim0, crd0,
+                          static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl), static_cast<int32_t>(prefetch));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+
+  struct PREFETCH
+  {
+    MUTE_HOST_DEVICE static void
+    copy(void const* desc_ptr,
+         int32_t const& crd0,
+         int32_t const& dim0)
+    {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+      uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+      __musa_tme_tile_prefetch_only_1d(gmem_int_desc, dim0, crd0);
+#else
+      MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+    }
+  };
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct MP31_TME_LOAD_2D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& dim0,
+       int32_t const& dim1)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    mute::v2i32_t crd {crd0, crd1};
+    mute::v2i32_t dim {dim0, dim1};
+    __musa_tme_ld_tile_2d(bar_id, make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim, crd,
+                          static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl), static_cast<int32_t>(prefetch));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+
+  struct PREFETCH
+  {
+    MUTE_HOST_DEVICE static void
+    copy(void const* desc_ptr,
+         int32_t const& crd0,
+         int32_t const& crd1,
+         int32_t const& dim0,
+         int32_t const& dim1)
+    {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+      uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+      mute::v2i32_t crd {crd0, crd1};
+      mute::v2i32_t dim {dim0, dim1};
+      __musa_tme_tile_prefetch_only_2d(gmem_int_desc, dim, crd);
+#else
+      MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+    }
+  };
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct MP31_TME_LOAD_3D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    mute::v3i32_t crd {crd0, crd1, crd2};
+    mute::v3i32_t dim {dim0, dim1, dim2};
+    __musa_tme_ld_tile_3d(bar_id, make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim, crd,
+                          static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl), static_cast<int32_t>(prefetch));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+
+  struct PREFETCH
+  {
+    MUTE_HOST_DEVICE static void
+    copy(void const* desc_ptr,
+         int32_t const& crd0,
+         int32_t const& crd1,
+         int32_t const& crd2,
+         int32_t const& dim0,
+         int32_t const& dim1,
+         int32_t const& dim2)
+    {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+      uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+      mute::v3i32_t crd {crd0, crd1, crd2};
+      mute::v3i32_t dim {dim0, dim1, dim2};
+      __musa_tme_tile_prefetch_only_3d(gmem_int_desc, dim, crd);
+#else
+      MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+    }
+  };
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct MP31_TME_LOAD_4D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& crd3,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2,
+       int32_t const& dim3)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    mute::v4i32_t crd {crd0, crd1, crd2, crd3};
+    mute::v4i32_t dim {dim0, dim1, dim2, dim3};
+    __musa_tme_ld_tile_4d(bar_id, make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim, crd,
+                          static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl), static_cast<int32_t>(prefetch));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+
+  struct PREFETCH
+  {
+    MUTE_HOST_DEVICE static void
+    copy(void const* desc_ptr,
+         int32_t const& crd0,
+         int32_t const& crd1,
+         int32_t const& crd2,
+         int32_t const& crd3,
+         int32_t const& dim0,
+         int32_t const& dim1,
+         int32_t const& dim2,
+         int32_t const& dim3)
+    {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+      uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+      mute::v4i32_t crd {crd0, crd1, crd2, crd3};
+      mute::v4i32_t dim {dim0, dim1, dim2, dim3};
+      __musa_tme_tile_prefetch_only_4d(gmem_int_desc, dim, crd);
+#else
+      MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+    }
+  };
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct MP31_TME_LOAD_5D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& crd3,
+       int32_t const& crd4,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2,
+       int32_t const& dim3,
+       int32_t const& dim4)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    mute::v5i32_t crd {crd0, crd1, crd2, crd3, crd4};
+    mute::v5i32_t dim {dim0, dim1, dim2, dim3, dim4};
+    __musa_tme_ld_tile_5d(bar_id, make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim, crd,
+                          static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl), static_cast<int32_t>(prefetch));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+
+  struct PREFETCH
+  {
+    MUTE_HOST_DEVICE static void
+    copy(void const* desc_ptr,
+         int32_t const& crd0,
+         int32_t const& crd1,
+         int32_t const& crd2,
+         int32_t const& crd3,
+         int32_t const& crd4,
+         int32_t const& dim0,
+         int32_t const& dim1,
+         int32_t const& dim2,
+         int32_t const& dim3,
+         int32_t const& dim4)
+    {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+      uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+      mute::v5i32_t crd {crd0, crd1, crd2, crd3, crd4};
+      mute::v5i32_t dim {dim0, dim1, dim2, dim3, dim4};
+      __musa_tme_tile_prefetch_only_5d(gmem_int_desc, dim, crd);
+#else
+      MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+    }
+  };
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct MP31_TME_LOAD_ENTRY
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& dim0)
+  {
+    return MP31_TME_LOAD_1D<sg, ss, sl, prefetch>::copy(desc_ptr, bar_id, smem_ptr, crd0, dim0);
+  }
+
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& dim0,
+       int32_t const& dim1)
+  {
+    return MP31_TME_LOAD_2D<sg, ss, sl, prefetch>::copy(desc_ptr, bar_id, smem_ptr, crd0, crd1, dim0, dim1);
+  }
+
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2)
+  {
+    return MP31_TME_LOAD_3D<sg, ss, sl, prefetch>::copy(desc_ptr, bar_id, smem_ptr, crd0, crd1, crd2, dim0, dim1, dim2);
+  }
+
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& crd3,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2,
+       int32_t const& dim3)
+  {
+    return MP31_TME_LOAD_4D<sg, ss, sl, prefetch>::copy(desc_ptr, bar_id, smem_ptr, crd0, crd1, crd2, crd3, dim0, dim1, dim2, dim3);
+  }
+
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr, uint32_t const& bar_id,
+       void      * smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& crd3,
+       int32_t const& crd4,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2,
+       int32_t const& dim3,
+       int32_t const& dim4)
+  {
+    return MP31_TME_LOAD_5D<sg, ss, sl, prefetch>::copy(desc_ptr, bar_id, smem_ptr, crd0, crd1, crd2, crd3, crd4, dim0, dim1, dim2, dim3, dim4);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// TME_STORE : Initiates a TME store, in tile mode, from shared memory to global memory
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP31_TME_STORE {};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl
+>
+struct MP31_TME_STORE_1D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& dim0)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    __musa_tme_st_1d(make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim0, crd0,
+                     static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl
+>
+struct MP31_TME_STORE_2D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& dim0,
+       int32_t const& dim1)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    mute::v2i32_t crd {crd0, crd1};
+    mute::v2i32_t dim {dim0, dim1};
+    __musa_tme_st_2d(make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim, crd,
+                     static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl
+>
+struct MP31_TME_STORE_3D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    mute::v3i32_t crd {crd0, crd1, crd2};
+    mute::v3i32_t dim {dim0, dim1, dim2};
+    __musa_tme_st_3d(make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim, crd,
+                     static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl
+>
+struct MP31_TME_STORE_4D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& crd3,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2,
+       int32_t const& dim3)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    mute::v4i32_t crd {crd0, crd1, crd2, crd3};
+    mute::v4i32_t dim {dim0, dim1, dim2, dim3};
+    __musa_tme_st_4d(make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim, crd,
+                     static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl
+>
+struct MP31_TME_STORE_5D
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& crd3,
+       int32_t const& crd4,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2,
+       int32_t const& dim3,
+       int32_t const& dim4)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    uint64_t gmem_int_desc = reinterpret_cast<uint64_t>(desc_ptr);
+    uint32_t smem_int_ptr  = cast_smem_ptr_to_uint(smem_ptr);
+    mute::v5i32_t crd {crd0, crd1, crd2, crd3, crd4};
+    mute::v5i32_t dim {dim0, dim1, dim2, dim3, dim4};
+    __musa_tme_st_5d(make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr), gmem_int_desc, dim, crd,
+                     static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+};
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl
+>
+struct MP31_TME_STORE_ENTRY
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& dim0)
+  {
+    return MP31_TME_STORE_1D<sg, ss, sl>::copy(desc_ptr, smem_ptr, crd0, dim0);
+  }
+
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& dim0,
+       int32_t const& dim1)
+  {
+    return MP31_TME_STORE_2D<sg, ss, sl>::copy(desc_ptr, smem_ptr, crd0, crd1, dim0, dim1);
+  }
+
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2)
+  {
+    return MP31_TME_STORE_3D<sg, ss, sl>::copy(desc_ptr, smem_ptr, crd0, crd1, crd2, dim0, dim1, dim2);
+  }
+
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& crd3,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2,
+       int32_t const& dim3)
+  {
+    return MP31_TME_STORE_4D<sg, ss, sl>::copy(desc_ptr, smem_ptr, crd0, crd1, crd2, crd3, dim0, dim1, dim2, dim3);
+  }
+
+  MUTE_HOST_DEVICE static void
+  copy(void const* desc_ptr,
+       void const* smem_ptr,
+       int32_t const& crd0,
+       int32_t const& crd1,
+       int32_t const& crd2,
+       int32_t const& crd3,
+       int32_t const& crd4,
+       int32_t const& dim0,
+       int32_t const& dim1,
+       int32_t const& dim2,
+       int32_t const& dim3,
+       int32_t const& dim4)
+  {
+    return MP31_TME_STORE_5D<sg, ss, sl>::copy(desc_ptr, smem_ptr, crd0, crd1, crd2, crd3, crd4, dim0, dim1, dim2, dim3, dim4);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// BLK_COPY : Copy a bulk of memory between shared memory and global memory
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP31_BLK_COPY_G2S
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* gmem_ptr, int32_t  const& bar_id,
+       void      * smem_ptr, uint32_t const& load_bytes)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    constexpr TME::SmemSwizzleGranularity sg = TME::SmemSwizzleGranularity::NONE;
+    constexpr TME::SmemSwizzleStride      ss = TME::SmemSwizzleStride::B256;
+    constexpr TME::SmemSwizzleLine        sl = TME::SmemSwizzleLine::B256;
+    constexpr PrefetchSize          prefetch = PrefetchSize::B128;
+
+    uint64_t gmem_int_ptr = reinterpret_cast<uint64_t>(gmem_ptr);
+    uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
+    __musa_tme_ld_blk(make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr),
+                      make_ptr_with_address_space<AddressSpace::Global>(gmem_int_ptr),
+                      load_bytes, bar_id,
+                      static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl), static_cast<int32_t>(prefetch));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+};
+
+struct MP31_BLK_COPY_S2G
+{
+  MUTE_HOST_DEVICE static void
+  copy(void const* smem_ptr,
+       void      * gmem_ptr, uint32_t const& store_bytes)
+  {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+    constexpr TME::SmemSwizzleGranularity sg = TME::SmemSwizzleGranularity::NONE;
+    constexpr TME::SmemSwizzleStride      ss = TME::SmemSwizzleStride::B256;
+    constexpr TME::SmemSwizzleLine        sl = TME::SmemSwizzleLine::B256;
+
+    uint64_t gmem_int_ptr = reinterpret_cast<uint64_t>(gmem_ptr);
+    uint32_t smem_int_ptr = cast_smem_ptr_to_uint(smem_ptr);
+    __musa_tme_st_blk(make_ptr_with_address_space<AddressSpace::Shared>(smem_int_ptr),
+                      make_ptr_with_address_space<AddressSpace::Global>(gmem_int_ptr),
+                      store_bytes,
+                      static_cast<int32_t>(sg), static_cast<int32_t>(ss), static_cast<int32_t>(sl));
+#else
+    MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+  }
+};
+
+struct MP31_BLK_COPY_AUTO {};
+
+// Indicate arrival of warp issuing TME_STORE
+MUTE_HOST_DEVICE static void
+tme_store_arrive() {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+  __musa_tme_store_commit();
+#else
+  MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+}
+
+// Wait until at most Count committed TME_STOREs are pending and all prior commits are complete
+MUTLASS_HOST_DEVICE static void
+tme_store_wait() {
+#if defined(MUTE_ARCH_TME_MP31_ACTIVATED)
+  __musa_tme_store_read_wait();
+#else
+  MUTE_INVALID_CONTROL_PATH("Trying to use tme without MUTE_ARCH_TME_MP31_ACTIVATED.");
+#endif
+}
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/mma.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/mma.hpp
new file mode 100644
index 000000000..a0fa652e0
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/mma.hpp
@@ -0,0 +1,65 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>           // MUTE_HOST_DEVICE
+#include <mute/numeric/complex.hpp>  // mute::fma
+#include <mute/numeric/real.hpp>     // mute::fma
+
+namespace mute
+{
+
+//
+// Direct FMA for any type
+//
+
+template <class D, class A = D, class B = A, class C = D>
+struct UniversalFMA
+{
+  using DRegisters = D[1];
+  using ARegisters = A[1];
+  using BRegisters = B[1];
+  using CRegisters = C[1];
+
+  MUTE_HOST_DEVICE static constexpr void
+  fma(D      & d,
+      A const& a,
+      B const& b,
+      C const& c)
+  {
+    // Forward to an ADL/mute free function for these types
+    using mute::fma;
+    fma(d, a, b, c);
+  }
+};
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp22.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp22.hpp
new file mode 100644
index 000000000..f86fc0d0c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp22.hpp
@@ -0,0 +1,428 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/arch/mma.hpp>
+
+
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 220))
+#define MUTE_ARCH_MMA_MP22_ENABLED
+#endif
+
+
+namespace mute {
+
+//
+// MP22 HMMA 32x32x16
+//
+
+struct MP22_32x32x16_F32F16F16F32_TT
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_fmma_m32n32k16_mma(d, a, b, c, 0, 0, 0, 0, 1, 0);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x16_F32F16F16F32_TT without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x16_F32F16F16F32_TN
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_fmma_m32n32k16_mma(d, a, b, c, 0, 0, 0, 0, 1, 1);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x16_F32F16F16F32_TN without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x16_F32F16F16F32_NT
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_fmma_m32n32k16_mma(d, a, b, c, 0, 0, 0, 0, 1, 2);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x16_F32F16F16F32_NT without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x16_F32F16F16F32_NN
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_fmma_m32n32k16_mma(d, a, b, c, 0, 0, 0, 0, 1, 3);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x16_F32F16F16F32_NN without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// MP22 BFMMA 32x32x16
+//
+
+struct MP22_32x32x16_F32BF16BF16F32_TT
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_bfmma_m32n32k16_mma(d, a, b, c, 0, 0, 0, 0, 1, 0);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x16_F32BF16BF16F32_TT without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x16_F32BF16BF16F32_TN
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_bfmma_m32n32k16_mma(d, a, b, c, 0, 0, 0, 0, 1, 1);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x16_F32BF16BF16F32_TN without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x16_F32BF16BF16F32_NT
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_bfmma_m32n32k16_mma(d, a, b, c, 0, 0, 0, 0, 1, 2);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x16_F32BF16BF16F32_NT without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x16_F32BF16BF16F32_NN
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_bfmma_m32n32k16_mma(d, a, b, c, 0, 0, 0, 0, 1, 3);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x16_F32BF16BF16F32_NN without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// MP22 TFMMA 32x32x8
+//
+
+struct MP22_32x32x8_F32TF32TF32F32_TT
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_tfmma_m32n32k8_mma(d, a, b, c, 0, 0, 0, 0, 1, 0);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x8_F32TF32TF32F32_TT without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x8_F32TF32TF32F32_TN
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_tfmma_m32n32k8_mma(d, a, b, c, 0, 0, 0, 0, 1, 1);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x8_F32TF32TF32F32_TN without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x8_F32TF32TF32F32_NT
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_tfmma_m32n32k8_mma(d, a, b, c, 0, 0, 0, 0, 1, 2);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x8_F32TF32TF32F32_NT without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x8_F32TF32TF32F32_NN
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_tfmma_m32n32k8_mma(d, a, b, c, 0, 0, 0, 0, 1, 3);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x8_F32TF32TF32F32_NN without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// MP22 IMMA 32x32x32
+//
+
+struct MP22_32x32x32_S32S8S8S32_TT
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_imma_m32n32k32_mma(d, a, b, c, 0, 0, 0, 0, 1, 0);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x32_S32S8S8S32_TT without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x32_S32S8S8S32_TN
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_imma_m32n32k32_mma(d, a, b, c, 0, 0, 0, 0, 1, 1);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x32_S32S8S8S32_TN without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x32_S32S8S8S32_NT
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_imma_m32n32k32_mma(d, a, b, c, 0, 0, 0, 0, 1, 2);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x32_S32S8S8S32_NT without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct MP22_32x32x32_S32S8S8S32_NN
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP22_ENABLED)
+    __musa_imma_m32n32k32_mma(d, a, b, c, 0, 0, 0, 0, 1, 3);
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP22_32x32x32_S32S8S8S32_NN without MUTE_ARCH_MMA_MP22_ENABLED");
+#endif
+  }
+};
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31.hpp
new file mode 100644
index 000000000..b08748e04
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31.hpp
@@ -0,0 +1,4541 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/config.hpp>
+#include <mute/arch/mma.hpp>
+#include <mute/arch/tce_desc.hpp>
+
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 310))
+#define MUTE_ARCH_MMA_MP31_ENABLED
+#endif
+
+namespace mute {
+
+namespace MP31 {
+
+// TCE Src Operand enum
+enum class TceABDtype {
+  U8U8     = 0x00,
+  S8S8     = 0x01,
+  F16F16   = 0x02,
+  BF16BF16 = 0x03,
+  TF32TF32 = 0x04,
+  E4M3E4M3 = 0x05,
+  E5M2E5M2 = 0x06,
+  F16S4    = 0x07,
+  F16S8    = 0x08,
+  S4F16    = 0x09,
+  S8F16    = 0x0A,
+  E4M3E5M2 = 0x0B,
+  E5M2E4M3 = 0x0C,
+  BF16S4   = 0x0D,
+  BF16S8   = 0x0E,
+  S4BF16   = 0x0F,
+  S8BF16   = 0x10,
+};
+
+} // namespace MP31
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x8_F32F16F16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k8_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x8_F32F16F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x8_F32BF16BF16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k8_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x8_F32BF16BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32F16F16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32F16F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32BF16BF16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32BF16BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32F16F16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32F16F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32BF16BF16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32BF16BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32F16F16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32F16F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32BF16BF16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32BF16BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32F16F16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32F16F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32BF16BF16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32BF16BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_S32S8S8S32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 1, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_S32S8S8S32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_U32U8U8U32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_U32U8U8U32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32E5M2E5M2F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32E5M2E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32E4M3E4M3F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32E4M3E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32E4M3E5M2F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32E4M3E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32E5M2E4M3F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32E5M2E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_S32S8S8S32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 1, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_S32S8S8S32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_U32U8U8U32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_U32U8U8U32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32E5M2E5M2F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32E5M2E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32E4M3E4M3F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32E4M3E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32E4M3E5M2F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32E4M3E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32E5M2E4M3F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32E5M2E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_S32S8S8S32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 1, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_S32S8S8S32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_U32U8U8U32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_U32U8U8U32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32E5M2E5M2F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32E5M2E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32E4M3E4M3F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32E4M3E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32E4M3E5M2F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32E4M3E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32E5M2E4M3F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32E5M2E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_S32S8S8S32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 1, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_S32S8S8S32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_U32U8U8U32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_U32U8U8U32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32E5M2E5M2F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32E5M2E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32E4M3E4M3F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32E4M3E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32E4M3E5M2F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32E4M3E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32E5M2E4M3F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32E5M2E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x64_S32S8S8S32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k64_mma(d, a, b, c, 0, 0, 0, 0, 1, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x64_S32S8S8S32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x64_U32U8U8U32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k64_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x64_U32U8U8U32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x64_F32E5M2E5M2F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k64_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x64_F32E5M2E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x64_F32E4M3E4M3F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k64_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x64_F32E4M3E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x64_F32E4M3E5M2F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k64_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x64_F32E4M3E5M2F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x64_F32E5M2E4M3F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k64_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x64_F32E5M2E4M3F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x4_F32TF32TF32F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k4_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x4_F32TF32TF32F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x8_F32TF32TF32F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k8_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x8_F32TF32TF32F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32TF32TF32F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32TF32TF32F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32F16S4F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32F16S4F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32S4F16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32S4F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32BF16S4F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32BF16S4F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32S4BF16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32S4BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32F16S8F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32F16S8F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32S8F16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32S8F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32BF16S8F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32BF16S8F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x8x16_F32S8BF16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n8k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x8x16_F32S8BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32F16S8F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32F16S8F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32S8F16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32S8F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32BF16S8F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32BF16S8F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_8x16x16_F32S8BF16F32
+{
+  using DRegisters = int32_t[4];
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[4];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m8n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_8x16x16_F32S8BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32F16S8F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32F16S8F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32S8F16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32S8F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32BF16S8F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[2];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32BF16S8F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x16_F32S8BF16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[2];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k16_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x16_F32S8BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32F16S8F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32F16S8F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32S8F16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32S8F16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32BF16S8F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[8];
+  using BRegisters = int32_t[4];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32BF16S8F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x16x32_F32S8BF16F32
+{
+  using DRegisters = int32_t[8];
+  using ARegisters = int32_t[4];
+  using BRegisters = int32_t[8];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t      * d,
+      int32_t const* a,
+      int32_t const* b,
+      int32_t const* c)
+  {
+#if defined(MUTE_ARCH_MMA_MP31_ENABLED)
+    constexpr int layout_num = ((static_cast<int>(tnspA) << 1) | static_cast<int>(tnspB)) ^ 0b01;
+    __musa_wmma_m16n16k32_mma(d, a, b, c, 0, 0, 0, 0, 0, layout_num, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x16x32_F32S8BF16F32 without MUTE_ARCH_MMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mute
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include <mute/arch/mma_mp31_desc.hpp>
+#include <mute/arch/mma_mp31_sqmma.hpp>
+
+namespace mute {
+namespace MP31::SQMMA {
+
+template <
+  class ElementA,
+  class ElementB,
+  class ElementC,
+  class TileShape_MNK,
+  TCE::Major MajorA = TCE::Major::K,
+  TCE::Major MajorB = TCE::Major::K,
+  class MaxInstructionM = _128,
+  class MaxInstructionN = _128,
+  auto... Args                         // e.g. SQMMA::ScaleOut::One, [SQMMA::ScaleIn::One, SQMMA::ScaleIn::One]
+                                       // But most commonly leave empty for defaults
+>
+MUTE_HOST_DEVICE constexpr
+auto
+ss_op_selector()
+{
+  static_assert(is_static<TileShape_MNK>::value, "TileShape_MNK must be static.");
+  static_assert(rank(TileShape_MNK{}) == 3, "TileShape_MNK must be rank 3.");
+  auto Tile_M = size<0>(TileShape_MNK{});
+  auto Tile_N = size<1>(TileShape_MNK{});
+  auto Tile_K = size<2>(TileShape_MNK{});
+  auto MaxInstM = MaxInstructionM{};
+  auto MaxInstN = MaxInstructionN{};
+  static_assert(MaxInstM >= 16 && MaxInstM <= 128 && (MaxInstM & (MaxInstM-1)) == 0);
+  static_assert(MaxInstN >= 16 && MaxInstN <= 128 && (MaxInstN & (MaxInstN-1)) == 0);
+
+  // Input A: uint8_t, Input B: uint8_t
+  if constexpr (is_same_v<ElementA, uint8_t> && is_same_v<ElementB, uint8_t>) {
+    static_assert(is_same_v<ElementC, uint32_t>, "ElementC should be uint32_t");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x128x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x64x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x32x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x128x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x64x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x32x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x16x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x128x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x64x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x32x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_16x64x128_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_U32U8U8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: int8_t, Input B: int8_t
+  else if constexpr (is_same_v<ElementA, int8_t> && is_same_v<ElementB, int8_t>) {
+    static_assert(is_same_v<ElementC, int32_t>, "ElementC should be int32_t");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x128x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x64x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x32x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x128x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x64x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x32x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x16x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x128x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x64x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x32x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_16x64x128_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_S32S8S8_SS<MajorA, MajorB>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: half_t, Input B: half_t
+  else if constexpr (is_same_v<ElementA, half_t> && is_same_v<ElementB, half_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_128x128x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_128x64x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_128x32x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x128x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x64x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x32x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x16x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_32x128x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_32x64x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_32x32x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_16x64x16_F32F16F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: bfloat16_t, Input B: bfloat16_t
+  else if constexpr (is_same_v<ElementA, bfloat16_t> && is_same_v<ElementB, bfloat16_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_128x128x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_128x64x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_128x32x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x128x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x64x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x32x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x16x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_32x128x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_32x64x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_32x32x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_16x64x16_F32BF16BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 16 == 0, "Tile_K must be a multiple of 16.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: tfloat32_t, Input B: tfloat32_t
+  else if constexpr (is_same_v<ElementA, tfloat32_t> && is_same_v<ElementB, tfloat32_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr ((Tile_M % 128 == 0)     &&
+                  MajorA == TCE::Major::K &&
+                  MajorB == TCE::Major::K &&
+                  MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_128x128x16_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 8 == 0) {
+          return MP31_128x128x8_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 8 == 0, "Tile_K must be a multiple of 8.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_128x64x16_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 8 == 0) {
+          return MP31_128x64x8_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 8 == 0, "Tile_K must be a multiple of 8.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x64x16_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 8 == 0) {
+          return MP31_64x64x8_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 8 == 0, "Tile_K must be a multiple of 8.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x32x16_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 8 == 0) {
+          return MP31_64x32x8_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 8 == 0, "Tile_K must be a multiple of 8.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_64x16x16_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 8 == 0) {
+          return MP31_64x16x8_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 8 == 0, "Tile_K must be a multiple of 8.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_32x64x16_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 8 == 0) {
+          return MP31_32x64x8_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 8 == 0, "Tile_K must be a multiple of 8.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_32x32x16_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 8 == 0) {
+          return MP31_32x32x8_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 8 == 0, "Tile_K must be a multiple of 8.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 16 == 0) {
+          return MP31_16x64x16_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 8 == 0) {
+          return MP31_16x64x8_F32TF32TF32_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 8 == 0, "Tile_K must be a multiple of 8.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: float_e4m3_t, Input B: float_e4m3_t
+  else if constexpr (is_same_v<ElementA, float_e4m3_t> && is_same_v<ElementB, float_e4m3_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x128x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x64x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x32x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x128x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x64x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x32x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x16x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x128x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x64x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x32x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_16x64x128_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32E4M3E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: float_e5m2_t, Input B: float_e5m2_t
+  else if constexpr (is_same_v<ElementA, float_e5m2_t> && is_same_v<ElementB, float_e5m2_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x128x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x64x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x32x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x128x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x64x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x32x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x16x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x128x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x64x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x32x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_16x64x128_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32E5M2E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: half_t, Input B: int4b_t
+  else if constexpr (is_same_v<ElementA, half_t> && is_same_v<ElementB, int4b_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32F16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: half_t, Input B: int8_t
+  else if constexpr (is_same_v<ElementA, half_t> && is_same_v<ElementB, int8_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32F16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: int4b_t, Input B: half_t
+  else if constexpr (is_same_v<ElementA, int4b_t> && is_same_v<ElementB, half_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32S4F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 32 == 0, "Tile_M must be a multiple of 32.");
+    }
+  }
+  // Input A: int8_t, Input B: half_t
+  else if constexpr (is_same_v<ElementA, int8_t> && is_same_v<ElementB, half_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32S8F16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 32 == 0, "Tile_M must be a multiple of 32.");
+    }
+  }
+  // Input A: float_e4m3_t, Input B: float_e5m2_t
+  else if constexpr (is_same_v<ElementA, float_e4m3_t> && is_same_v<ElementB, float_e5m2_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x128x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x64x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x32x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x128x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x64x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x32x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x16x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x128x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x64x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x32x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_16x64x128_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32E4M3E5M2_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: float_e5m2_t, Input B: float_e4m3_t
+  else if constexpr (is_same_v<ElementA, float_e5m2_t> && is_same_v<ElementB, float_e4m3_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x128x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x64x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_128x32x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x128x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x64x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x32x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_64x16x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x128x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x64x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_32x32x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 128 == 0) {
+          return MP31_16x64x128_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32E5M2E4M3_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: bfloat16_t, Input B: int4b_t
+  else if constexpr (is_same_v<ElementA, bfloat16_t> && is_same_v<ElementB, int4b_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32BF16S4_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: bfloat16_t, Input B: int8_t
+  else if constexpr (is_same_v<ElementA, bfloat16_t> && is_same_v<ElementB, int8_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 16 == 0 && MaxInstM >= 16) {
+      if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_16x64x64_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_16x64x32_F32BF16S8_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 64 == 0, "Tile_N must be a multiple of 64.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 16 == 0, "Tile_M must be a multiple of 16.");
+    }
+  }
+  // Input A: int4b_t, Input B: bfloat16_t
+  else if constexpr (is_same_v<ElementA, int4b_t> && is_same_v<ElementB, bfloat16_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32S4BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 32 == 0, "Tile_M must be a multiple of 32.");
+    }
+  }
+  // Input A: int8_t, Input B: bfloat16_t
+  else if constexpr (is_same_v<ElementA, int8_t> && is_same_v<ElementB, bfloat16_t>) {
+    static_assert(is_same_v<ElementC, float>, "ElementC should be float");
+    if constexpr (Tile_M % 128 == 0 && MaxInstM >= 128) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x128x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x128x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x64x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x64x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_128x32x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_128x32x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else if constexpr (Tile_M % 64 == 0 && MaxInstM >= 64) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x128x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x128x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x64x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x64x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x32x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x32x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 16 == 0 && MaxInstN >= 16) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_64x16x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_64x16x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 16 == 0, "Tile_N must be a multiple of 16.");
+      }
+    }
+    else if constexpr (Tile_M % 32 == 0 && MaxInstM >= 32) {
+      if constexpr (Tile_N % 128 == 0 && MaxInstN >= 128) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x128x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x128x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 64 == 0 && MaxInstN >= 64) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x64x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x64x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else if constexpr (Tile_N % 32 == 0 && MaxInstN >= 32) {
+        if constexpr (Tile_K % 64 == 0) {
+          return MP31_32x32x64_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else if constexpr (Tile_K % 32 == 0) {
+          return MP31_32x32x32_F32S8BF16_SS<MajorA, MajorB, Args...>{};
+        }
+        else {
+          static_assert(Tile_K % 32 == 0, "Tile_K must be a multiple of 32.");
+        }
+      }
+      else {
+        static_assert(Tile_N % 32 == 0, "Tile_N must be a multiple of 32.");
+      }
+    }
+    else {
+      static_assert(Tile_M % 32 == 0, "Tile_M must be a multiple of 32.");
+    }
+  }
+  else {
+    static_assert(sizeof(ElementA) == 0, "Unknown TCE AB Dtype");
+  }
+}
+
+} // namespace MP31::SQMMA
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31_desc.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31_desc.hpp
new file mode 100644
index 000000000..11db16867
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31_desc.hpp
@@ -0,0 +1,175 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/arch/mma.hpp>
+
+
+namespace mute {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// SQMMA Descriptor and utilities
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA enums and utilities
+namespace MP31::SQMMA
+{
+
+enum class LeadingStrideType : uint8_t {
+  B8   = 1,
+  B16  = 2,
+  B32  = 3,
+  B64  = 4,
+  B128 = 5,
+  B256 = 6,
+};
+
+enum class SwizzleGranularityType : uint8_t {
+  SWIZZLE_GRANULARITY_NONE = 0,
+  SWIZZLE_GRANULARITY_B16  = 1,
+  SWIZZLE_GRANULARITY_B32  = 2,
+  SWIZZLE_GRANULARITY_B64  = 3,
+};
+
+
+MUTE_HOST_DEVICE char const* to_string(LeadingStrideType const& t) {
+  switch (t) {
+    case LeadingStrideType::B8:    return "8B";
+    case LeadingStrideType::B16:   return "16B";
+    case LeadingStrideType::B32:   return "32B";
+    case LeadingStrideType::B64:   return "64B";
+    case LeadingStrideType::B128:  return "128B";
+    case LeadingStrideType::B256:  return "256B";
+  }
+  return nullptr;
+}
+
+MUTE_HOST_DEVICE char const* to_string(SwizzleGranularityType const& t) {
+  switch (t) {
+    case SwizzleGranularityType::SWIZZLE_GRANULARITY_NONE: return "NONE";
+    case SwizzleGranularityType::SWIZZLE_GRANULARITY_B16:  return "SG_16B";
+    case SwizzleGranularityType::SWIZZLE_GRANULARITY_B32:  return "SG_32B";
+    case SwizzleGranularityType::SWIZZLE_GRANULARITY_B64:  return "SG_64B";
+  }
+  return nullptr;
+}
+
+#if !defined(__MUSACC_RTC__)
+MUTE_HOST std::ostream& operator<<(std::ostream& os, LeadingStrideType const& t) {
+  char const* s = to_string(t);
+  if (s) {
+    std::operator<<(os, s);  // Explicit call to avoid ambiguity
+  } else {
+    os.setstate(std::ios_base::failbit);
+  }
+  return os;
+}
+
+MUTE_HOST std::ostream& operator<<(std::ostream& os, SwizzleGranularityType const& t) {
+  char const* s = to_string(t);
+  if (s) {
+    std::operator<<(os, s);  // Explicit call to avoid ambiguity
+  } else {
+    os.setstate(std::ios_base::failbit);
+  }
+  return os;
+}
+#endif // !defined(__MUSACC_RTC__)
+
+
+union SqmmaDescriptor
+{
+
+  MUTE_HOST_DEVICE constexpr
+  SqmmaDescriptor() noexcept : desc_(0) {}
+  MUTE_HOST_DEVICE constexpr
+  SqmmaDescriptor(uint32_t desc) noexcept : desc_(desc) {}
+  MUTE_HOST_DEVICE constexpr
+  SqmmaDescriptor(SqmmaDescriptor const& t) noexcept : desc_(t.desc_) {}
+  MUTE_HOST_DEVICE constexpr
+  SqmmaDescriptor(SqmmaDescriptor && t) noexcept : desc_(t.desc_) {}
+
+  MUTE_HOST_DEVICE constexpr
+  SqmmaDescriptor& operator=(SqmmaDescriptor const& t) noexcept {
+    desc_ = t.desc_;
+    return *this;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  SqmmaDescriptor& operator=(SqmmaDescriptor && t) noexcept {
+    desc_ = t.desc_;
+    return *this;
+  }
+
+  int32_t desc_;
+
+  // Bitfield implementation avoids the need for shifts in assignment
+  struct BitField {
+    // start_address, bit [0, 18)
+    uint32_t start_address_ : 18;
+    // leading dimension byte offset type, bit [18, 21)
+    // 8B:1,16B:2,32B:3,64B:4,128B:5,256B:6
+    uint8_t leading_stride_type_ : 3;
+    // swizzle granularity type, bit [21, 23)
+    // SG_NONE=0,SG_16B=1,SG_32B=2,SG_64B=3
+    uint8_t swizzle_granularity_type_ : 2;
+  } bitfield;
+
+  static_assert(sizeof_bits_v<BitField> == 32, "The underlying bits of BitField must be 32.");
+
+  // Decay to a int32_t
+  MUTE_HOST_DEVICE constexpr
+  operator int32_t() const noexcept { return desc_; }
+
+  // Printer
+  MUTE_HOST_DEVICE friend void print(SqmmaDescriptor const& t)
+  {
+    #if !defined(__MUSACC_RTC__)
+    printf("SqmmaDescriptor: 0x%08x\n", t.desc_);
+    printf("     start_addr: 0x%05x\n", t.bitfield.start_address_);
+    printf(" leading_stride: 0x%01x(%s)\n", t.bitfield.leading_stride_type_,
+           to_string(static_cast<SQMMA::LeadingStrideType>(t.bitfield.leading_stride_type_)));
+    printf(" sw_granularity: 0x%01x(%s)\n", t.bitfield.swizzle_granularity_type_,
+           to_string(static_cast<SQMMA::SwizzleGranularityType>(t.bitfield.swizzle_granularity_type_)));
+    #endif
+  }
+};
+
+} // namespace MP31::SQMMA
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mute
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31_sqmma.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31_sqmma.hpp
new file mode 100644
index 000000000..7f82288f2
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/mma_mp31_sqmma.hpp
@@ -0,0 +1,14598 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/config.hpp>
+#include <mute/arch/mma.hpp>
+#include <mute/arch/mma_mp31.hpp>
+
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 310))
+#define MUTE_ARCH_SQMMA_MP31_ENABLED
+#endif
+
+namespace mute {
+
+// Wait all previous in-flight SQMMA to complete
+template <int N = 0>
+MUTE_HOST_DEVICE
+void
+warpsquad_wait()
+{
+  static_assert(N == 0, "SQMMA wait: N must be 0 now");
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+  __musa_sqmma_wait();
+#else
+  MUTE_INVALID_CONTROL_PATH("Attempting to use sqmma wait without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+}
+
+
+namespace MP31::SQMMA {
+
+enum class ScaleOut {
+  One  = 0,
+  Zero = 1,
+};
+
+enum class ScaleIn {
+  One = 0,
+  Neg = 1,
+};
+
+} // namespace SQMMA
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// SQMMA Intrinsic definitions:  C = (scaleA * A) * (scaleB * B) + (scaleD * C)
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x64x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x64x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x64x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x32x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x32x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x32x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x64x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x64x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x64x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x128x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x128x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x128x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x16x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x16x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x16x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x32x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x32x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x32x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x64x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x64x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x64x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x128x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x128x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x128x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x32x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x32x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x32x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x64x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x64x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x64x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x128x32_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x128x64_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x128 U32+=U8*U8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x128x128_U32U8U8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::U8U8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x128_U32U8U8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x64x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x64x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_16x64x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x32x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x32x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x32x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x64x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x64x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x64x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x128x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x128x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_32x128x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x16x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x16x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x16x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x32x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x32x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x32x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x64x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x64x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x64x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x128x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x128x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_64x128x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x32x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x32x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x32x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x64x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x64x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x64x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x128x32_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x128x64_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x128 S32+=S8*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB
+>
+struct MP31_128x128x128_S32S8S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x128_S32S8S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x16 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x16_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x16_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=F16*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32F16F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32F16F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x16 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x16_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x16_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=BF16*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32BF16BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32BF16BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x8 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x8_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 8, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x8_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x16 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x16_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x16_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x8 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x8_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 8, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x8_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x16 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x16_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x16_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x8 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x8_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 8, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x8_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x16 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x16_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x16_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x8 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x8_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 8, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x8_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x16 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x16_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x16_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x8 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x8_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 8, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x8_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x16 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x16_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x16_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x8 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x8_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 8, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x8_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x16 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x16_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x16_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x8 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x8_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  static_assert(tnspA == TCE::Major::K,
+      "TF32 SQMMA 128x64 operand A must have K major layout.");
+  static_assert(tnspB == TCE::Major::K,
+      "TF32 SQMMA 128x64 operand B must have K major layout.");
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, 0, 1,
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 8, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x8_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x16 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x16_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  static_assert(tnspA == TCE::Major::K,
+      "TF32 SQMMA 128x64 operand A must have K major layout.");
+  static_assert(tnspB == TCE::Major::K,
+      "TF32 SQMMA 128x64 operand B must have K major layout.");
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, 0, 1,
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x16_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  static_assert(tnspA == TCE::Major::K,
+      "TF32 SQMMA 128x64 operand A must have K major layout.");
+  static_assert(tnspB == TCE::Major::K,
+      "TF32 SQMMA 128x64 operand B must have K major layout.");
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, 0, 1,
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x8 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x8_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  static_assert(tnspA == TCE::Major::K,
+      "TF32 SQMMA 128x128 operand A must have K major layout.");
+  static_assert(tnspB == TCE::Major::K,
+      "TF32 SQMMA 128x128 operand B must have K major layout.");
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, 0, 1,
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 8, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x8_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x16 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x16_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  static_assert(tnspA == TCE::Major::K,
+      "TF32 SQMMA 128x128 operand A must have K major layout.");
+  static_assert(tnspB == TCE::Major::K,
+      "TF32 SQMMA 128x128 operand B must have K major layout.");
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, 0, 1,
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 16, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x16_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=TF32*TF32
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32TF32TF32_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  static_assert(tnspA == TCE::Major::K,
+      "TF32 SQMMA 128x128 operand A must have K major layout.");
+  static_assert(tnspB == TCE::Major::K,
+      "TF32 SQMMA 128x128 operand B must have K major layout.");
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, 0, 1,
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::TF32TF32));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32TF32TF32_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x128 F32+=E4M3*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x128_F32E4M3E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x128_F32E4M3E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x128 F32+=E5M2*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x128_F32E5M2E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x128_F32E5M2E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=F16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32F16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32F16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=F16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32F16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::F16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32F16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=S4*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32S4F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32S4F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=S8*F16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32S8F16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8F16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32S8F16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x128 F32+=E4M3*E5M2
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x128_F32E4M3E5M2_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E4M3E5M2));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x128_F32E4M3E5M2_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x128 F32+=E5M2*E4M3
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x128_F32E5M2E4M3_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 128, /* mma_type = */ static_cast<int>(MP31::TceABDtype::E5M2E4M3));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x128_F32E5M2E4M3_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=BF16*S4
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32BF16S4_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S4));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32BF16S4_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 16x64x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_16x64x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m16n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_16x64x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=BF16*S8
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleA = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32BF16S8_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      static_cast<int32_t>(scaleA), 0, static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::BF16S8));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32BF16S8_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=S4*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32S4BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S4BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32S4BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x32x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x32x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x32x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x64x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x64x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x64x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 32x128x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_32x128x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m32n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_32x128x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x16x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x16x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[8];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n16_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x16x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x32x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x32x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[16];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x32x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x64x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x64x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x64x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 64x128x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_64x128x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m64n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_64x128x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x32x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x32x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[32];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n32_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x32x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x64x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x64x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[64];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n64_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x64x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x32 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x32_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 32, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x32_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA 128x128x64 F32+=S8*BF16
+template <
+  TCE::Major tnspA,
+  TCE::Major tnspB,
+  MP31::SQMMA::ScaleIn scaleB = MP31::SQMMA::ScaleIn::One
+>
+struct MP31_128x128x64_F32S8BF16_SS
+{
+  using DRegisters = void;
+  using ARegisters = int32_t[1];
+  using BRegisters = int32_t[1];
+  using CRegisters = int32_t[128];
+
+  MUTE_HOST_DEVICE static void
+  fma(int32_t const& desc_a,
+      int32_t const& desc_b,
+      int32_t      * d,
+      MP31::SQMMA::ScaleOut scaleD = MP31::SQMMA::ScaleOut::One)
+  {
+#if defined(MUTE_ARCH_SQMMA_MP31_ENABLED)
+    __musa_sqmma_m128n128_mma(
+      d, &desc_a, &desc_b, d, static_cast<int32_t>(tnspA), 1 - static_cast<int32_t>(tnspB),
+      0, static_cast<int32_t>(scaleB), static_cast<int32_t>(scaleD),
+      /* sat = */ 0, /* k = */ 64, /* mma_type = */ static_cast<int>(MP31::TceABDtype::S8BF16));
+#else
+    MUTE_INVALID_CONTROL_PATH("Attempting to use MP31_128x128x64_F32S8BF16_SS without MUTE_ARCH_SQMMA_MP31_ENABLED");
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/tce_desc.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/tce_desc.hpp
new file mode 100644
index 000000000..3414b36c2
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/tce_desc.hpp
@@ -0,0 +1,44 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+namespace mute {
+namespace TCE {
+
+// TCE MMA Major enum
+enum class Major {
+  K  = 0,
+  MN = 1
+};
+
+} // namespace TCE
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/arch/util.hpp b/vllm_musa/csrc/flash_attn/include/mute/arch/util.hpp
new file mode 100644
index 000000000..98248a298
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/arch/util.hpp
@@ -0,0 +1,257 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/numeric/integer_sequence.hpp>
+
+namespace mute
+{
+
+/// MUTE helper to cast SMEM pointer to unsigned
+MUTE_HOST_DEVICE
+uint32_t
+cast_smem_ptr_to_uint(void const* const ptr)
+{
+#if defined(__MUSA_ARCH__)
+  /// MUTE helper to get SMEM pointer
+  return static_cast<uint32_t>(__cvta_generic_to_shared(ptr));
+#else
+  (void) ptr;
+  printf("ERROR: cast_smem_ptr_to_uint not supported but used.\n");
+  return 0;
+#endif
+}
+
+namespace detail {
+
+//
+// Wrapper for MMAOp::fma
+//
+
+template <class MmaOp>
+struct CallFMA {
+  template <class... Args>
+  MUTE_HOST_DEVICE constexpr void
+  operator()(Args&&... args) const {
+    return MmaOp::fma(static_cast<Args&&>(args)...);
+  }
+};
+
+//
+// Wrapper for CopyOp::copy
+//
+
+template <class CopyOp>
+struct CallCOPY {
+  template <class... Args>
+  MUTE_HOST_DEVICE constexpr void
+  operator()(Args&&... args) const {
+    return CopyOp::copy(static_cast<Args&&>(args)...);
+  }
+};
+
+//
+// Utility for exploding pointers/arrays/tensors into functions
+//
+
+template <class Fn,
+          class PtrA, int... I>
+MUTE_HOST_DEVICE constexpr
+void
+explode(Fn fn,
+        PtrA&& a, int_sequence<I...>)
+{
+  return fn(a[I]...);
+}
+
+template <class Fn,
+          class PtrS, int... Is,
+          class PtrD, int... Id>
+MUTE_HOST_DEVICE constexpr
+void
+explode(Fn fn,
+        PtrS&& s, int_sequence<Is...>,
+        PtrD&& d, int_sequence<Id...>)
+{
+  return fn(s[Is]..., d[Id]...);
+}
+
+template <class Fn,
+          class PtrA, int... Ia,
+          class PtrB, int... Ib,
+          class PtrC, int... Ic>
+MUTE_HOST_DEVICE constexpr
+void
+explode(Fn fn,
+        PtrA&& a, int_sequence<Ia...>,
+        PtrB&& b, int_sequence<Ib...>,
+        PtrC&& c, int_sequence<Ic...>)
+{
+  return fn(a[Ia]..., b[Ib]..., c[Ic]...);
+}
+
+template <class Fn,
+          class PtrD, int... Id,
+          class PtrA, int... Ia,
+          class PtrB, int... Ib,
+          class PtrC, int... Ic>
+MUTE_HOST_DEVICE constexpr
+void
+explode(Fn fn,
+        PtrD&& d, int_sequence<Id...>,
+        PtrA&& a, int_sequence<Ia...>,
+        PtrB&& b, int_sequence<Ib...>,
+        PtrC&& c, int_sequence<Ic...>)
+{
+  return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]...);
+}
+
+template <class Fn,
+          class PtrD, int... Id,
+          class PtrA, int... Ia,
+          class PtrB, int... Ib,
+          class PtrC, int... Ic,
+          class PtrE, int... Ie>
+MUTE_HOST_DEVICE constexpr
+void
+explode(Fn fn,
+        PtrD&& d, int_sequence<Id...>,
+        PtrA&& a, int_sequence<Ia...>,
+        PtrB&& b, int_sequence<Ib...>,
+        PtrC&& c, int_sequence<Ic...>,
+        PtrE&& e, int_sequence<Ie...>)
+{
+  return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., e[Ie]...);
+}
+
+template <class Fn,
+          class PtrD,   int... Id,
+          class PtrA,   int... Ia,
+          class PtrB,   int... Ib,
+          class PtrC,   int... Ic,
+          class PtrSFA, int... Isfa,
+          class PtrSFB, int... Isfb>
+MUTE_HOST_DEVICE constexpr
+void
+explode(Fn fn,
+        PtrD&& d,     int_sequence<Id...>,
+        PtrA&& a,     int_sequence<Ia...>,
+        PtrB&& b,     int_sequence<Ib...>,
+        PtrC&& c,     int_sequence<Ic...>,
+        PtrSFA&& sfa, int_sequence<Isfa...>,
+        PtrSFB&& sfb, int_sequence<Isfb...>)
+{
+  return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., sfa[Isfa]..., sfb[Isfb]...);
+}
+
+template <class Fn,
+          class PtrD, int... Id,
+          class PtrA, int... Ia,
+          class PtrB, int... Ib,
+          class PtrC, int... Ic,
+          class PtrE, int... Ie,
+          class PtrF, int... If,
+          class PtrG, int... Ig>
+MUTE_HOST_DEVICE constexpr
+void
+explode(Fn fn,
+        PtrD&& d, int_sequence<Id...>,
+        PtrA&& a, int_sequence<Ia...>,
+        PtrB&& b, int_sequence<Ib...>,
+        PtrC&& c, int_sequence<Ic...>,
+        PtrE&& e, int_sequence<Ie...>,
+        PtrF&& f, int_sequence<If...>,
+        PtrG&& g, int_sequence<Ig...>)
+{
+  return fn(d[Id]..., a[Ia]..., b[Ib]..., c[Ic]..., e[Ie]..., f[If]..., g[Ig]...);
+}
+//
+// Utility for exploding tuples into functions
+//
+
+template <class Fn,
+          class TupleA, int... I>
+MUTE_HOST_DEVICE constexpr
+void
+explode_tuple(Fn fn,
+              TupleA&& a, int_sequence<I...>)
+{
+  return fn(get<I>(a)...);
+}
+
+template <class Fn,
+          class TupleA, int... Ia,
+          class TupleB, int... Ib>
+MUTE_HOST_DEVICE constexpr
+void
+explode_tuple(Fn fn,
+              TupleA&& a, int_sequence<Ia...>,
+              TupleB&& b, int_sequence<Ib...>)
+{
+  return fn(get<Ia>(a)..., get<Ib>(b)...);
+}
+
+template <class Fn,
+          class TupleA, int... Ia,
+          class TupleB, int... Ib,
+          class TupleC, int... Ic>
+MUTE_HOST_DEVICE constexpr
+void
+explode_tuple(Fn fn,
+              TupleA&& a, int_sequence<Ia...>,
+              TupleB&& b, int_sequence<Ib...>,
+              TupleC&& c, int_sequence<Ic...>)
+{
+  return fn(get<Ia>(a)..., get<Ib>(b)..., get<Ic>(c)...);
+}
+
+template <class Fn,
+          class TupleA, int... Ia,
+          class TupleB, int... Ib,
+          class TupleC, int... Ic,
+          class TupleD, int... Id>
+MUTE_HOST_DEVICE constexpr
+void
+explode_tuple(Fn fn,
+              TupleA&& a, int_sequence<Ia...>,
+              TupleB&& b, int_sequence<Ib...>,
+              TupleC&& c, int_sequence<Ic...>,
+              TupleD&& d, int_sequence<Id...>)
+{
+  return fn(get<Ia>(a)..., get<Ib>(b)..., get<Ic>(c)..., get<Id>(d)...);
+}
+
+} // end namespace detail
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/copy_atom.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_atom.hpp
new file mode 100644
index 000000000..6219da0d0
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_atom.hpp
@@ -0,0 +1,753 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE
+#include <mute/tensor.hpp>                     // mute::Tensor
+#include <mute/util/type_traits.hpp>           // mute::__MUTE_REQUIRES
+#include <mute/container/tuple.hpp>            // mute::is_tuple
+#include <mute/numeric/integral_constant.hpp>  // mute::is_constant, mute::is_integral
+#include <mute/atom/copy_traits.hpp>           // mute::Copy_Traits
+#include <mute/atom/mma_atom.hpp>              // mute::TiledMMA
+
+namespace mute
+{
+
+template <class... Args>
+struct Copy_Atom;
+
+template <class CopyOperation, class CopyInternalType>
+struct Copy_Atom<CopyOperation, CopyInternalType> : Copy_Atom<Copy_Traits<CopyOperation>, CopyInternalType>
+{};
+
+template <class... Args, class CopyInternalType>
+struct Copy_Atom<Copy_Traits<Args...>, CopyInternalType>
+  : Copy_Traits<Args...>
+{
+  using Traits = Copy_Traits<Args...>;
+
+  // Bit and Thr layouts from the Copy_Traits
+  using ThrID        = typename Traits::ThrID;
+  using BitLayoutSrc = typename Traits::SrcLayout;
+  using BitLayoutDst = typename Traits::DstLayout;
+  using BitLayoutRef = typename Traits::RefLayout;
+
+  using ValType = CopyInternalType;
+
+  using ValLayoutSrc = decltype(recast_layout<uint1_t, ValType>(BitLayoutSrc{}));
+  using ValLayoutDst = decltype(recast_layout<uint1_t, ValType>(BitLayoutDst{}));
+  using ValLayoutRef = decltype(recast_layout<uint1_t, ValType>(BitLayoutRef{}));
+
+  MUTE_STATIC_ASSERT_V(size<0>(ValLayoutSrc{}) == size(ThrID{}), "CopyOperation is not valid for Src of ValType.");
+  MUTE_STATIC_ASSERT_V(size<0>(ValLayoutDst{}) == size(ThrID{}), "CopyOperation is not valid for Dst of ValType.");
+  MUTE_STATIC_ASSERT_V(size<0>(ValLayoutRef{}) == size(ThrID{}), "CopyOperation is not valid for Ref of ValType.");
+
+  static constexpr int NumValSrc = size<1>(ValLayoutSrc{});
+  static constexpr int NumValDst = size<1>(ValLayoutDst{});
+
+  // Additional Trait parameters/transformations
+  template <class... TraitsArgs>
+  MUTE_HOST_DEVICE
+  auto
+  with(TraitsArgs&&... args) const {
+    auto traits = Traits::with(static_cast<TraitsArgs&&>(args)...);
+    return Copy_Atom<decltype(traits), CopyInternalType>{traits};
+  }
+
+  //
+  // Tensor call interfaces
+  //
+
+  // Check and call instruction, or recurse
+  template <class SEngine, class SLayout,
+            class DEngine, class DLayout>
+  MUTE_HOST_DEVICE
+  void
+  call(Tensor<SEngine,SLayout> const& src,
+       Tensor<DEngine,DLayout>      & dst) const
+  {
+    static_assert(SLayout::rank == 1, "Expected rank-1 src tensor");
+    static_assert(DLayout::rank == 1, "Expected rank-1 dst tensor");
+
+    if constexpr (is_constant<NumValSrc, decltype(size(src))>::value ||
+                  is_constant<NumValDst, decltype(size(dst))>::value) {
+      // Dispatch to unpack to execute instruction
+      return copy_unpack(static_cast<Traits const&>(*this), src, dst);
+    } else if constexpr (is_tuple<decltype(shape(src))>::value &&
+                         is_tuple<decltype(shape(dst))>::value) {
+      // If the size of the src/dst doesn't match the instruction,
+      //   recurse this rank-1 layout by peeling off the mode
+      //   ((A,B,C,...)) -> (A,B,C,...)
+      return copy(*this, tensor<0>(src), tensor<0>(dst));
+    } else {
+      static_assert(dependent_false<SEngine>,
+                    "CopyAtom: Src/Dst partitioning does not match the instruction requirement.");
+    }
+  }
+
+  // Accept mutable temporaries
+  template <class SEngine, class SLayout,
+            class DEngine, class DLayout>
+  MUTE_HOST_DEVICE
+  void
+  call(Tensor<SEngine,SLayout> const& src,
+       Tensor<DEngine,DLayout>     && dst) const
+  {
+    return call(src, dst);
+  }
+};
+
+//
+// A tiling of copy atoms
+//
+
+template <class TiledCopy, class ThrIdx>
+struct ThrCopy;
+
+template <class Copy_Atom,
+          class LayoutCopy_TV,  // (tid,vid) -> coord   [Need not be 2D...]
+          class ShapeTiler_MN>  // coord space
+struct TiledCopy : Copy_Atom
+{
+  // Layout information from the CopyAtom
+  using AtomThrID     = typename Copy_Atom::ThrID;        // thrid -> thr_idx
+  using AtomLayoutSrc = typename Copy_Atom::ValLayoutSrc; // (thr,val) -> offset
+  using AtomLayoutDst = typename Copy_Atom::ValLayoutDst; // (thr,val) -> offset
+  using AtomLayoutRef = typename Copy_Atom::ValLayoutRef; // (thr,val) -> offset
+
+  using AtomNumThr = decltype(size<0>(AtomLayoutRef{}));
+  using AtomNumVal = decltype(size<1>(AtomLayoutRef{}));
+
+  // Layout information for the TiledCopy
+  using Tiler_MN       = ShapeTiler_MN;
+  using TiledLayout_TV = LayoutCopy_TV;
+  using TiledNumThr    = decltype(size<0>(TiledLayout_TV{}));
+  using TiledNumVal    = decltype(size<1>(TiledLayout_TV{}));
+
+  MUTE_STATIC_ASSERT_V(TiledNumThr{} % AtomNumThr{} == Int<0>{}, "TiledCopy uses too few thrs for selected CopyAtom");
+  MUTE_STATIC_ASSERT_V(TiledNumVal{} % AtomNumVal{} == Int<0>{}, "TiledCopy uses too few vals for selected CopyAtom");
+
+  // Tile a tensor or a layout from shape
+  //   (M,N,...)
+  // to shape
+  //   ((ThrV,ThrX),FrgV,(RestM,RestN,...))
+  // where
+  //   ThrV:  The threads local to a COPY_ATOM Src.
+  //   ThrX:  The threads tiled across COPY_ATOMs Src.
+  //   FrgV:  The values local to a COPY_ATOM Src.
+  //   RestM: The values tiled in M.
+  //   RestN: The values tiled in N.
+  template <class STensor>
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  tidfrg_S(STensor&& stensor)
+  {
+    MUTE_STATIC_ASSERT_V(rank(stensor) >= rank(Tiler_MN{}), "Rank of tensor to be partitioned too small.");
+
+    // Tile the stensor and compute the (src-thr, src-val) -> (ref-thr, ref-val) layout
+    return tile2thrfrg(zipped_divide(stensor,Tiler_MN{}), right_inverse(AtomLayoutRef{}).compose(AtomLayoutSrc{}));
+  }
+
+  // Tile a tensor or a layout from shape
+  //   (M,N,...)
+  // to shape
+  //   ((ThrV,ThrX),FrgV,(RestM,RestN,...))
+  // where
+  //   ThrV:  The threads local to a COPY_ATOM Dst.
+  //   ThrX:  The threads tiled across COPY_ATOMs Dst.
+  //   FrgV:  The values local to a COPY_ATOM Dst.
+  //   RestM: The values tiled in M.
+  //   RestN: The values tiled in N.
+  template <class DTensor>
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  tidfrg_D(DTensor&& dtensor)
+  {
+    MUTE_STATIC_ASSERT_V(rank(dtensor) >= rank(Tiler_MN{}), "Rank of tensor to be partitioned too small.");
+
+    // Tile the dtensor and compute the (dst-thr, dst-val) -> (ref-thr, ref-val) layout
+    return tile2thrfrg(zipped_divide(dtensor,Tiler_MN{}), right_inverse(AtomLayoutRef{}).compose(AtomLayoutDst{}));
+  }
+
+  // Tile a tensor or a layout from shape
+  //   ((TileM,TileN,...), (RestM,RestN,...))
+  // to shape
+  //   ((ThrV,ThrX),FrgV,(RestM,RestN,...))
+  template <class Tensor, class Ref2TrgLayout>
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  tile2thrfrg(Tensor&& tensor, Ref2TrgLayout const& ref2trg)
+  {
+    // Take the thrs/vals that the atom is interested in
+    // NOTE: Assumes the AtomNumThr are contiguous and identity within TiledThrID
+    auto atom_layout_TV = zipped_divide(TiledLayout_TV{}, make_shape(AtomNumThr{}, AtomNumVal{}));
+    // ((atom_tid,atom_val),(rest_tid,rest_val)) -> (m,n)
+
+    // Transform to the trg layout
+    auto trg_layout_TV = atom_layout_TV.compose(ref2trg, _);
+    // ((trg_tid,trg_val),(rest_tid,rest_val)) -> (m,n)
+
+    // Transform the thrs mode from thrid to thr_idx
+    // NOTE: Assumes the AtomNumThr are contiguous and identity within TiledThrID
+    auto thrval2mn = coalesce(zip(trg_layout_TV), Shape<_1,Shape<_1,_1>>{});
+    // ((trg_tid,rest_tid),(trg_val,rest_val)) -> (m,n)
+
+    /// ==================
+
+    // Transform the tile mode
+    auto tv_tensor = tensor.compose(thrval2mn, _);
+    // ((thrid,val),(RestM,RestN,...))
+
+    // Unfold and return
+    return tv_tensor(make_coord(_,_), _);
+  }
+
+  // retile_S and retile_D assume they are working with the reference layout -- they are the same
+  template <class Tensor>
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  retile(Tensor&& tensor)
+  {
+    constexpr int R = remove_cvref_t<Tensor>::rank;
+    // Assert that AtomLayoutSrc|Dst is identity so we can skip the Ref transformation
+
+    // Assume the first size<0>(tensor) elements are the first val_ids in TiledLayout_TV.
+    // Then, we only need the shape+layout of those size<0>(tensor) elements in TiledLayout_TV
+    //   and that shape is what we gather from the other modes of tensor
+
+    auto V = size<0>(tensor);
+
+    auto frg_layout_mn = upcast<TiledNumThr{} * V>(right_inverse(TiledLayout_TV{}).with_shape(shape(Tiler_MN{})));
+    // (m,n) -> v_idx -- The shape and order of the V inside of TiledLayout_TV
+
+    auto frg_layout_v = zipped_divide(logical_product(make_layout(V), right_inverse(frg_layout_mn)), make_layout(AtomNumVal{}));
+    // (atom_vals,rest_vals) -> (v,m,n)
+
+    /// =======
+
+    // Tile the tensor for TileFrg
+    auto t_tensor = zipped_divide(tensor, prepend(product_each(shape(frg_layout_mn)), V));
+    // ((TileV,TileM,TileN,...),(1,RestM,RestN,...))
+
+    // Transform the tile mode
+    auto v_tensor = t_tensor.compose(frg_layout_v, _);
+    // ((atom_vals,rest_vals),(1,RM,RN,...))
+
+    // Unfold and return
+    return v_tensor(_, append<R>(Int<0>{},_));
+  }
+
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  get_layoutS_TV()
+  {
+    // (M,N) -> (M,N)
+    auto ref_S = make_layout(make_shape(shape(Tiler_MN{}), Int<1>{}));
+    // (thr_idx,val_idx) -> (M,N)
+    return tile2thrfrg(ref_S, right_inverse(AtomLayoutRef{}).compose(AtomLayoutSrc{}))(_,_,Int<0>{});
+  }
+
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  get_layoutS_MN()
+  {
+    // (thr_idx,val_idx) -> (M,N)
+    auto layoutS_TV = get_layoutS_TV();
+    // (M,K) -> (thr_idx,val_idx)
+    auto layoutS_MK = right_inverse(layoutS_TV).with_shape(shape(Tiler_MN{}));
+
+    // athrid = (v,m,k) -> thr_idx
+    auto thrID_S = make_layout(size<0>(TiledLayout_TV{}));
+
+    return mute::make_tuple(layoutS_MK, thrID_S);
+  }
+
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  get_layoutD_TV()
+  {
+    // (M,N) -> (M,N)
+    auto ref_D = make_layout(make_shape(shape(Tiler_MN{}), Int<1>{}));
+    // (thr_idx,val_idx) -> (M,N)
+    return tile2thrfrg(ref_D, right_inverse(AtomLayoutRef{}).compose(AtomLayoutDst{}))(_,_,Int<0>{});
+  }
+
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  get_layoutD_MN()
+  {
+    // (thr_idx,val_idx) -> (M,N)
+    auto layoutD_TV = get_layoutD_TV();
+    // (M,K) -> (thr_idx,val_idx)
+    auto layoutD_MK = right_inverse(layoutD_TV).with_shape(shape(Tiler_MN{}));
+
+    // athrid = (v,m,k) -> thr_idx
+    auto thrID_D = make_layout(size<0>(TiledLayout_TV{}));
+
+    return mute::make_tuple(layoutD_MK, thrID_D);
+  }
+
+  template <class ThrIdx,
+            __MUTE_REQUIRES(is_integral<ThrIdx>::value)>
+  MUTE_HOST_DEVICE static
+  auto
+  get_slice(ThrIdx const& thr_idx)
+  {
+    return ThrCopy<TiledCopy, ThrIdx>(thr_idx);
+  }
+
+  template <class ThrIdx,
+            __MUTE_REQUIRES(is_integral<ThrIdx>::value)>
+  MUTE_HOST_DEVICE  static
+  auto
+  get_thread_slice(ThrIdx const& thr_idx)
+  {
+    return get_slice(thr_idx);
+  }
+};
+
+template <class TiledCopy, class ThrIdx>
+struct ThrCopy
+{
+  ThrIdx thr_idx_;
+
+  MUTE_HOST_DEVICE
+  ThrCopy(ThrIdx const& thr_idx) : thr_idx_(thr_idx) {}
+
+  template <class STensor>
+  MUTE_HOST_DEVICE
+  auto
+  partition_S(STensor&& stensor) const {
+    //static_assert(sizeof(typename remove_cvref_t<STensor>::value_type) == sizeof(typename TiledCopy::ValType),
+    //              "Expected ValType for tiling SrcTensor.");
+    auto thr_tensor = make_tensor(static_cast<STensor&&>(stensor).data(), TiledCopy::tidfrg_S(stensor.layout()));
+    return thr_tensor(thr_idx_, _, repeat<rank_v<STensor>>(_));
+  }
+
+  template <class DTensor>
+  MUTE_HOST_DEVICE
+  auto
+  partition_D(DTensor&& dtensor) const {
+    //static_assert(sizeof(typename remove_cvref_t<DTensor>::value_type) == sizeof(typename TiledCopy::ValType),
+    //              "Expected ValType for tiling DstTensor.");
+    auto thr_tensor = make_tensor(static_cast<DTensor&&>(dtensor).data(), TiledCopy::tidfrg_D(dtensor.layout()));
+    return thr_tensor(thr_idx_, _, repeat<rank_v<DTensor>>(_));
+  }
+
+  template <class STensor>
+  MUTE_HOST_DEVICE static
+  auto
+  retile_S(STensor&& stensor) {
+    // static_assert(sizeof(typename remove_cvref_t<STensor>::value_type) == sizeof(typename TiledCopy::ValType),
+    //               "Expected ValType for tiling SrcTensor.");
+    return make_tensor(static_cast<STensor&&>(stensor).data(), TiledCopy::retile(stensor.layout()));
+  }
+
+  template <class DTensor>
+  MUTE_HOST_DEVICE static
+  auto
+  retile_D(DTensor&& dtensor) {
+    // static_assert(sizeof(typename remove_cvref_t<DTensor>::value_type) == sizeof(typename TiledCopy::ValType),
+    //               "Expected ValType for tiling DstTensor.");
+    return make_tensor(static_cast<DTensor&&>(dtensor).data(), TiledCopy::retile(dtensor.layout()));
+  }
+};
+
+
+template <class... Args,
+          class LayoutCopy_TV,
+          class Tiler>
+MUTE_HOST_DEVICE
+auto
+make_tiled_copy_impl(Copy_Atom<Args...> const& atom,
+                     LayoutCopy_TV      const&,
+                     Tiler              const&)
+{
+  return TiledCopy<Copy_Atom<Args...>, LayoutCopy_TV, Tiler>{atom};
+}
+
+//
+// These tile the Copy_Atom as a whole
+//
+
+template <class... CArgs, class... MArgs>
+MUTE_HOST_DEVICE
+auto
+make_tiled_copy_A(Copy_Atom<CArgs...> const& copy_atom,
+                  TiledMMA<MArgs...>  const& mma)
+{
+  return make_tiled_copy_impl(copy_atom, mma.get_layoutA_TV(), make_shape(tile_size<0>(mma),tile_size<2>(mma)));
+}
+
+template <class... CArgs, class... MArgs>
+MUTE_HOST_DEVICE
+auto
+make_tiled_copy_B(Copy_Atom<CArgs...> const& copy_atom,
+                  TiledMMA<MArgs...>  const& mma)
+{
+  return make_tiled_copy_impl(copy_atom, mma.get_layoutB_TV(), make_shape(tile_size<1>(mma),tile_size<2>(mma)));
+}
+
+template <class... CArgs, class... MArgs>
+MUTE_HOST_DEVICE
+auto
+make_tiled_copy_C(Copy_Atom<CArgs...> const& copy_atom,
+                  TiledMMA<MArgs...>  const& mma)
+{
+  return make_tiled_copy_impl(copy_atom, mma.get_layoutC_TV(), make_shape(tile_size<0>(mma),tile_size<1>(mma)));
+}
+
+// returns the smallest tiled copy that can retile LayoutC_TV
+// for use with pipelined epilogues with subtiled stores
+template <class... CArgs, class... MArgs>
+MUTE_HOST_DEVICE
+auto
+make_tiled_copy_C_atom(Copy_Atom<CArgs...> const& copy_atom,
+                       TiledMMA<MArgs...>  const& mma)
+{
+  // Truncate the V-layout to just the Copy_Atom, keep the V-order
+  auto layoutC_TV = mma.get_layoutC_TV();
+  auto copy_V     = Int<Copy_Atom<CArgs...>::NumValSrc>{};
+  MUTE_STATIC_ASSERT_V(copy_V <= size<1>(layoutC_TV));
+  auto layout_TV  = composition(layoutC_TV, make_layout(make_shape(size<0>(layoutC_TV), copy_V)));
+
+  // Recompute tiler and restride the TV layout for the new tiler
+
+  // Tiler -- Find the active elements in the MMA tensor and generate a tiler to extract them
+  // Convert to the awkward by-mode tiler to preserve the modes of the tiled MMA
+  auto mma_tiler = make_shape(tile_size<0>(mma),tile_size<1>(mma));
+  auto mma_zeros = repeat_like(mma_tiler, Int<0>{});
+
+  auto tiler = transform(make_seq<rank(mma_tiler)>{}, [&](auto i) {
+    return filter(composition(make_layout(mma_tiler, replace<i>(mma_zeros, Int<1>{})), layout_TV));
+  });
+
+  // Layout_TV -- Find the (tid,vid) -> tile coord transformation
+  // Apply the tiler to a reference and transform the codomain
+  // tile_coord -> mma_coord
+  auto tile2mma = composition(make_layout(mma_tiler), tiler);
+
+  // (tid,vid) -> tile_coord
+  auto layout_tv = composition(left_inverse(tile2mma), layout_TV);
+
+  return make_tiled_copy_impl(copy_atom, layout_tv, tiler);
+}
+
+/** Produce a TiledCopy from logical thread and values layouts.
+ * The thread and value layouts map coordinates to thr_idx and val_idx.
+ *    The product of these layouts is taken to produce the TV layout and the Tiler.
+ * Useful when threads and values need very specific mappings onto coordinates
+ *    in the target tensors.
+ */
+template <class... Args,
+          class ThrLayout,
+          class ValLayout = Layout<_1>>
+MUTE_HOST_DEVICE
+auto
+make_tiled_copy(Copy_Atom<Args...> const& copy_atom,
+                ThrLayout          const& thr_layout = {},     // (m,n) -> thr_idx
+                ValLayout          const& val_layout = {})     // (m,n) -> val_idx
+{
+  // Take the raked_products to compute the Layout_MN
+  // (M,N) -> (thr_idx, val_idx)
+  auto layout_mn = raked_product(thr_layout, val_layout);
+  // (thr_idx, val_idx) -> (M,N)
+  auto layout_tv = right_inverse(layout_mn).with_shape(make_shape(size(thr_layout), size(val_layout)));
+  // Tiler for extracting relevant elements
+  // (M,N) -> tensor coord
+  auto tiler = product_each(shape(layout_mn));
+
+#if 0
+  print("thr_layout: "); print(thr_layout); print("\n");
+  print("val_layout: "); print(val_layout); print("\n");
+  print("layout_mn : "); print(layout_mn);  print("\n");
+  print("layout_tv : "); print(layout_tv);  print("\n");
+  print("tiler     : "); print(tiler);      print("\n");
+#endif
+
+  return make_tiled_copy_impl(copy_atom, layout_tv, tiler);
+}
+
+/** Produce a TiledCopy from thread and value offset maps.
+ * The TV Layout maps threads and values to the codomain of the data_layout.
+ * It is verified that the intended codomain is valid within data_layout.
+ * Useful when threads and values don't care about owning specific coordinates, but
+ *   care more about the vector-width and offsets between them.
+ */
+template <class... Args, class AtomTVLayout, class DataLayout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_cotiled_copy(Copy_Atom<Args...> const& copy_atom,
+                  AtomTVLayout const& atom_tv_layout,   // atom (thr,val) -> data addr
+                  DataLayout   const& data_layout)      // coord          -> data addr    The target layout
+{
+  static_assert(is_static<AtomTVLayout>::value);
+  static_assert(is_static<DataLayout>::value);
+
+  // data addr -> data coord    Append 1:0 so off-the-ends get the stride-0
+  auto inv_data_layout = make_layout(left_inverse(data_layout), Layout<_1,_0>{});
+
+  // (tid,vid) -> data_coord
+  auto layout_tv_data = composition(inv_data_layout, atom_tv_layout);
+
+  // Check validity
+  MUTE_STATIC_ASSERT_V(coalesce(composition(data_layout, layout<1>(layout_tv_data))) == coalesce(layout<1>(atom_tv_layout)),
+                       "The memory pointed to by AtomTVLayout does not exist in the DataLayout.");
+
+#if 0
+  if (thread0()) {
+    print("data_layout        : "); print(data_layout); print("\n");
+    print("atom_tv_layout     : "); print(atom_tv_layout); print("\n");
+    print("layout_tv_data     : "); print(layout_tv_data); print("\n");
+  }
+#endif
+
+  //
+  // Tiler -- Find the active elements in the DATA tensor and generate a tiler to extract them
+  //
+
+  // Convert to the awkward by-mode tiler to preserve the modes of the tiled DATA
+  auto flat_data_shape = product_each(shape(data_layout));
+  auto flat_data_zeros = repeat<rank(flat_data_shape)>(Int<0>{});
+
+  auto tiler = transform(make_seq<rank(flat_data_shape)>{}, [&](auto i) {
+    return filter(composition(make_layout(flat_data_shape, replace<i>(flat_data_zeros, Int<1>{})), layout_tv_data));
+  });
+
+  //
+  // Layout_TV -- Find the (tid,vid) -> tile coord transformation
+  //
+
+  // Apply the tiler to a reference and transform the codomain
+  // tile_coord -> data_coord
+  auto tile2data = composition(make_layout(flat_data_shape), tiler);
+
+  // (tid,vid) -> tile_coord
+  auto layout_tv = composition(left_inverse(tile2data), layout_tv_data);
+
+#if 0
+  if (thread0()) {
+    print("tiler              : "); print(tiler); print("\n");
+    print("tile2data          : "); print(tile2data); print("\n");
+    print("layout_tv          : "); print(layout_tv); print("\n");
+  }
+#endif
+
+  return make_tiled_copy_impl(copy_atom, layout_tv, tiler);
+}
+
+// Make a TiledCopy out of the copy_atom that matches the Src-Layout of tiled_copy
+template <class... Args,
+          class TiledCopy>
+MUTE_HOST_DEVICE
+auto
+make_tiled_copy_S(Copy_Atom<Args...> const& copy_atom,
+                  TiledCopy          const& tiled_copy)
+{
+  return make_tiled_copy_impl(copy_atom, tiled_copy.get_layoutS_TV(), typename TiledCopy::Tiler_MN{});
+}
+
+// Make a TiledCopy out of the copy_atom that matches the Dst-Layout of tiled_copy
+template <class... Args,
+          class TiledCopy>
+MUTE_HOST_DEVICE
+auto
+make_tiled_copy_D(Copy_Atom<Args...> const& copy_atom,
+                  TiledCopy          const& tiled_copy)
+{
+  return make_tiled_copy_impl(copy_atom, tiled_copy.get_layoutD_TV(), typename TiledCopy::Tiler_MN{});
+}
+
+//
+// Size
+//
+
+// The logical size of a TileCopy
+template <int... I, class... Args>
+MUTE_HOST_DEVICE constexpr
+auto
+tile_size(TiledCopy<Args...> const&)
+{
+  return size<I...>(typename TiledCopy<Args...>::Tiler_MN{});
+}
+
+// The number of threads involved in a TiledCopy
+template <class... Args>
+MUTE_HOST_DEVICE constexpr
+auto
+size(TiledCopy<Args...> const&)
+{
+  return typename TiledCopy<Args...>::TiledNumThr{};
+}
+
+//
+// Display utilities
+//
+
+template <class... Args, class T>
+MUTE_HOST_DEVICE
+void
+print(Copy_Atom<Copy_Traits<Args...>, T> const&)
+{
+  using Atom = Copy_Atom<Copy_Traits<Args...>, T>;
+  print("Copy_Atom\n");
+  print("  ThrID:        "); print(typename Atom::ThrID{});        print("\n");
+  print("  ValLayoutSrc: "); print(typename Atom::ValLayoutSrc{}); print("\n");
+  print("  ValLayoutDst: "); print(typename Atom::ValLayoutDst{}); print("\n");
+  print("  ValLayoutRef: "); print(typename Atom::ValLayoutRef{}); print("\n");
+  print("  ValueType:    "); print(sizeof_bits<typename Atom::ValType>::value); print("b\n");
+}
+
+template <class Atom, class... Args>
+MUTE_HOST_DEVICE
+void
+print(TiledCopy<Atom, Args...> const& copy, char const* pad = "")
+{
+  using Copy = TiledCopy<Atom, Args...>;
+  print("TiledCopy\n");
+  print("  Tiler_MN:       "); print(typename Copy::Tiler_MN{});       print("\n");
+  print("  TiledLayout_TV: "); print(typename Copy::TiledLayout_TV{}); print("\n");
+  print(static_cast<Atom const&>(copy));
+}
+
+template <class TiledCopy, class ThrIdx>
+MUTE_HOST_DEVICE
+void
+print(ThrCopy<TiledCopy, ThrIdx> const& thr_copy)
+{
+  print("ThrCopy\n");
+  print("  ThrIdx: "); print(thr_copy.thr_idx_); print("\n");
+  print(TiledCopy{});
+}
+
+// TiledCopy to LaTeX TikZ
+template <class... Args, class TikzColorFn = TikzColor_TV>
+MUTE_HOST_DEVICE
+auto
+print_latex(TiledCopy<Args...> const& copy,
+            TikzColorFn color = {})
+{
+  auto [layoutS_MN, thrID_S] = copy.get_layoutS_MN();
+  auto [layoutD_MN, thrID_D] = copy.get_layoutD_MN();
+
+  print_latex_copy(layoutS_MN, thrID_S,
+                   layoutD_MN, thrID_D);
+}
+
+// MNK Copy Layout to Latex TIKZ
+template <class LayoutS, class ThrIDS,
+          class LayoutD, class ThrIDD,
+          class TikzColorFn = TikzColor_TV>
+MUTE_HOST_DEVICE
+void
+print_latex_copy(LayoutS const& S, ThrIDS const& TS,  // (m,n) -> (tid,vid)  and  tid -> thr_idx
+                 LayoutD const& D, ThrIDD const& TD,  // (m,n) -> (tid,vid)  and  tid -> thr_idx
+                 TikzColorFn color = {})              // lambda(thr_idx,val_idx) -> tikz color string
+{
+  MUTE_STATIC_ASSERT_V(rank(S) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(D) == Int<2>{});
+
+  assert(size<0>(S) == size<0>(D));
+  assert(size<1>(S) == size<1>(D));
+
+  // Commented prints
+  printf("%% LayoutS: "); print(S);  printf("\n");
+  printf("%% ThrIDS : "); print(TS); printf("\n");
+  printf("%% LayoutD: "); print(D);  printf("\n");
+  printf("%% ThrIDD : "); print(TD); printf("\n\n");
+
+  // Header
+  printf("\\documentclass[convert]{standalone}\n"
+         "\\usepackage{tikz}\n\n"
+         "\\begin{document}\n"
+         "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
+
+  // S starting at 0,0
+  for (int i = 0; i < size<0>(S); ++i) {
+    for (int j = 0; j < size<1>(S); ++j) {
+      int thrid   = S(i,j) % size(TS);
+      int val_idx = S(i,j) / size(TS);
+      int thr_idx = TS(thrid);
+
+      printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
+             color(thr_idx, val_idx),
+             i, j,
+             thr_idx, val_idx);
+    }
+  }
+  // Grid
+  printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
+         0, 0, int(size<0>(S)), int(size<1>(S)));
+  // S Labels
+  for (int i =  0, j = -1; i < size<0>(S); ++i) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
+  }
+  for (int i = -1, j =  0; j < size<1>(S); ++j) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j);
+  }
+
+  // D starting at 0,size<1>(S)+3
+  for (int i = 0; i < size<0>(D); ++i) {
+    for (int j = 0; j < size<1>(D); ++j) {
+      int thrid   = D(i,j) % size(TD);
+      int val_idx = D(i,j) / size(TD);
+      int thr_idx = TD(thrid);
+
+      printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
+             color(thr_idx, val_idx),
+             i, j + size<1>(S) + 3,
+             thr_idx, val_idx);
+    }
+  }
+
+  // Grid
+  printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (%d,%d) grid (%d,%d);\n\n",
+         0, int(size<1>(S)+3), int(size<0>(D)), int(size<1>(D)+size<1>(S)+3));
+  // D Labels
+  for (int i = 0, j = size<1>(D); i < size<0>(D); ++i) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j + size<1>(S) + 3, i);
+  }
+  for (int i = -1, j =          0; j < size<1>(D); ++j) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j + size<1>(S) + 3, j);
+  }
+
+  // Footer
+  printf("\\end{tikzpicture}\n"
+         "\\end{document}\n");
+}
+
+} // end namespace mute
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include <mute/atom/copy_traits_mp31.hpp>
+#include <mute/atom/copy_traits_mp31_tme.hpp>
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits.hpp
new file mode 100644
index 000000000..70f099cd7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits.hpp
@@ -0,0 +1,164 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/arch/copy.hpp>
+
+#include <mute/tensor.hpp>
+
+namespace mute
+{
+
+/**
+ * concept Copy_Traits
+ * {
+ *   using ThrID     =    // Logical thread id (tid) -> tidx
+ *
+ *   using SrcLayout =    // (Logical src thread id (tid), Logical src value id (vid)) -> bit
+ *   using DstLayout =    // (Logical dst thread id (tid), Logical dst value id (vid)) -> bit
+ *   using RefLayout =    // (Logical ref thread id (tid), Logical ref value id (vid)) -> bit
+ * };
+ *
+ * The abstract bit ordering of the Copy_Traits (the codomain of SrcLayout, DstLayout, and RefLayout)
+ * is arbitrary and only used to construct maps
+ *   (ref-tid,ref-vid) -> (src-tid,src-vid)
+ *   (ref-tid,ref-vid) -> (dst-tid,dst-vid)
+ * in TiledCopy. The Layout_TV in TiledCopy is in accordance with the RefLayout of a Traits, then mapped to
+ * the Src or Dst (tid,vid) representation on demand.
+ *
+ */
+
+template <class CopyOperation, class... CopyOpArgs>
+struct Copy_Traits
+{
+  static_assert(dependent_false<CopyOperation>, "Copy_Traits not implemented for this CopyOperation.");
+};
+
+template <class S, class D>
+struct Copy_Traits<UniversalCopy<S,D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1,Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1,Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+};
+
+template <int MaxVecBits>
+struct Copy_Traits<AutoVectorizingCopyWithAssumedAlignment<MaxVecBits>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1,_1>, Stride<_0,_0>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1,_1>, Stride<_0,_0>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+};
+
+// Extract a CPY_Op from a CPY_Traits
+template <class CPY_Traits>
+struct CPY_Op {};
+
+template <class CPY_Op_Arg, class... Args>
+struct CPY_Op<Copy_Traits<CPY_Op_Arg, Args...>> {
+  using type = CPY_Op_Arg;
+};
+
+//
+// Generic copy_unpack for common argument-based Copy_Traits
+//
+
+template <class AnyCPYTraits,
+          class SEngine, class SLayout,
+          class DEngine, class DLayout>
+MUTE_HOST_DEVICE constexpr
+void
+copy_unpack(AnyCPYTraits            const&,
+            Tensor<SEngine,SLayout> const& src,
+            Tensor<DEngine,DLayout>      & dst)
+{
+  using CopyOp       = typename CPY_Op<AnyCPYTraits>::type;
+  using RegistersSrc = typename CopyOp::SRegisters;
+  using RegistersDst = typename CopyOp::DRegisters;
+  using RegTypeSrc   = typename remove_extent<RegistersSrc>::type;
+  using RegTypeDst   = typename remove_extent<RegistersDst>::type;
+  constexpr int RegNumSrc = extent<RegistersSrc>::value;
+  constexpr int RegNumDst = extent<RegistersDst>::value;
+
+  Tensor rS = recast<RegTypeSrc>(src);
+  Tensor rD = recast<RegTypeDst>(dst);
+
+  MUTE_STATIC_ASSERT_V(size(rS) == Int<RegNumSrc>{},
+    "Copy_Traits: src failed to vectorize into registers. Layout is incompatible with this CopyOp.");
+  MUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumDst>{},
+    "Copy_Traits: dst failed to vectorize into registers. Layout is incompatible with this CopyOp.");
+
+  detail::explode(detail::CallCOPY<CopyOp>{},
+                  rS, make_int_sequence<RegNumSrc>{},
+                  rD, make_int_sequence<RegNumDst>{});
+}
+
+// Accept mutable temporaries
+template <class AnyCPYTraits,
+          class SEngine, class SLayout,
+          class DEngine, class DLayout>
+MUTE_HOST_DEVICE constexpr
+void
+copy_unpack(AnyCPYTraits            const& traits,
+            Tensor<SEngine,SLayout> const& src,
+            Tensor<DEngine,DLayout>     && dst)
+{
+  copy_unpack(traits, src, dst);
+}
+
+// Element copy selector
+template <class SrcTensor, class DstTensor>
+MUTE_HOST_DEVICE constexpr
+auto
+select_elementwise_copy(SrcTensor const&, DstTensor const&)
+{
+  using SrcType = typename SrcTensor::value_type;
+  using DstType = typename DstTensor::value_type;
+
+  return UniversalCopy<SrcType,DstType>{};
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31.hpp
new file mode 100644
index 000000000..3daa9f124
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31.hpp
@@ -0,0 +1,372 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/arch/copy_mp31.hpp>
+#include <mute/atom/copy_traits.hpp>
+
+
+namespace mute {
+
+//
+// ROBUST LOAD
+//
+
+template <class S, class D>
+struct MP31_ROBUST_LOAD_OP : MP31_ROBUST_LOAD<S, D> {};
+
+template <class S, class D>
+struct MP31_ROBUST_PRED_LOAD_OP : MP31_ROBUST_LOAD<S, D> {};
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_PRED_LOAD_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  // Predicate value that determines whether to load or return zero
+  bool pred;
+
+  // Overload copy_unpack to pass robust buffer access descriptor
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst)
+  {
+    static_assert(is_gmem<TS>::value, "Expected gmem src for MP31_ROBUST_LOAD");
+    static_assert(is_rmem<TD>::value, "Expected rmem dst for MP31_ROBUST_LOAD");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    MUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+        "This src layout is incompatible with this tiled copy.");
+    MUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+        "This dst layout is incompatible with this tiled copy.");
+
+    MP31_ROBUST_LOAD<S, D>::copy(rS[0], rD[0], traits.pred, traits.src_desc);
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_LOAD_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP31_ROBUST_PRED_LOAD_OP<S, D>>
+  with(bool pred) const {
+    return {src_desc, pred};
+  }
+};
+
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_LOAD<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Construct an executable MP31_ROBUST_LOAD with robust buffer access descriptor
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP31_ROBUST_LOAD_OP<S, D>>
+  with(RobustDescriptor desc) const {
+    return {desc};
+  }
+
+  // Don't try to execute a copy with MP31_ROBUST_LOAD before calling .with()
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst) = delete;
+
+};
+
+//
+// ROBUST STORE
+//
+
+template <class S, class D>
+struct MP31_ROBUST_STORE_OP : MP31_ROBUST_STORE<S, D> {};
+
+template <class S, class D>
+struct MP31_ROBUST_PRED_STORE_OP : MP31_ROBUST_STORE<S, D> {};
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_PRED_STORE_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor dst_desc;
+
+  // Predicate value that determines whether to store or discard
+  bool pred;
+
+  // Overload copy_unpack to pass robust buffer access descriptor
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst)
+  {
+    static_assert(is_rmem<TS>::value, "Expected rmem src for MP31_ROBUST_STORE");
+    static_assert(is_gmem<TD>::value, "Expected gmem dst for MP31_ROBUST_STORE");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    MUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+        "This src layout is incompatible with this tiled copy.");
+    MUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+        "This dst layout is incompatible with this tiled copy.");
+
+    MP31_ROBUST_STORE<S, D>::copy(rS[0], rD[0], traits.pred, traits.dst_desc);
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_STORE_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor dst_desc;
+
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP31_ROBUST_PRED_STORE_OP<S, D>>
+  with(bool pred) const {
+    return {dst_desc, pred};
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_STORE<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Construct an executable MP31_ROBUST_STORE with robust buffer access descriptor
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP31_ROBUST_STORE_OP<S, D>>
+  with(RobustDescriptor desc) const {
+    return {desc};
+  }
+
+  // Don't try to execute a copy with MP31_ROBUST_STORE before calling .with()
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst) = delete;
+
+};
+
+//
+// ROBUST LDGSTS
+//
+
+template <class S, class D>
+struct MP31_ROBUST_LDGSTS_OP : MP31_ROBUST_LDGSTS<S, D> {};
+
+template <class S, class D>
+struct MP31_ROBUST_PRED_LDGSTS_OP : MP31_ROBUST_LDGSTS<S, D> {};
+
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_PRED_LDGSTS_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  // Predicate value that determines whether to load or return zero
+  bool pred;
+
+  // Overload copy_unpack to pass robust buffer access descriptor
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst)
+  {
+    static_assert(is_gmem<TS>::value, "Expected gmem src for MP31_ROBUST_LDGSTS");
+    static_assert(is_smem<TD>::value, "Expected smem dst for MP31_ROBUST_LDGSTS");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    MUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+        "This src layout is incompatible with this tiled copy.");
+    MUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+        "This dst layout is incompatible with this tiled copy.");
+
+    MP31_ROBUST_LDGSTS<S, D>::copy(rS[0], rD[0], traits.pred, traits.src_desc);
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_LDGSTS_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP31_ROBUST_PRED_LDGSTS_OP<S, D>>
+  with(bool pred) const {
+    return {src_desc, pred};
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP31_ROBUST_LDGSTS<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Construct an executable MP31_ROBUST_LDGSTS with robust buffer access descriptor
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP31_ROBUST_LDGSTS_OP<S, D>>
+  with(RobustDescriptor desc) const {
+    return {desc};
+  }
+
+  // Don't try to execute a copy with MP31_ROBUST_LDGSTS before calling .with()
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst) = delete;
+
+};
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31_tme.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31_tme.hpp
new file mode 100644
index 000000000..bd892e28e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31_tme.hpp
@@ -0,0 +1,893 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#if !defined(__MUSACC_RTC__)
+#include <musa.h>
+#endif
+
+#include <mute/atom/copy_traits_mp31_tme_swizzle.hpp>
+#include <mute/atom/copy_traits.hpp>
+#include <mute/atom/copy_atom.hpp>
+
+#include <mute/numeric/integral_ratio.hpp>
+
+namespace mute
+{
+
+template <class GmemTmeBasisStrides_, class TmeGmemBasis_, class TmeSwizzle_, class TmeBlockDim_>
+struct AuxTmeParams {
+  using GmemStrides  = GmemTmeBasisStrides_;    // Strides for Gmem mode -> Tme coord mode, may be dynamic
+  GmemStrides g_stride_;
+  using TmeGmemBasis = TmeGmemBasis_;
+  static_assert(is_static<TmeGmemBasis>::value);
+  using TmeSwizzle   = TmeSwizzle_;
+  static_assert(is_static<TmeSwizzle>::value);
+  using TmeBlockDim  = TmeBlockDim_;
+  static_assert(is_static<TmeBlockDim>::value);
+  TmeBlockDim block_dim_;
+};
+
+// Utility for unpacking TME_LOAD arguments into a CopyOp
+template <class CopyOp>
+struct TME_LOAD_Unpack
+{
+  template <class... Args,
+            class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr void
+  copy_unpack(Copy_Traits<CopyOp, Args...> const& traits,
+              Tensor<TS,SLayout>           const& src,
+              Tensor<TD,DLayout>                & dst)
+  {
+    auto src_coord = src.data().coord_;
+    static_assert(is_smem<TD>::value, "MP31_TME_LOAD requires the destination be shared memory.");
+    void* dst_ptr = mute::raw_pointer_cast(dst.data());
+
+    auto aux_params = get<2>(traits.opargs_);
+    auto block_dim = mute::flatten(aux_params.block_dim_);
+
+    return detail::explode_tuple(detail::CallCOPY<CopyOp>{},
+                                 traits.opargs_, make_seq<2>{},
+                                 make_tuple(dst_ptr), seq<0>{},
+                                 src_coord, tuple_seq<decltype(src_coord)>{},
+                                 block_dim, tuple_seq<decltype(block_dim)>{});
+  }
+};
+
+//////////////////////////////////////////////////////////////////////////////
+///////////////////////////// TME_LOAD ///////////////////////////////////////
+//////////////////////////////////////////////////////////////////////////////
+
+template <
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct MP31_TME_LOAD_OP : MP31_TME_LOAD_ENTRY<sg, ss, sl, prefetch> {};
+
+
+// The non-executable MP31_TME_LOAD_ENTRY with tme_desc and no tme_bar
+// Use .with(tme_bar) to construct an executable version
+template <
+  class NumBitsPerTME,
+  class AuxParams_,
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct Copy_Traits<MP31_TME_LOAD_ENTRY<sg, ss, sl, prefetch>, NumBitsPerTME, AuxParams_>
+{
+  using ThrID    = Layout<_1>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1,NumBitsPerTME>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1,NumBitsPerTME>>;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // MP31_TME_LOAD arguments
+  TmeDescriptor tme_desc_;
+  using AuxParams = AuxParams_;
+  AuxParams aux_params_;
+
+  // Return TmeDescriptor
+  MUTE_HOST_DEVICE constexpr
+  TmeDescriptor const*
+  get_tme_descriptor() const {
+    return &tme_desc_;
+  }
+
+  // Construct an executable MP31_TME_LOAD with tme_bar
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP31_TME_LOAD_OP<sg, ss, sl, prefetch>, NumBitsPerTME, AuxParams>
+  with(uint32_t tme_bar) const {
+    return {{}, {&tme_desc_, tme_bar, aux_params_}};
+  }
+
+  // Generate the TME coord tensor
+  template <class GShape>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_tme_tensor(GShape const& g_shape) const {
+    static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
+    return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
+  }
+
+  // Don't try to execute a copy with SM90_TME_LOAD before calling .with()
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr void
+  copy_unpack(Copy_Traits        const& traits,
+              Tensor<TS,SLayout> const& src,
+              Tensor<TD,DLayout>      & dst) = delete;
+};
+
+// The executable MP31_TME_LOAD with tme_desc, tme_bar & tme_params
+template <
+  class NumBitsPerTME,
+  class AuxParams,
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl,
+  PrefetchSize          prefetch
+>
+struct Copy_Traits<MP31_TME_LOAD_OP<sg, ss, sl, prefetch>, NumBitsPerTME, AuxParams>
+     : TME_LOAD_Unpack<MP31_TME_LOAD_OP<sg, ss, sl, prefetch>>
+{
+  using ThrID    = Layout<_1>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1,NumBitsPerTME>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1,NumBitsPerTME>>;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // MP31_TME_LOAD arguments
+  tuple<
+    TmeDescriptor const*,
+    uint32_t const&,
+    AuxParams const&
+  > const opargs_;
+};
+
+//////////////////////////////////////////////////////////////////////////////
+///////////////////////////// TME_STORE //////////////////////////////////////
+//////////////////////////////////////////////////////////////////////////////
+
+// The executable MP31_TME_STORE_ENTRY with tme_desc
+template <
+  class NumBitsPerTME,
+  class AuxParams_,
+  TME::SmemSwizzleGranularity sg,
+  TME::SmemSwizzleStride      ss,
+  TME::SmemSwizzleLine        sl
+>
+struct Copy_Traits<MP31_TME_STORE_ENTRY<sg, ss, sl>, NumBitsPerTME, AuxParams_>
+{
+  using ThrID    = Layout<_1>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1,NumBitsPerTME>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1,NumBitsPerTME>>;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // MP31_TME_STORE arguments
+  TmeDescriptor tme_desc_;
+  using AuxParams = AuxParams_;
+  AuxParams aux_params_;
+
+  // Return TmeDescriptor
+  MUTE_HOST_DEVICE constexpr
+  TmeDescriptor const*
+  get_tme_descriptor() const {
+    return &tme_desc_;
+  }
+
+  // Generate the TME coord tensor
+  template <class GShape>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_tme_tensor(GShape const& g_shape) const {
+    static_assert(is_congruent<decltype(g_shape), decltype(aux_params_.g_stride_)>::value);
+    return make_counting_tensor(make_layout(g_shape, aux_params_.g_stride_));
+  }
+
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr void
+  copy_unpack(Copy_Traits        const& traits,
+              Tensor<TS,SLayout> const& src,
+              Tensor<TD,DLayout>      & dst)
+  {
+    static_assert(is_smem<TS>::value, "Expected smem src for MP31_TME_STORE");
+    //static_assert(is_gmem<TD>::value, "Expected gmem dst for MP31_TME_STORE");
+
+    void const* const desc_ptr = &(traits.tme_desc_);
+    void const* const smem_ptr = mute::raw_pointer_cast(src.data());
+    auto dst_coord = dst.data().coord_;
+    auto block_dim = mute::flatten(traits.aux_params_.block_dim_);
+
+#if 0
+    auto [c0,c1,c2,c3,c4] = append<5>(dst_coord, 0);
+    if (thread0()) {
+      printf("TME CRD (%d,%d,%d,%d,%d) SMEM_ADDR (%p)\n",
+             int32_t(c0), int32_t(c1), int32_t(c2), int32_t(c3), int32_t(c4), smem_ptr);
+    }
+#endif
+
+    return detail::explode_tuple(detail::CallCOPY<MP31_TME_STORE_ENTRY<sg, ss, sl>>{},
+                                 make_tuple(desc_ptr, smem_ptr), seq<0,1>{},
+                                 dst_coord, tuple_seq<decltype(dst_coord)>{},
+                                 block_dim, tuple_seq<decltype(block_dim)>{});
+  }
+};
+
+
+//////////////////////////////////////////////////////////////////////////////
+///////////////////////////// BLK_COPY ///////////////////////////////////////
+//////////////////////////////////////////////////////////////////////////////
+
+template <class NumBitsPerTME, class... OpArgs>
+struct Copy_Traits<MP31_BLK_COPY_G2S, NumBitsPerTME, OpArgs...>
+{
+
+  using ThrID    = Layout<_1>;
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1,NumBitsPerTME>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1,NumBitsPerTME>>;
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // MP31_BLK_COPY_G2S arguments
+  // 0: uint32_t& barrier_id
+  mute::tuple<OpArgs...> blk_load_bar_;
+
+  // Record the async barrier for the instruction
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP31_BLK_COPY_G2S, NumBitsPerTME, uint32_t const&>
+  with(uint32_t tme_bar) const {
+    return {{tme_bar}};
+  }
+
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits        const& traits,
+              Tensor<TS,SLayout> const& src,
+              Tensor<TD,DLayout>      & dst)
+  {
+    static_assert(is_same<mute::tuple<OpArgs...>, mute::tuple<uint32_t const&>>::value,
+                  "Extra arguments not set. Set .with() before use.");
+    static_assert(is_gmem<TS>::value, "Expected gmem src for MP31_BLK_COPY_G2S");
+    static_assert(is_smem<TS>::value, "Expected smem dst for MP31_BLK_COPY_G2S");
+    MP31_BLK_COPY_G2S::copy(raw_pointer_cast(src.data()), get<0>(traits.blk_load_bar_),
+                            raw_pointer_cast(dst.data()), int32_t(NumBitsPerTME::value / 8));
+
+  }
+
+};
+
+//
+// MAKE_TME_COPY and related
+//
+
+namespace detail {
+
+// Custom version of coalesce that greedily combines modes only up to size-65536
+// Look at each element and the back of the stack (in order of priority)
+// back(NewLayout)  get<I>(OldLayout)
+//      s0:d0           _1:d1     =>  continue
+//      _1:d0           s1:d1     =>  replace_back     s1:d1
+//      s0:d0           s1:s0*d0  =>  replace_back  s0*s1:d0   if s0*s1 <= 65536
+//      s0:d0           s1:d1     =>  append           s1:d1
+//
+// @pre OldShape and OldStride are flat
+template <int I, class OldShape, class OldStride, class NewShape, class NewStride>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce_65536_impl(OldShape const& old_shape, OldStride const& old_stride,
+                    NewShape const& new_shape, NewStride const& new_stride)
+{
+  if constexpr (I == rank_v<OldShape>) {
+    // Base case, we're done
+    if constexpr (is_constant<1, NewShape>::value) {
+      return Layout<_1,_0>{};
+    } else {
+      return Layout<NewShape,NewStride>{new_shape,new_stride};
+    }
+  } else if constexpr (is_constant<1, decltype(get<I>(old_shape))>::value) {
+    // shape<I>(layout) == _1, skip it and continue
+    return coalesce_65536_impl<I+1>(old_shape, old_stride, new_shape, new_stride);
+  } else if constexpr (is_constant<1, NewShape>::value) {
+    // Replace our shape-1 with anything (Can only happen on input new_shape/new_stride)
+    return coalesce_65536_impl<I+1>(old_shape, old_stride, get<I>(old_shape), get<I>(old_stride));
+  } else if constexpr (is_constant<true, decltype(back(new_shape) * back(new_stride) == get<I>(old_stride) &&
+                                                  get<I>(old_shape) * back(new_shape) <= Int<65536>{})>::value) {
+    // Merge modes because the shapes and strides match and the merge is 65536 or less
+    return coalesce_65536_impl<I+1>(old_shape, old_stride,
+                                  replace_back(new_shape, get<I>(old_shape) * back(new_shape)),
+                                  new_stride);
+  } else {
+    // Can't replace or merge, so append a new mode
+    return coalesce_65536_impl<I+1>(old_shape, old_stride,
+                                  append(new_shape,  get<I>(old_shape)),
+                                  append(new_stride, get<I>(old_stride)));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Combine all the modes that are possible to combine
+// Does not respect the profile of the layout, but does preserve total size
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce_65536(Layout<Shape,Stride> const& layout)
+{
+  auto flat_shape  = flatten(layout.shape());
+  auto flat_stride = flatten(layout.stride());
+  return coalesce_65536_impl<1>(flat_shape, flat_stride, get<0>(flat_shape), get<0>(flat_stride));
+}
+
+
+template <class TmeInternalType,
+          class GEngine, class GLayout,
+          class SShape, class SStride,
+          class VShape, class VStride>
+MUTE_HOST_DEVICE constexpr
+auto
+construct_tme_gbasis(Tensor<GEngine,GLayout> const& gtensor,       // The original GMEM Tensor
+                     Layout<SShape,SStride>  const& slayout,       // The layout of SMEM
+                     Layout<VShape,VStride>  const& cta_v_map)     // smem_coord to hier gmode
+{
+  //
+  // TME parameter checking
+  //
+  MUTE_STATIC_ASSERT_V(product_each(shape(slayout)) == product_each(shape(cta_v_map)),
+                       "TME requires CTA_Tile and SLayout top-level shape equivalence.");
+
+#if 0
+  print("gtensor         : "); print(gtensor); print("\n");
+  print("slayout         : "); print(slayout); print("\n");
+  print("cta_v_map       : "); print(cta_v_map); print("\n");
+#endif
+
+  //
+  // TME slayout manipulation
+  //
+
+  // Invert the smem to get the largest contiguous vector in the smem layout
+  // smem idx -> smem coord
+  auto inv_smem_layout = right_inverse(get_nonswizzle_portion(slayout));
+
+  // Compose with the V-Map to convert smem coord (CTA val idx) to gmem mode
+  // smem idx -> gmem mode
+  auto sidx2gmode_full = coalesce(composition(cta_v_map, inv_smem_layout));
+
+#if 0
+  print("inv_smem_layout : "); print(inv_smem_layout); print("\n");
+  print("sidx2gmode_full : "); print(sidx2gmode_full); print("\n");
+#endif
+
+  //
+  // TME gtensor truncation
+  //
+
+  // Truncate any incompatibilities -- no starting in the middle of gmodes
+  auto smem_rank = find_if(stride(sidx2gmode_full), [](auto e) {
+    [[maybe_unused]] auto v = basis_value(e);
+    return not is_constant<1,decltype(v)>{};
+  });
+  static_assert(smem_rank > 0, "Could not find a common tile-gmem vectorization. Does the Tile select out major GMEM modes?");
+
+  // Keep only the static-1 basis modes into gmem
+  auto sidx2gmode = take<0,smem_rank>(sidx2gmode_full);
+
+#if 0
+  print("smem_rank  : "); print(smem_rank); print("\n");
+  print("sidx2gmode : "); print(sidx2gmode); print("\n");
+#endif
+
+  //
+  // TME gtensor manipulation
+  //
+
+  // The smem vector is the same units as gtensor, so compose first and then recast
+  // tme_val_idx:gmem_strides
+  auto tile_gstride = recast<TmeInternalType>(gtensor.compose(sidx2gmode)).layout();
+  // Coalesce modes up to size-65536
+  // tme_box_shape:gmem_strides
+  auto tme_gstride  = coalesce_65536(tile_gstride);
+
+  // Perform the tiling, recast, and coalesce to the gmem vector again, but with indirections to the gtensor modes
+  auto gbasis = make_identity_layout(shape(gtensor));
+  auto tile_gbasis_tmp = gbasis.compose(sidx2gmode);
+
+  // Instead of the recast (gbasis doesn't have type info), replace the shape with the already-recasted shape
+  // tme_box_shape:gmem_mode
+  auto tile_gbasis = make_layout(shape(tile_gstride), stride(tile_gbasis_tmp));
+
+  // "Coalesce" the tile basis into a compatible shape with the tme_gstride
+  auto tme_gbasis_tile = tile_gbasis.compose(make_layout(wrap(shape(tme_gstride))));
+
+  // Recast the original tensor for shape/stride inspections
+  Tensor gtensor_T = recast<TmeInternalType>(gtensor);
+
+  // Find missing bases that don't appear in tile_gbasis
+  auto tile_gbasis_remaining_stride = filter_tuple(flatten(shape (gtensor_T)), flatten(stride(gtensor_T)),
+                                                   flatten(stride(gbasis)),
+                                                   [&](auto s, auto d, auto e)
+  {
+    if constexpr (is_constant<1, decltype(s)>::value || is_constant<0, decltype(d)>::value) {
+      return mute::tuple<>{};          // If size-1 or stride-0, then don't append
+    } else {
+      using E = decltype(e);
+      auto has_e = any_of(flatten(stride(tme_gbasis_tile)), [] (auto tb) { return tb == E{}; });
+      if constexpr (decltype(has_e)::value) {
+        return mute::tuple<>{};        // If d was found, then don't append
+      } else {
+        return mute::tuple<E>(e);      // Else, this is missing so append
+      }
+    }
+  });
+
+  // Append the remaining basis modes that contribute to the TME with size-1
+  auto tile_gbasis_remaining_shape = repeat<rank(tile_gbasis_remaining_stride)>(Int<1>{});
+  auto tme_gbasis_full = make_layout(tuple_cat(wrap( shape(tme_gbasis_tile)), wrap(tile_gbasis_remaining_shape )),
+                                     tuple_cat(wrap(stride(tme_gbasis_tile)), wrap(tile_gbasis_remaining_stride)));
+
+  // Group the trailing modes to make this max rank-5 -- TME rank limitation
+  // tme_box_shape:gmem_mode
+  auto tme_gbasis = group<mute::min(rank(tme_gbasis_full),4),-1>(tme_gbasis_full);
+
+#if 0
+  print("tile_gstride : "); print(tile_gstride); print("\n");
+  print("tme_gstride  : "); print(tme_gstride); print("\n");
+  print("gbasis       : "); print(gbasis); print("\n");
+  print("tile_gbasis  : "); print(tme_gbasis_tile); print("\n");
+  print("tme_gbasis   : "); print(tme_gbasis); print("\n");
+#endif
+
+  return tme_gbasis;
+}
+
+template <class GEngine, class GLayout,
+          class TmeGmemBasisStride,
+          class ShapeT, size_t TmeRank>
+MUTE_HOST_DEVICE constexpr
+void
+fill_tme_gmem_shape_stride(Tensor<GEngine,GLayout>   const& gtensor,           // Gmem Shapes and Strides, in units of TmeInternalType
+                           TmeGmemBasisStride        const& tme_gbasis_stride, // Map Tme mode idx -> Gmem mode(s)
+                           mute::array<ShapeT,   TmeRank> & gmem_prob_shape,   // Tme Shapes, uint32_t or uin64_t
+                           mute::array<uint64_t, TmeRank> & gmem_prob_stride)  // Tme Strides
+{
+  static_assert(is_tuple<TmeGmemBasisStride>::value);
+  static_assert(is_same<uint32_t, ShapeT>::value || is_same<uint64_t, ShapeT>::value);
+
+  using TmeInternalType = typename GEngine::value_type;
+  constexpr int tme_rank = decltype(rank(tme_gbasis_stride))::value;
+  static_assert(TmeRank >= tme_rank);
+
+  auto gmem_shape  =  shape(gtensor);
+  auto gmem_stride = stride(gtensor);
+  // Use the indirections in tme_gbasis_stride into gtensor to construct the tme gmem shapes/strides
+  for_each(make_seq<tme_rank>{}, [&](auto i) {
+    constexpr int tme_i_rank = decltype(rank<i>(tme_gbasis_stride))::value;
+    if constexpr (tme_i_rank == 1) {
+      // Trivial contribution of this gmem mode to this tme mode
+      auto ej = unwrap(get<i>(tme_gbasis_stride));
+      gmem_prob_shape[i]  = basis_get(ej, gmem_shape);
+      gmem_prob_stride[i] = basis_get(ej, gmem_stride);
+    } else {
+      // Apply a recurrence to each gmem mode that contributes to this tme mode
+      for_each(get<i>(tme_gbasis_stride), [&](auto ej) {
+        // Problem shape
+        uint64_t shape_j  = basis_get(ej, gmem_shape);
+        // Problem stride (in bytes)
+        uint64_t stride_j = basis_get(ej, gmem_stride);
+        uint64_t old_stride = gmem_prob_stride[i];
+        gmem_prob_stride[i] = gcd(gmem_prob_stride[i], stride_j);
+
+        if (gmem_prob_stride[i] != 0) {
+          // Recurrence: g_shape = (s_i - 1) * (d_i / gcd_j d_j) + 1
+          gmem_prob_shape[i] = (gmem_prob_shape[i]-1) * (old_stride / gmem_prob_stride[i])
+                             +            (shape_j-1) * (stride_j   / gmem_prob_stride[i])
+                             + 1;
+        } else {
+          gmem_prob_shape[i] = shape_j;
+        }
+      });
+    }
+  });
+}
+
+
+template <class TmeInternalType,
+          class CopyOp,
+          class GEngine, class GLayout,
+          class TShape,  class TStride,
+          int B, int M, int S>
+MUTE_HOST_RTC
+auto
+make_tme_copy_desc(Tensor<GEngine,GLayout> const& gtensor,         // The original GMEM Tensor
+                   Layout<TShape,TStride>  const& tme_gbasis,      // TME mode -> GMEM mode mapping
+                   Swizzle<B,M,S>          const& swizzle)         // Swizzle fn on smem_idx
+{
+  //
+  // TME desc creation
+  //
+
+  constexpr int tme_dim = decltype(rank(tme_gbasis))::value;
+
+  //
+  // TME gmem desc info
+  //
+
+  // Recast the original tensor for shape/stride inspections
+  Tensor gtensor_T = recast<TmeInternalType>(gtensor);
+
+  void* gmem_address = (void*)raw_pointer_cast(gtensor_T.data());
+  auto  gmem_layout  = gtensor_T.layout();
+
+  mute::array<uint64_t, 5> gmem_prob_shape  = {1,1,1,1,1};
+  mute::array<uint64_t, 5> gmem_prob_stride = {0,0,0,0,0};
+
+  fill_tme_gmem_shape_stride(gtensor_T, stride(tme_gbasis), gmem_prob_shape, gmem_prob_stride);
+
+  // TME descriptor does not store the zeroth stride and assumes it is 1 (TmaInternalType element).
+  assert(gmem_prob_stride[0] == 1 && "Majorness of smem doesn't match majorness of gmem");
+
+  // convert strides to byte strides
+  for(uint64_t& stride : gmem_prob_stride) {
+    stride = (stride * sizeof_bits_v<TmeInternalType>) / 8;
+  }
+
+
+  //
+  // Construct the descriptor
+  //
+
+  TmeDescriptor tme_desc{};
+
+  //
+  // TME general info
+  //
+
+  MUtensorDescriptorDataType   tme_format     = TME::to_MUtensorDescriptorDataType<TmeInternalType>();
+  MUtensorDescriptorInterleave tme_interleave = MU_TENSOR_DESCRIPTOR_INTERLEAVE_NONE;
+  uint64_t                     tme_oobFill    = 0;
+
+  MUresult result = muTensorDescriptorEncode(
+      &tme_desc,
+      tme_format,
+      tme_dim,
+      gmem_address,
+      gmem_prob_shape.data(),
+      gmem_prob_stride.data() + 1, // gmem_prob_stride[0] implicitly 1
+      tme_interleave,
+      tme_oobFill);
+
+  if (result != MUSA_SUCCESS) {
+    std::cerr << "TME Desc Addr:    " << &tme_desc
+              << "\nformat          " << tme_format
+              << "\ndim             " << tme_dim
+              << "\ngmem_address    " << gmem_address
+              << "\nglobalDim       " << gmem_prob_shape
+              << "\nglobalStrides   " << gmem_prob_stride
+              << "\ninterleave      " << tme_interleave
+              << "\noobFill         " << tme_oobFill << std::endl;
+    std::cerr << "Error: Failed to initialize the TME descriptor: "
+              << result << std::endl;
+    //assert(false);
+  }
+
+  //
+  // TME auxiliary params
+  //
+  auto recast_ratio = mute::trait_ratio(sizeof_bits<typename GEngine::value_type>{},
+                                        sizeof_bits<             TmeInternalType>{});
+
+  auto gbasis = make_basis_like(shape(gtensor));
+
+  // Finally, get the inverse permutation of the E<i> bases for the mocked gmem stride
+  auto gmem_tme_basis_stride = transform_leaf(gbasis, [&](auto ei) {
+    auto si = basis_get(ei,  shape(gmem_layout));
+    auto di = basis_get(ei, stride(gmem_layout));
+    if constexpr (is_constant<1, decltype(si)>::value || is_constant<0, decltype(di)>::value) {
+      return Int<0>{};                  // If size-1 or stride-0, return arithmetic identity -- no contribution to the TME
+    } else {
+      auto tme_gmem_basis_stride = stride(tme_gbasis);
+      // Find j such that E<i> is in stride<j>(tme_gbasis)
+      using EI = decltype(ei);
+      [[maybe_unused]] auto j = find_if(tme_gmem_basis_stride, [&](auto tme_stride_j) { return any_of(tme_stride_j, [&](auto dj) { return dj == EI{}; }); });
+      if constexpr (decltype(j == rank(tme_gmem_basis_stride))::value) {
+        return Int<0>{};               // If not-found, return arithmetic identity -- no contribution to the TME
+      } else
+      if constexpr (decltype(j == Int<0>{})::value) {
+        auto scale = recast_ratio * basis_get(ei, stride(gtensor));
+        return E<j>{} * scale;         // Return TME Coord basis -- with a recast scale factor
+      } else
+      if constexpr (decltype(rank<j>(tme_gmem_basis_stride) == Int<1>{})::value) {
+        return E<j>{};                 // Return TME Coord basis -- known scale of Int<1>{}
+      } else {
+        int32_t scale = ceil_div(int32_t(di * sizeof_bits_v<TmeInternalType> / mute::max(gmem_prob_stride[j], uint64_t{16})), 8);
+        return E<j>{} * scale;         // Return TME Coord basis -- with a dynamic scale factor
+      }
+    }
+  });
+
+  constexpr auto attributes             = get_tme_swizzle_attributes(Swizzle<B, M, S>{});
+  constexpr auto tme_SwizzleGranularity = get<0>(attributes);
+  constexpr auto tme_SwizzleStride      = get<1>(attributes);
+  constexpr auto tme_SwizzleLine        = get<2>(attributes);
+  constexpr auto tme_l2Prefetch         = PrefetchSize::B128;
+
+
+  static_assert(is_same_v<CopyOp, MP31_TME_LOAD> || is_same_v<CopyOp, MP31_TME_STORE>, "Unsupported CopyOp");
+
+  using CopyOpEntry = mute::conditional_t<is_same_v<CopyOp, MP31_TME_LOAD>,
+                                          MP31_TME_LOAD_ENTRY<tme_SwizzleGranularity, tme_SwizzleStride, tme_SwizzleLine, tme_l2Prefetch>,
+                                          MP31_TME_STORE_ENTRY<tme_SwizzleGranularity, tme_SwizzleStride, tme_SwizzleLine>>;
+
+  auto block_dim = product_each(take<0, tme_dim>(tme_gbasis).shape());
+
+#if 0
+  print("block dim                :"); print(block_dim); print("\n");
+  print("gmem_tme_basis_stride   : "); print(gmem_tme_basis_stride); print("\n");
+  print("tme swizzle line        : "); print(to_string(tme_SwizzleLine)); print("\n");
+  print("tme swizzle stride      : "); print(to_string(tme_SwizzleStride)); print("\n");
+  print("tme swizzle granularity : "); print(to_string(tme_SwizzleGranularity)); print("\n");
+  print("tme prefetch size       : "); print(to_string(tme_l2Prefetch)); print("\n");
+#endif
+
+  using AuxParams = AuxTmeParams<decltype(gmem_tme_basis_stride),
+                                 decltype(tme_gbasis),
+                                 decltype(swizzle),
+                                 decltype(block_dim)>;
+  return mute::make_tuple(tme_desc,
+                          AuxParams{gmem_tme_basis_stride, block_dim},
+                          CopyOpEntry{});
+}
+
+template <class TmeInternalType,
+          class CopyOp,
+          class GEngine, class GLayout,
+          class SLayout,
+          class VShape, class VStride>
+MUTE_HOST_RTC
+auto
+make_tme_copy_atom(CopyOp,
+                   Tensor<GEngine,GLayout> const& gtensor,       // Full GMEM Tensor
+                   SLayout                 const& slayout,       // CTA Tile of SMEM, potentially swizzled
+                   Layout<VShape,VStride>  const& cta_v_map)     // V: CTA val idx -> gmem mode
+{
+  //
+  // TME truncated layout
+  //
+
+  auto smem_swizzle = get_swizzle_portion(slayout);
+  auto smem_layout  = get_nonswizzle_portion(slayout);
+
+  auto tme_gbasis = detail::construct_tme_gbasis<TmeInternalType>(gtensor, smem_layout, cta_v_map);
+
+  auto [tme_desc, aux_params, copy_op] = detail::make_tme_copy_desc<TmeInternalType, CopyOp>(
+                                                                                     gtensor,
+                                                                                     tme_gbasis,
+                                                                                     smem_swizzle);
+  //
+  // Construct the Copy_Traits
+  //
+
+  constexpr int num_bits_per_tme = size(tme_gbasis) * sizeof_bits_v<TmeInternalType>;
+  using Traits = Copy_Traits<decltype(copy_op), mute::C<num_bits_per_tme>, decltype(aux_params)>;
+  using Atom   = Copy_Atom<Traits, typename GEngine::value_type>;
+
+
+  Traits tme_traits{tme_desc, aux_params};
+
+#if 0
+  print("num_bits_per_tme    : "); print(num_bits_per_tme); print("\n");
+  print("g_stride_bases      : "); print(tme_traits.aux_params_.g_stride_); print("\n");
+  print("smem block dim      : "); print(tme_traits.aux_params_.block_dim_); print("\n");
+#endif
+
+  // Return the Copy_Atom
+  return Atom{tme_traits};
+
+}
+
+template <class TmeInternalType,
+          class CopyOp,
+          class GEngine, class GLayout,
+          class SLayout,
+          class VShape, class VStride>
+MUTE_HOST_RTC
+auto
+make_tme_copy_tiled(CopyOp                   const& copy_op,
+                    Tensor<GEngine, GLayout> const& gtensor,    // Full GMEM Tensor
+                    SLayout                  const& slayout,    // CTA Tile of SMEM
+                    Layout<VShape, VStride>  const& cta_v_map)
+{
+  Copy_Atom atom = make_tme_copy_atom<TmeInternalType>(copy_op, gtensor, slayout, cta_v_map);
+
+  //
+  // Construct the TiledCopy
+  //
+
+  [[maybe_unused]] auto cta_tiler = product_each(shape(cta_v_map));
+
+  auto num_elems_per_tme = size<1>(typename decltype(atom)::RefLayout{}) / static_value<sizeof_bits<typename GEngine::value_type>>();
+
+  // smem idx -> smem coord
+  auto inv_smem_layout = right_inverse(get_nonswizzle_portion(slayout));
+  // CTA V -> smem_coord
+  auto layout_v = composition(inv_smem_layout, num_elems_per_tme);
+  // Scale that up to cover all of the smem_coords
+  auto layout_V = tile_to_shape(make_layout(layout_v), size(cta_v_map));
+
+  // Simplely construct thread layout.
+  // TME is a warp-level instruction and only the first active lane will issue.
+  auto layout_T = Layout<_1, _0>{};
+  // Combine with the T mapping
+  [[maybe_unused]] auto layout_TV = make_layout(layout_T, layout_V);
+
+#if 0
+  print("cta_tiler : "); print(cta_tiler); print("\n");
+  print("layout_v  : "); print(layout_v); print("\n");
+  print("layout_V  : "); print(layout_V); print("\n");
+  print("layout_T  : "); print(layout_T); print("\n");
+  print("layout_TV : "); print(layout_TV); print("\n");
+#endif
+
+  return TiledCopy<decltype(atom), decltype(layout_TV), decltype(cta_tiler)>{atom};
+
+}
+
+} // namespace detail
+
+
+template <class TmeInternalType = void,
+          class CopyOp,
+          class GEngine, class GLayout,
+          class SLayout,
+          class CTA_Tiler>
+MUTE_HOST_RTC
+auto
+make_tme_copy(CopyOp                   const& copy_op,
+              Tensor<GEngine, GLayout> const& gtensor,
+              SLayout                  const& slayout,
+              CTA_Tiler                const& cta_tiler)
+{
+  auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler);
+
+  using TmeType = conditional_t<is_same<void, TmeInternalType>::value, typename GEngine::value_type, TmeInternalType>;
+  return detail::make_tme_copy_tiled<TmeType>(copy_op,
+                                              gtensor, slayout,
+                                              cta_v_tile);
+}
+
+// Explicit defaulting
+template <class CopyOp,
+          class GEngine, class GLayout,
+          class SLayout>
+MUTE_HOST_RTC
+auto
+make_tme_copy(CopyOp                  const& copy_op,
+              Tensor<GEngine,GLayout> const& gtensor,
+              SLayout                 const& slayout)
+{
+  return make_tme_copy(copy_op, gtensor, slayout, product_each(shape(slayout)));
+}
+
+//////////////////////////////////////////////////
+// Experimental make TME Atom and TME Partitioner
+//////////////////////////////////////////////////
+
+
+template <class TmeInternalType = void,
+          class CopyOp,
+          class GEngine, class GLayout,
+          class SLayout,
+          class CTA_Tiler>
+MUTE_HOST_RTC
+auto
+make_tme_atom(CopyOp                   const& copy_op,
+              Tensor<GEngine, GLayout> const& gtensor,
+              SLayout                  const& slayout,
+              CTA_Tiler                const& cta_tiler)
+{
+  auto cta_v_tile = make_identity_layout(shape(gtensor)).compose(cta_tiler);
+
+  using TmeType = conditional_t<is_same<void, TmeInternalType>::value, typename GEngine::value_type, TmeInternalType>;
+  return detail::make_tme_copy_atom<TmeType>(copy_op,
+                                             gtensor, slayout,
+                                             cta_v_tile);
+}
+
+template <class... Args,
+          class WarpCoord,
+          class WarpLayout,
+          class SEngine, class SLayout,
+          class GEngine, class GLayout>
+MUTE_HOST_DEVICE
+auto
+tme_partition(Copy_Atom<Args...>       const& copy_atom,
+              WarpCoord                const& warp_coord,
+              WarpLayout               const& warp_layout,  // warp id
+              Tensor<SEngine, SLayout> const& stensor,      // SMEM Tensor (TMETile, Rest...)
+              Tensor<GEngine, GLayout> const& gtensor)      // GMEM Tensor (TMETile, Rest...)
+{
+  MUTE_STATIC_ASSERT_V(size<0>(stensor) == size<0>(gtensor));
+
+  // Invert the smem to get the largest contiguous vector in the smem layout
+  Layout inv_smem_layout = right_inverse(get_nonswizzle_portion(layout<0>(stensor)));
+
+  // Scale that up to conver all of the smem_coords
+  Layout layout_v = tile_to_shape(make_layout(inv_smem_layout), size<0>(stensor));
+
+  // Factor out the single-instruction portion
+  Layout tme_layout_v = make_layout(Int<Copy_Atom<Args...>::NumValSrc>{});
+  auto layout_V = logical_divide(layout_v, tme_layout_v);
+
+  // TME Iters must be divided by issue_warps
+  static_assert(size<1>(layout_V) % cosize_v<WarpLayout> == 0, "TME Iters must be divided by issue warps");
+
+  Layout tme_layout_v_ext = make_layout(Shape<Int<Copy_Atom<Args...>::NumValSrc>, Int<cosize_v<WarpLayout>>>{});
+  auto layout_V_ext = make_tile(flat_divide(layout_v, tme_layout_v_ext));
+
+  // Append with _ until we conver all Rest... modes
+  auto glayout_V = append<GLayout::rank>(layout_V_ext, _);
+  auto slayout_V = append<SLayout::rank>(layout_V_ext, _);
+  // Transform tile mode and coalesce
+  Tensor gtensor_v = coalesce(gtensor.compose(glayout_V), Shape<Shape<_1, _1>>{}); // ((TME, IssueWarps, TME Iter/IssueWarps), Rest...)
+  Tensor stensor_v = coalesce(stensor.compose(slayout_V), Shape<Shape<_1, _1>>{}); // ((TME, IssueWarps, TME Iter/IssueWarps), Rest...)
+
+  // Offset inside the TME-mode for the warp slice
+  auto warp_offset = warp_layout(warp_coord);
+  auto slice_coord = make_coord(make_coord(Underscore{}, warp_offset, Underscore{}));
+
+  auto scoord = append<SLayout::rank>(slice_coord, Underscore{});
+  auto gcoord = append<GLayout::rank>(slice_coord, Underscore{});
+
+  Tensor sresult = stensor_v(scoord);
+  Tensor gresult = gtensor_v(gcoord);
+
+  return mute::make_tuple(gresult, sresult);
+}
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31_tme_swizzle.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31_tme_swizzle.hpp
new file mode 100644
index 000000000..041a31218
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/copy_traits_mp31_tme_swizzle.hpp
@@ -0,0 +1,184 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/arch/copy_mp31_desc.hpp>
+#include <mute/swizzle_layout.hpp>
+
+
+namespace mute::detail {
+
+template <int B_, int M_, int S_, class T = uint8_t>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle_granularity(Swizzle<B_,M_,S_> _)
+{
+  // Recast to byte unit
+  using TmeSwizzle = decltype(recast_layout<T, uint8_t>(_));
+  constexpr int M = TmeSwizzle::num_base;
+
+  static_assert(4 <= M && M <= 6, "Unknown Swizzle Granularity.");
+  switch (M) {
+    case  4: return TME::SmemSwizzleGranularity::B16;
+    case  5: return TME::SmemSwizzleGranularity::B32;
+    case  6: return TME::SmemSwizzleGranularity::B64;
+  }
+}
+
+template <TME::SmemSwizzleGranularity sg>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle_base()
+{
+  static_assert(0 <= static_cast<uint8_t>(sg) && static_cast<uint8_t>(sg) <= 3, "Unknown Swizzle Granularity.");
+  switch (sg) {
+    case  TME::SmemSwizzleGranularity::B16:  return 4;
+    case  TME::SmemSwizzleGranularity::B32:  return 5;
+    case  TME::SmemSwizzleGranularity::B64:  return 6;
+  }
+}
+
+template <int B_, int M_, int S_, class T = uint8_t>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle_line(Swizzle<B_,M_,S_> _)
+{
+  // Recast to byte unit
+  using TmeSwizzle = decltype(recast_layout<T, uint8_t>(_));
+  constexpr int M = TmeSwizzle::num_base;
+  constexpr int S = TmeSwizzle::num_shft;
+
+  static_assert((M+S) == 8 || (M+S) == 7, "Unknown Swizzle Line.");
+  return (M + S) == 8 ? TME::SmemSwizzleLine::B256
+                      : TME::SmemSwizzleLine::B128;
+}
+
+template <TME::SmemSwizzleLine sl, int M>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle_shft()
+{
+  static_assert(4 <= M && M <= 6, "Unknown Swizzle Granularity.");
+  static_assert(0 <= static_cast<uint8_t>(sl) && static_cast<uint8_t>(sl) <= 1, "Unknown Swizzle Line.");
+  switch (sl) {
+    case  TME::SmemSwizzleLine::B256: return  8 - M;
+    case  TME::SmemSwizzleLine::B128: return  7 - M;
+  }
+}
+
+template <int B_, int M_, int S_, class T = uint8_t>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle_stride(Swizzle<B_,M_,S_> _)
+{
+  // Recast to byte unit
+  using TmeSwizzle = decltype(recast_layout<T, uint8_t>(_));
+  constexpr int B = TmeSwizzle::num_bits;
+  constexpr int M = TmeSwizzle::num_base;
+
+  static_assert(5 <= (M+B) && (M+B) <= 8, "Unknown Swizzle Stride.");
+  switch ((M+B)) {
+    case 5:  return TME::SmemSwizzleStride::B32;
+    case 6:  return TME::SmemSwizzleStride::B64;
+    case 7:  return TME::SmemSwizzleStride::B128;
+    case 8:  return TME::SmemSwizzleStride::B256;
+  }
+}
+
+template <TME::SmemSwizzleStride ss, int M>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle_bits()
+{
+  static_assert(4 <= M && M <= 6, "Unknown Swizzle Granularity.");
+  static_assert(0 <= static_cast<uint8_t>(ss) && static_cast<uint8_t>(ss) <= 3, "Unknown Swizzle Stride.");
+  switch (ss) {
+    case  TME::SmemSwizzleStride::B32:  return 5 - M;
+    case  TME::SmemSwizzleStride::B64:  return 6 - M;
+    case  TME::SmemSwizzleStride::B128: return 7 - M;
+    case  TME::SmemSwizzleStride::B256: return 8 - M;
+  }
+}
+
+
+template <int B_, int M_, int S_, class T = uint8_t>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle_attributes(Swizzle<B_,M_,S_> _)
+{
+  // Recast to byte unit
+  using TmeSwizzle = decltype(recast_layout<T, uint8_t>(_));
+  constexpr int B = TmeSwizzle::num_bits;
+  constexpr int M = TmeSwizzle::num_base;
+  constexpr int S = TmeSwizzle::num_shft;
+
+  // If B==0 then we don't do swizzle
+  if constexpr (B == 0) {
+    return mute::make_tuple(TME::SmemSwizzleGranularity::NONE,
+                            TME::SmemSwizzleStride::B256,
+                            TME::SmemSwizzleLine::B256);
+  }
+  else {
+    constexpr TME::SmemSwizzleGranularity sg = get_tme_swizzle_granularity(TmeSwizzle{});
+    constexpr TME::SmemSwizzleStride      ss = get_tme_swizzle_stride(TmeSwizzle{});
+    constexpr TME::SmemSwizzleLine        sl = get_tme_swizzle_line(TmeSwizzle{});
+
+    // Check validation SG < SS <= SL
+    static_assert(sg < ss && ss <= sl, "Invalid Smem Swizzle Pattern.");
+    return mute::make_tuple(sg, ss, sl);
+  }
+}
+
+template <class Layout, class T = uint8_t>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle_attributes(Layout const& layout) {
+  return get_tme_swizzle_attributes<T>(get_swizzle_portion(layout));
+}
+
+template <TME::SmemSwizzleLine sl, TME::SmemSwizzleStride ss, TME::SmemSwizzleGranularity sg>
+MUTE_HOST_DEVICE constexpr
+auto
+get_tme_swizzle() {
+  static_assert(sg < ss && ss <= sl, "Invalid Smem Swizzle Pattern.");
+  if constexpr (sg == TME::SmemSwizzleGranularity::NONE) {
+    return Swizzle<0, 4, 3>{};
+  }
+  else {
+    constexpr int M = get_tme_swizzle_base<sg>();
+    constexpr int B = get_tme_swizzle_bits<ss, M>();
+    constexpr int S = get_tme_swizzle_shft<sl, M>();
+    return Swizzle<B, M, S>{};
+  }
+}
+
+} // namespace mute::detail
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/mma_atom.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_atom.hpp
new file mode 100644
index 000000000..c3becd18f
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_atom.hpp
@@ -0,0 +1,948 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/arch/mma.hpp>
+
+#include <mute/atom/mma_traits.hpp>
+#include <mute/tensor.hpp>
+#include <mute/util/type_traits.hpp>
+
+namespace mute {
+
+template <class... Args>
+struct MMA_Atom;
+
+template <class MMAOperation>
+struct MMA_Atom<MMAOperation> : MMA_Atom<MMA_Traits<MMAOperation>>
+{};
+
+template <class MMAOperation, class... Args>
+struct MMA_Atom<MMA_Traits<MMAOperation, Args...>>
+  : MMA_Traits<MMAOperation, Args...>
+{
+  using MMA_Op = MMAOperation;
+  using Traits = MMA_Traits<MMAOperation, Args...>;
+
+  // Element value types from the MMA_Traits
+  using ValTypeD = typename Traits::ValTypeD;
+  using ValTypeA = typename Traits::ValTypeA;
+  using ValTypeB = typename Traits::ValTypeB;
+  using ValTypeC = typename Traits::ValTypeC;
+
+  // Thr-Val layouts from the MMA_Traits
+  using Shape_MNK  = typename Traits::Shape_MNK;
+  using ThrID      = typename Traits::ThrID;
+  using LayoutC_TV = typename Traits::CLayout;
+  using LayoutA_TV = typename Traits::ALayout;
+  using LayoutB_TV = typename Traits::BLayout;
+
+  // Fragment value types from the MMA_Traits (optional, defaults to Val type)
+  using FrgTypeD = typename detail::FrgTypeC_or_Default<Traits>::type;
+  using FrgTypeA = typename detail::FrgTypeA_or_Default<Traits>::type;
+  using FrgTypeB = typename detail::FrgTypeB_or_Default<Traits>::type;
+  using FrgTypeC = typename detail::FrgTypeC_or_Default<Traits>::type;
+
+  // Additional Trait parameters/transformations
+  template <class... TraitsArgs>
+  MUTE_HOST_DEVICE
+  auto
+  with(TraitsArgs&&... args) const {
+    auto traits = Traits::with(static_cast<TraitsArgs&&>(args)...);
+    return MMA_Atom<decltype(traits)>{traits};
+  }
+
+  //
+  // Tensor call interfaces
+  //
+
+  // Cast, check, and call fma
+  template <class TD, class DLayout,
+            class TA, class ALayout,
+            class TB, class BLayout,
+            class TC, class CLayout>
+  MUTE_HOST_DEVICE constexpr
+  void
+  call(Tensor<TD, DLayout>      & D,
+       Tensor<TA, ALayout> const& A,
+       Tensor<TB, BLayout> const& B,
+       Tensor<TC, CLayout> const& C) const
+  {
+    static_assert(DLayout::rank == 1, "Expected rank-1 D tensor");
+    static_assert(ALayout::rank == 1, "Expected rank-1 A tensor");
+    static_assert(BLayout::rank == 1, "Expected rank-1 B tensor");
+    static_assert(CLayout::rank == 1, "Expected rank-1 C tensor");
+
+    return mma_unpack(*this, D, A, B, C);
+  }
+
+  // Three arguments reproduces C
+  template <class TA, class ALayout,
+            class TB, class BLayout,
+            class TC, class CLayout>
+  MUTE_HOST_DEVICE constexpr
+  void
+  call(Tensor<TA, ALayout> const& A,
+       Tensor<TB, BLayout> const& B,
+       Tensor<TC, CLayout>      & C) const
+  {
+    return call(C, A, B, C);
+  }
+
+  //
+  // make_fragment_A|B|C
+  //   These functions are awkward as they expect already-partitioned tensors
+  //     resulting from a previous call to partition_A|B|C
+  //   The reasoning is that we can inspect the layout of the partitioned data
+  //     and attempt to match it in generated fragment to promote vectorization
+  //     when copying from partition to fragment.
+  //
+
+  template <class CTensor>
+  MUTE_HOST_DEVICE static constexpr
+  auto
+  make_fragment_C(CTensor&& ctensor)
+  {
+    // Check that this tensor is likely already partitioned
+    MUTE_STATIC_ASSERT_V(rank(ctensor) >= Int<3>{});  // VMN
+    MUTE_STATIC_ASSERT_V(size<0>(ctensor) == size<1>(LayoutC_TV{}));
+    // C is a bit special because we are after accumulators here
+    // The input/output type doesn't have to match the accumulator type
+    //static_assert(std::is_same<ValTypeC, typename remove_cvref_t<CTensor>::value_type>::value, "Expecting ValTypeC type");
+
+    // We'll never base the accumulator layout on the input tensor layout, so just return a FrgTypeC tensor
+    return make_tensor<FrgTypeC>(shape(ctensor));
+  }
+
+  template <class ATensor>
+  MUTE_HOST_DEVICE static constexpr
+  auto
+  make_fragment_A(ATensor&& atensor)
+  {
+    // Check that this tensor is likely already partitioned
+    MUTE_STATIC_ASSERT_V(rank(atensor) >= Int<3>{});  // VMK
+    MUTE_STATIC_ASSERT_V(size<0>(atensor) == size<1>(LayoutA_TV{}));
+
+    if constexpr (has_dereference<FrgTypeA>::value) {
+      // If the intended FrgTypeA is a view (of the current tensor), forward the whole
+      static_assert(is_same<ValTypeA, typename remove_cvref_t<ATensor>::value_type>::value
+                      , "Expecting ValTypeA type");
+      return make_tensor<FrgTypeA>(static_cast<ATensor&&>(atensor));
+    } else {
+      // Else, the intended FrgTypeA is a value type, construct a new tensor with a fragment layout
+      return make_fragment_like<FrgTypeA>(atensor);
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+
+  template <class BTensor>
+  MUTE_HOST_DEVICE static constexpr
+  auto
+  make_fragment_B(BTensor&& btensor)
+  {
+    // Check that this tensor is likely already partitioned
+    MUTE_STATIC_ASSERT_V(rank(btensor) >= Int<3>{});  // VNK
+    MUTE_STATIC_ASSERT_V(size<0>(btensor) == size<1>(LayoutB_TV{}));
+
+    if constexpr (has_dereference<FrgTypeB>::value) {
+      // If the intended FrgTypeB is a view (of the current tensor), forward the whole
+      static_assert(is_same<ValTypeB, typename remove_cvref_t<BTensor>::value_type>::value
+                      , "Expecting ValTypeB type");
+      return make_tensor<FrgTypeB>(static_cast<BTensor&&>(btensor));
+    } else {
+      // Else, the intended FrgTypeB is a value type, construct a new tensor with a fragment layout
+      return make_fragment_like<FrgTypeB>(btensor);
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+};
+
+//
+// A tiling of mma atoms
+//
+
+template <class TiledMMA, class ThrCoord>
+struct ThrMMA;
+
+// @tparam MMA_Atom The MMA_Atom to use in the TiledMMA
+// @tparam AtomLayoutMNK The MNK-tiling of the Atom to be performed.
+// @tparam PermuationsMNK Permutations to apply to each MNK-mode before tiling for the Atom.
+template <class MMA_Atom,
+          class AtomLayoutMNK,
+          class PermutationMNK = Tile<Underscore,Underscore,Underscore>>
+struct TiledMMA : MMA_Atom
+{
+  using Atom           = MMA_Atom;
+  using AtomShape_MNK  = typename MMA_Atom::Shape_MNK;
+  using AtomThrID      = typename MMA_Atom::ThrID;
+  using AtomLayoutC_TV = typename MMA_Atom::LayoutC_TV;
+  using AtomLayoutA_TV = typename MMA_Atom::LayoutA_TV;
+  using AtomLayoutB_TV = typename MMA_Atom::LayoutB_TV;
+
+  static_assert(   rank_v<AtomLayoutMNK>  == 3,   "TiledMMA requires rank-3 AtomLayoutMNK");
+  static_assert(   rank_v<PermutationMNK> == 3,   "TiledMMA requires rank-3 PermutationMNK");
+  static_assert( is_tuple<PermutationMNK>::value, "TiledMMA requires independent permutations of MNK.");
+  static_assert(is_static<PermutationMNK>::value, "TiledMMA requires static permutations of MNK.");
+
+  using ThrLayoutVMNK = decltype(tiled_product(AtomThrID{}, AtomLayoutMNK{}));
+  ThrLayoutVMNK thr_layout_vmnk_;
+
+  MUTE_HOST_DEVICE constexpr
+  TiledMMA(MMA_Atom const& mma_atom = {}, AtomLayoutMNK const& thr_layout_mnk = {})
+    : MMA_Atom(mma_atom),
+      thr_layout_vmnk_(tiled_product(AtomThrID{}, thr_layout_mnk)) {}
+
+  MUTE_HOST_DEVICE constexpr auto
+  get_thr_layout_vmnk() const {
+    return thr_layout_vmnk_;
+  }
+
+  // Tile a tensor or a layout from shape
+  //   (M,N,...)
+  // to shape
+  //   ((ThrV,(ThrM,ThrN)),(FrgV,(RestM,RestN,...)))
+  // where
+  //   ThrV:  The threads local to an MMA. layout<0>(ThrLayoutVMNK): ThrV -> thread_idx
+  //   ThrM:  The threads tiled in M.      layout<1>(ThrLayoutVMNK): ThrM -> thread_idx
+  //   ThrN:  The threads tiled in N.      layout<2>(ThrLayoutVMNK): ThrN -> thread_idx
+  //   FrgV:  The values local to an MMA.
+  //   RestM: The values tiled in M.
+  //   RestN: The values tiled in N.
+  template <class CTensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  thrfrg_C(CTensor&& ctensor) const
+  {
+    MUTE_STATIC_ASSERT_V(rank(ctensor) >= Int<2>{});
+    //MUTE_STATIC_ASSERT_V(size<0>(ctensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
+    //MUTE_STATIC_ASSERT_V(size<1>(ctensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
+
+    // Reorder the tensor for the TiledAtom
+    auto t_tile = make_tile(get<0>(PermutationMNK{}),
+                            get<1>(PermutationMNK{}));
+    auto t_tensor = logical_divide(ctensor, t_tile);                 // (PermM,PermN)
+
+    // Tile the tensor for the Atom
+    auto a_tile = make_tile(make_layout(size<0>(AtomShape_MNK{})),
+                            make_layout(size<1>(AtomShape_MNK{})));
+    auto a_tensor = zipped_divide(t_tensor, a_tile);                 // ((AtomM,AtomN),(RestM,RestN))
+
+    // Transform the Atom mode from (M,K) to (Thr,Val)
+    auto tv_tensor = a_tensor.compose(AtomLayoutC_TV{},_);           // ((ThrV,FrgV),(RestM,RestN))
+
+    // Tile the tensor for the C-threads
+    auto thr_tile = make_tile(_,
+                              make_tile(make_layout(size<1>(thr_layout_vmnk_)),
+                                        make_layout(size<2>(thr_layout_vmnk_))));
+    auto thr_tensor = zipped_divide(tv_tensor, thr_tile);            // ((ThrV,(ThrM,ThrN)),(FrgV,(RestM,RestN)))
+
+    return thr_tensor;
+  }
+
+  // Tile a tensor or a layout from shape
+  //   (M,K,...)
+  // to shape
+  //   ((ThrV,(ThrM,ThrK)),(FrgV,(RestM,RestK,...)))
+  // where
+  //   ThrV: The threads local to an MMA. layout<0>(ThrLayoutVMNK): ThrV -> thread_idx
+  //   ThrM: The threads tiled in M.      layout<1>(ThrLayoutVMNK): ThrM -> thread_idx
+  //   ThrK: The threads tiled in K.      layout<3>(ThrLayoutVMNK): ThrK -> thread_idx
+  //   FrgV:  The values local to an MMA.
+  //   RestM: The values tiled in M.
+  //   RestK: The values tiled in K.
+  template <class ATensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  thrfrg_A(ATensor&& atensor) const
+  {
+    MUTE_STATIC_ASSERT_V(rank(atensor) >= Int<2>{});
+    //MUTE_STATIC_ASSERT_V(size<0>(atensor) % size<0>(TiledShape_MNK{}) == Int<0>{});
+    //MUTE_STATIC_ASSERT_V(size<1>(atensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
+
+    // Reorder the tensor for the TiledAtom
+    auto t_tile = make_tile(get<0>(PermutationMNK{}),
+                            get<2>(PermutationMNK{}));
+    auto t_tensor = logical_divide(atensor, t_tile);                 // (PermM,PermK)
+
+    // Tile the tensor for the Atom
+    auto a_tile = make_tile(make_layout(size<0>(AtomShape_MNK{})),
+                            make_layout(size<2>(AtomShape_MNK{})));
+    auto a_tensor = zipped_divide(t_tensor, a_tile);                 // ((AtomM,AtomK),(RestM,RestK))
+
+    // Transform the Atom mode from (M,K) to (Thr,Val)
+    auto tv_tensor = a_tensor.compose(AtomLayoutA_TV{},_);           // ((ThrV,FrgV),(RestM,RestK))
+
+    // Tile the tensor for the Thread
+    auto thr_tile = make_tile(_,
+                              make_tile(make_layout(size<1>(thr_layout_vmnk_)),
+                                        make_layout(size<3>(thr_layout_vmnk_))));
+    auto thr_tensor = zipped_divide(tv_tensor, thr_tile);            // ((ThrV,(ThrM,ThrK)),(FrgV,(RestM,RestK)))
+
+    return thr_tensor;
+  }
+
+  // Tile a tensor or a layout from shape
+  //   (N,K,...)
+  // to shape
+  //   ((ThrV,(ThrN,ThrK)),(FrgV,(RestN,RestK,...)))
+  // where
+  //   ThrV: The threads local to an MMA. layout<0>(ThrLayoutVMNK): ThrV -> thread_idx
+  //   ThrN: The threads tiled in N.      layout<2>(ThrLayoutVMNK): ThrN -> thread_idx
+  //   ThrK: The threads tiled in K.      layout<3>(ThrLayoutVMNK): ThrK -> thread_idx
+  //   FrgV:  The values local to an MMA.
+  //   RestN: The values tiled in N.
+  //   RestK: The values tiled in K.
+  template <class BTensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  thrfrg_B(BTensor&& btensor) const
+  {
+    MUTE_STATIC_ASSERT_V(rank(btensor) >= Int<2>{});
+    //MUTE_STATIC_ASSERT_V(size<0>(btensor) % size<1>(TiledShape_MNK{}) == Int<0>{});
+    //MUTE_STATIC_ASSERT_V(size<1>(btensor) % size<2>(TiledShape_MNK{}) == Int<0>{});
+
+    // Reorder the tensor for the TiledAtom
+    auto t_tile = make_tile(get<1>(PermutationMNK{}),
+                            get<2>(PermutationMNK{}));
+    auto t_tensor = logical_divide(btensor, t_tile);                 // (PermN,PermK)
+
+    // Tile the tensor for the Atom
+    auto a_tile = make_tile(make_layout(size<1>(AtomShape_MNK{})),
+                            make_layout(size<2>(AtomShape_MNK{})));
+    auto a_tensor = zipped_divide(t_tensor, a_tile);                 // ((AtomN,AtomK),(RestN,RestK))
+
+    // Transform the Atom mode from (M,K) to (Thr,Val)
+    auto tv_tensor = a_tensor.compose(AtomLayoutB_TV{},_);           // ((ThrV,FrgV),(RestN,RestK))
+
+    // Tile the tensor for the Thread
+    auto thr_tile = make_tile(_,
+                              make_tile(make_layout(size<2>(thr_layout_vmnk_)),
+                                        make_layout(size<3>(thr_layout_vmnk_))));
+    auto thr_tensor = zipped_divide(tv_tensor, thr_tile);            // ((ThrV,(ThrN,ThrK)),(FrgV,(RestN,RestK)))
+
+    return thr_tensor;
+  }
+
+  template <class ThrIdx,
+            __MUTE_REQUIRES(is_integral<ThrIdx>::value)>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_slice(ThrIdx const& thr_idx) const
+  {
+    auto thr_vmnk = thr_layout_vmnk_.get_flat_coord(thr_idx);
+    return ThrMMA<TiledMMA, decltype(thr_vmnk)>{*this, thr_vmnk};
+  }
+
+  template <class ThrIdx,
+            __MUTE_REQUIRES(is_integral<ThrIdx>::value)>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_thread_slice(ThrIdx const& thr_idx) const
+  {
+    return get_slice(thr_idx);
+  }
+
+  //
+  // Utility for printing and visualization
+  //
+
+  // The size of the MNK-mode
+  template <int I>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  tile_size_mnk() const {
+    static_assert(0 <= I && I < 3);
+    auto core_size = size<I>(AtomShape_MNK{}) * size<I+1>(get_thr_layout_vmnk());
+    [[maybe_unused]] auto perm_size = size<I>(PermutationMNK{});
+    if constexpr (is_underscore<decltype(perm_size)>::value) {
+      return core_size;
+    } else {
+      return mute::max(core_size, perm_size);
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_layoutC_MN() const
+  {
+    // (M,N) -> (M,N)
+    auto ref_C = make_layout(make_shape(tile_size_mnk<0>(), tile_size_mnk<1>()));
+    // (cthrid,val) -> (M,N)
+    auto layoutC_TV = thrfrg_C(ref_C);
+    // (M,N) -> (cthrid,frg)
+    auto layoutC_MN = right_inverse(layoutC_TV).with_shape(shape(ref_C));
+
+    // cthrid = (v,m,n) -> thr_idx
+    auto thrID_C = thr_layout_vmnk_(_,_,_,Int<0>{});
+
+    return mute::make_tuple(layoutC_MN, thrID_C);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_layoutC_TV() const
+  {
+    // (M,N) -> (M,N)
+    auto ref_C = make_layout(make_shape(tile_size_mnk<0>(), tile_size_mnk<1>()));
+    // (cthrid,val) -> (M,N)
+    auto layoutC_TV = thrfrg_C(ref_C);
+
+    // thr_idx -> (ThrV,ThrM,ThrN,ThrK)
+    auto thridx_2_thrid = right_inverse(thr_layout_vmnk_);
+
+    // (thr_idx,val) -> (M,N)
+    return layoutC_TV.compose(thridx_2_thrid, _);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_layoutA_MK() const
+  {
+    // (M,K) -> (M,K)
+    auto ref_A = make_layout(make_shape(tile_size_mnk<0>(), tile_size_mnk<2>()));
+    // (athrid,val) -> (M,K)
+    auto layoutA_TV = thrfrg_A(ref_A);
+    // (M,K) -> (athrid,frg)
+    auto layoutA_MK = right_inverse(layoutA_TV).with_shape(shape(ref_A));
+
+    // athrid = (v,m,k) -> thr_idx
+    auto thrID_A = thr_layout_vmnk_(_,_,Int<0>{},_);
+
+    return mute::make_tuple(layoutA_MK, thrID_A);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_layoutA_TV() const
+  {
+    // (M,K) -> (M,K)
+    auto ref_A = make_layout(make_shape(tile_size_mnk<0>(), tile_size_mnk<2>()));
+    // (athrid,val) -> (M,K)
+    auto layoutA_TV = thrfrg_A(ref_A);
+
+    // (ThrV,(ThrM,ThrK)) -> (ThrV,(ThrM,ThrN,ThrK))
+    auto atile = make_tile(_,
+                           make_tile(make_layout(make_shape (size<1>(thr_layout_vmnk_), size<2>(thr_layout_vmnk_)),
+                                                 make_stride(               Int<1>{} ,                Int<0>{} )),
+                                     _));
+
+    // thr_idx -> (ThrV,ThrM,ThrN,ThrK)
+    auto thridx_2_thrid = right_inverse(thr_layout_vmnk_);
+
+    // (thr_idx,val) -> (M,K)
+    return thrfrg_A(ref_A).compose(atile, _).compose(thridx_2_thrid, _);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_layoutB_NK() const
+  {
+    // (N,K) -> (N,K)
+    auto ref_B = make_layout(make_shape(tile_size_mnk<1>(), tile_size_mnk<2>()));
+    // (bthrid,val) -> (N,K)
+    auto layoutB_TV = thrfrg_B(ref_B);
+    // (N,K) -> (bthrid,frg)
+    auto layoutB_NK = right_inverse(layoutB_TV).with_shape(shape(ref_B));
+
+    // bthrid = (v,n,k) -> thr_idx
+    auto thrID_B = thr_layout_vmnk_(_,Int<0>{},_,_);
+
+    return mute::make_tuple(layoutB_NK, thrID_B);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_layoutB_TV() const
+  {
+    // (N,K) -> (N,K)
+    auto ref_B = make_layout(make_shape(tile_size_mnk<1>(), tile_size_mnk<2>()));
+    // (bthrid,val) -> (N,K)
+    auto layoutB_TV = thrfrg_B(ref_B);
+
+    // (ThrV,(ThrM,ThrK)) -> (ThrV,(ThrM,ThrN,ThrK))
+    auto btile = make_tile(_,
+                           make_tile(make_layout(make_shape (size<1>(thr_layout_vmnk_), size<2>(thr_layout_vmnk_)),
+                                                 make_stride(               Int<0>{} ,                Int<1>{} )),
+                                     _));
+
+    // thr_idx -> (ThrV,ThrM,ThrN,ThrK)
+    auto thridx_2_thrid = right_inverse(thr_layout_vmnk_);
+
+    // (thr_idx,val) -> (N,K)
+    return thrfrg_B(ref_B).compose(btile, _).compose(thridx_2_thrid, _);
+  }
+};
+
+template <class TiledMMA, class ThrVMNK>
+struct ThrMMA : TiledMMA
+{
+  ThrVMNK thr_vmnk_;
+
+  template <class CTensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  partition_C(CTensor&& ctensor) const
+  {
+    auto thr_tensor = make_tensor(static_cast<CTensor&&>(ctensor).data(), this->thrfrg_C(ctensor.layout()));
+
+    auto thr_vmn = make_coord(get<0>(thr_vmnk_), make_coord(get<1>(thr_vmnk_), get<2>(thr_vmnk_)));
+    return thr_tensor(thr_vmn, make_coord(_, repeat<rank<1,1>(thr_tensor)>(_)));
+  }
+
+  template <class ATensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  partition_A(ATensor&& atensor) const
+  {
+    auto thr_tensor = make_tensor(static_cast<ATensor&&>(atensor).data(), this->thrfrg_A(atensor.layout()));
+
+    auto thr_vmk = make_coord(get<0>(thr_vmnk_), make_coord(get<1>(thr_vmnk_), get<3>(thr_vmnk_)));
+    return thr_tensor(thr_vmk, make_coord(_, repeat<rank<1,1>(thr_tensor)>(_)));
+  }
+
+  template <class BTensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  partition_B(BTensor&& btensor) const
+  {
+    auto thr_tensor = make_tensor(static_cast<BTensor&&>(btensor).data(), this->thrfrg_B(btensor.layout()));
+
+    auto thr_vnk = make_coord(get<0>(thr_vmnk_), make_coord(get<2>(thr_vmnk_), get<3>(thr_vmnk_)));
+    return thr_tensor(thr_vnk, make_coord(_, repeat<rank<1,1>(thr_tensor)>(_)));
+  }
+
+  template <class CTensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  partition_fragment_C(CTensor&& ctensor) const
+  {
+    return TiledMMA::make_fragment_C(partition_C(ctensor));
+  }
+
+  template <class ATensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  partition_fragment_A(ATensor&& atensor) const
+  {
+    return TiledMMA::make_fragment_A(partition_A(atensor));
+  }
+
+  template <class BTensor>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  partition_fragment_B(BTensor&& btensor) const
+  {
+    return TiledMMA::make_fragment_B(partition_B(btensor));
+  }
+};
+
+//
+// These tile the MMA_Atom as a whole
+//
+
+template <class MMA_Op,
+          class MMAThrLayout = Layout<Shape<_1,_1,_1>>,
+          class Permutations = Tile<Underscore,Underscore,Underscore>>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tiled_mma(MMA_Atom<MMA_Op> const& mma_atom,
+               MMAThrLayout     const& thr_layout   = {},
+               Permutations     const& permutations = {})
+{
+  auto thr_layout_mnk  = append<3>(thr_layout, Layout<_1,_0>{});
+  auto permutation_mnk = append<3>(permutations, _);
+
+  return TiledMMA<MMA_Atom<MMA_Op>,
+                  decltype(thr_layout_mnk),
+                  decltype(permutation_mnk)>{mma_atom, thr_layout_mnk};
+}
+
+template <class MMA_Op,
+          class MMAThrLayout = Layout<Shape<_1,_1,_1>>,
+          class Permutations = Tile<Underscore,Underscore,Underscore>>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tiled_mma(MMA_Op       const&,
+               MMAThrLayout const& thr_layout   = {},
+               Permutations const& permutations = {})
+{
+  // Attempt to wrap in an MMA_Atom<> and forward
+  return make_tiled_mma(MMA_Atom<MMA_Op>{}, thr_layout, permutations);
+}
+
+//
+// partition_fragment_C -- static context
+//
+
+template <class... Args, class Shape_MN>
+MUTE_HOST_DEVICE constexpr
+auto
+partition_shape_C(TiledMMA<Args...> const& mma, Shape_MN const& shape_MN)
+{
+  constexpr int R = rank_v<Shape_MN>;
+  static_assert(R >= 2, "Must have at least rank-2");
+  auto atomMNK = typename TiledMMA<Args...>::AtomShape_MNK{};
+  auto thrVMNK = typename TiledMMA<Args...>::ThrLayoutVMNK{};
+  auto V = shape<1>(typename TiledMMA<Args...>::AtomLayoutC_TV{});
+  auto M = shape_div(size<0>(shape_MN), size<0>(atomMNK) * size<1>(thrVMNK));
+  auto N = shape_div(size<1>(shape_MN), size<1>(atomMNK) * size<2>(thrVMNK));
+  return mute::tuple_cat(make_shape(V,M,N), take<2,R>(shape_MN));
+}
+
+template <class... Args, class Shape_MN>
+MUTE_HOST_DEVICE constexpr
+auto
+partition_fragment_C(TiledMMA<Args...> const& mma, Shape_MN const& shapeMN)
+{
+  return make_tensor<typename TiledMMA<Args...>::FrgTypeC>(partition_shape_C(mma, shapeMN));
+}
+
+// partition_fragment_A and partition_fragment_B often depend on the
+//   layout of A and B and/or the thread_idx that is requesting the partition.
+// For these reasons, they should not be used in a static context.
+// See TiledMMA::get_slice(thr_idx).partition_fragment_A(tensorA) instead.
+
+template <class... Args, class Shape_MK>
+MUTE_HOST_DEVICE constexpr
+auto
+partition_shape_A(TiledMMA<Args...> const& mma, Shape_MK const& shape_MK)
+{
+  constexpr int R = rank_v<Shape_MK>;
+  static_assert(R >= 2, "Must have at least rank-2");
+  auto atomMNK = typename TiledMMA<Args...>::AtomShape_MNK{};
+  auto thrVMNK = typename TiledMMA<Args...>::ThrLayoutVMNK{};
+  auto V = shape<1>(typename TiledMMA<Args...>::AtomLayoutA_TV{});
+  auto M = shape_div(size<0>(shape_MK), size<0>(atomMNK) * size<1>(thrVMNK));
+  auto K = shape_div(size<1>(shape_MK), size<2>(atomMNK) * size<3>(thrVMNK));
+  return mute::tuple_cat(make_shape(V,M,K), take<2,R>(shape_MK));
+}
+
+template <class... Args, class Shape_NK>
+MUTE_HOST_DEVICE constexpr
+auto
+partition_shape_B(TiledMMA<Args...> const& mma, Shape_NK const& shape_NK)
+{
+  constexpr int R = rank_v<Shape_NK>;
+  static_assert(R >= 2, "Must have at least rank-2");
+  auto atomMNK = typename TiledMMA<Args...>::AtomShape_MNK{};
+  auto thrVMNK = typename TiledMMA<Args...>::ThrLayoutVMNK{};
+  auto V = shape<1>(typename TiledMMA<Args...>::AtomLayoutB_TV{});
+  auto N = shape_div(size<0>(shape_NK), size<1>(atomMNK) * size<2>(thrVMNK));
+  auto K = shape_div(size<1>(shape_NK), size<2>(atomMNK) * size<3>(thrVMNK));
+  return mute::tuple_cat(make_shape(V,N,K), take<2,R>(shape_NK));
+}
+
+//
+// Size
+//
+
+template <int I, class... Args>
+MUTE_HOST_DEVICE constexpr
+auto
+tile_size(TiledMMA<Args...> const& mma)
+{
+  return mma.template tile_size_mnk<I>();
+}
+
+template <class... Args>
+MUTE_HOST_DEVICE constexpr
+auto
+tile_shape(TiledMMA<Args...> const& mma)
+{
+  return make_shape(tile_size<0>(mma), tile_size<1>(mma), tile_size<2>(mma));
+}
+
+// Deprecate?
+template <int... I, class... Args>
+MUTE_HOST_DEVICE constexpr
+auto
+size(TiledMMA<Args...> const& mma)
+{
+  return size<I...>(mma.get_thr_layout_vmnk());
+}
+
+// Alias
+template <int... I, class... Args>
+MUTE_HOST_DEVICE constexpr
+auto
+thr_size(TiledMMA<Args...> const& mma)
+{
+  return size<I...>(mma.get_thr_layout_vmnk());
+}
+
+//
+// Display utilities
+//
+
+template <class... Args>
+MUTE_HOST_DEVICE
+void
+print(MMA_Atom<MMA_Traits<Args...>> const&)
+{
+  using Atom = MMA_Atom<MMA_Traits<Args...>>;
+  print("MMA_Atom\n");
+  print("  ThrID:      "); print(typename Atom::ThrID{});      print("\n");
+  print("  Shape_MNK:  "); print(typename Atom::Shape_MNK{});  print("\n");
+  print("  LayoutA_TV: "); print(typename Atom::LayoutA_TV{}); print("\n");
+  print("  LayoutB_TV: "); print(typename Atom::LayoutB_TV{}); print("\n");
+  print("  LayoutC_TV: "); print(typename Atom::LayoutC_TV{}); print("\n");
+}
+
+template <class Atom, class TiledThr, class TiledPerm>
+MUTE_HOST_DEVICE
+void
+print(TiledMMA<Atom, TiledThr, TiledPerm> const& mma)
+{
+  print("TiledMMA\n");
+  print("  ThrLayoutVMNK:  "); print(mma.get_thr_layout_vmnk());  print("\n");
+  print("  PermutationMNK: "); print(TiledPerm{}); print("\n");
+  print(static_cast<Atom const&>(mma));
+}
+
+template <class TiledMMA, class ThrVMNK>
+MUTE_HOST_DEVICE
+void
+print(ThrMMA<TiledMMA, ThrVMNK> const& thr_mma)
+{
+  print("ThrMMA\n");
+  print("  Thr VMNK: "); print(thr_mma.thr_vmnk_); print("\n");
+  print(static_cast<TiledMMA>(thr_mma));
+}
+
+template <class... Args>
+MUTE_HOST_DEVICE
+void
+print_latex(MMA_Atom<Args...> const& mma_atom)
+{
+  print_latex(make_tiled_mma(mma_atom));
+}
+
+template <class... Args>
+MUTE_HOST_DEVICE
+void
+print_latex(TiledMMA<Args...> const& mma)
+{
+  auto layout_and_thrid_C = mma.get_layoutC_MN();
+  auto layoutC_MN = get<0>(layout_and_thrid_C);
+  auto thrID_C    = get<1>(layout_and_thrid_C);
+
+  auto layout_and_thrid_A = mma.get_layoutA_MK();
+  auto layoutA_MK = get<0>(layout_and_thrid_A);
+  auto thrID_A    = get<1>(layout_and_thrid_A);
+
+  auto layout_and_thrid_B = mma.get_layoutB_NK();
+  auto layoutB_NK = get<0>(layout_and_thrid_B);
+  auto thrID_B    = get<1>(layout_and_thrid_B);
+
+  print_latex_mma(layoutC_MN, thrID_C,
+                  layoutA_MK, thrID_A,
+                  layoutB_NK, thrID_B);
+}
+
+// MNK MMA Layout to console printer
+template <class LayoutC, class ThrIDC,
+          class LayoutA, class ThrIDA,
+          class LayoutB, class ThrIDB>
+MUTE_HOST_DEVICE
+void
+print_layout_mma(LayoutC const& C, ThrIDC const& TC,  // (m,n) -> (tid,vid)  and  tid -> thr_idx
+                 LayoutA const& A, ThrIDA const& TA,  // (m,k) -> (tid,vid)  and  tid -> thr_idx
+                 LayoutB const& B, ThrIDB const& TB)  // (n,k) -> (tid,vid)  and  tid -> thr_idx
+{
+  MUTE_STATIC_ASSERT_V(rank(C) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(A) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(B) == Int<2>{});
+
+  assert(size<0>(A) == size<0>(C));
+  assert(size<0>(B) == size<1>(C));
+  assert(size<1>(A) == size<1>(B));
+
+  int a_width = size<1>(A) * 6 + 4;
+
+  // Print out B (white-shifted) k-by-n
+  for (int k = 0; k < size<1>(B); ++k) {
+    // Header
+    printf("%*s", a_width, "");
+    for (int n = 0; n < size<0>(B); ++n) printf("+-----");
+    printf("+\n");
+    // Values
+    printf("%*s", a_width, "");
+    for (int n = 0; n < size<0>(B); ++n) printf("|T%02dV%1d", int(TB(B(n,k) % size(TB))), int(B(n,k) / size(TB)));
+    printf("|\n");
+  }
+  // Footer
+  printf("%*s", a_width, "");
+  for (int n = 0; n < size<0>(B); ++n) printf("+-----");
+  printf("+\n\n");
+
+  // Print out A m-by-k and C m-by-n
+  for (int m = 0; m < size<0>(A); ++m) {
+    // Header
+    for (int k = 0; k < size<1>(A); ++k) printf("+-----");
+    printf("+   ");
+    for (int n = 0; n < size<1>(C); ++n) printf("+-----");
+    printf("+\n");
+    // Values
+    for (int k = 0; k < size<1>(A); ++k) printf("|T%02dV%1d", int(TA(A(m,k) % size(TA))), int(A(m,k) / size(TA)));
+    printf("|   ");
+    for (int n = 0; n < size<1>(C); ++n) printf("|T%02dV%1d", int(TC(C(m,n) % size(TC))), int(C(m,n) / size(TC)));
+    printf("|\n");
+  }
+  // Footer
+  for (int k = 0; k < size<1>(A); ++k) printf("+-----");
+  printf("+   ");
+  for (int n = 0; n < size<1>(C); ++n) printf("+-----");
+  printf("+\n");
+}
+
+// MNK MMA Layout to Latex TIKZ -- 8-value color coded by thread
+template <class LayoutC, class ThrIDC,
+          class LayoutA, class ThrIDA,
+          class LayoutB, class ThrIDB>
+MUTE_HOST_DEVICE
+void
+print_latex_mma(LayoutC const& C, ThrIDC const& TC,  // (m,n) -> (tid,vid)  and  tid -> thr_idx
+                LayoutA const& A, ThrIDA const& TA,  // (m,k) -> (tid,vid)  and  tid -> thr_idx
+                LayoutB const& B, ThrIDB const& TB)  // (n,k) -> (tid,vid)  and  tid -> thr_idx
+{
+  MUTE_STATIC_ASSERT_V(rank(C) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(A) == Int<2>{});
+  MUTE_STATIC_ASSERT_V(rank(B) == Int<2>{});
+
+  assert(size<0>(A) == size<0>(C));
+  assert(size<0>(B) == size<1>(C));
+  assert(size<1>(A) == size<1>(B));
+
+  char const* latex_header =
+      "\\documentclass{standalone}\n"
+      "\\usepackage{tikz}\n"
+      "\\usetikzlibrary{external}\n"
+      "\\tikzexternalize\n"
+      "\\begin{document}\n"
+      "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},box/.style={rectangle,draw=black,thick,minimum size=1cm,anchor=center}]\n\n";
+  char const* latex_footer =
+      "\\end{tikzpicture}\n"
+      "\\end{document}\n";
+
+  char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}",
+                              "{rgb,255:red,175;green,255;blue,175}",
+                              "{rgb,255:red,255;green,255;blue,175}",
+                              "{rgb,255:red,255;green,175;blue,175}",
+                              "{rgb,255:red,210;green,210;blue,255}",
+                              "{rgb,255:red,210;green,255;blue,210}",
+                              "{rgb,255:red,255;green,255;blue,210}",
+                              "{rgb,255:red,255;green,210;blue,210}"};
+
+  // Header
+  printf("%% LayoutC: "); print(C);  printf("\n");
+  printf("%% ThrIDC : "); print(TC); printf("\n");
+  printf("%% LayoutA: "); print(A);  printf("\n");
+  printf("%% ThrIDA : "); print(TA); printf("\n");
+  printf("%% LayoutB: "); print(B);  printf("\n");
+  printf("%% ThrIDB : "); print(TB); printf("\n\n");
+
+  printf(latex_header);
+
+  // C starting at 0,0
+  for (int m = 0; m < size<0>(C); ++m) {
+    for (int n = 0; n < size<1>(C); ++n) {
+      int thrid   = C(m,n) % size(TC);
+      int val_idx = C(m,n) / size(TC);
+      int thr_idx = TC(thrid);
+
+      printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
+             color_map[thr_idx % 8],
+             m, n,
+             thr_idx, val_idx);
+    }
+  }
+
+  // A starting at 0,-size<1>(A)-1
+  for (int m = 0; m < size<0>(A); ++m) {
+    for (int k = 0; k < size<1>(A); ++k) {
+      int thrid   = A(m,k) % size(TA);
+      int val_idx = A(m,k) / size(TA);
+      int thr_idx = TA(thrid);
+
+      printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
+             color_map[thr_idx % 8],
+             m, k-1-size<1>(A),
+             thr_idx, val_idx);
+    }
+  }
+
+  // B starting at -size<1>(B)-1,0
+  for (int n = 0; n < size<0>(B); ++n) {
+    for (int k = 0; k < size<1>(B); ++k) {
+      int thrid   = B(n,k) % size(TB);
+      int val_idx = B(n,k) / size(TB);
+      int thr_idx = TB(thrid);
+
+      printf("\\node[box,fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
+             color_map[thr_idx % 8],
+             k-1-size<1>(B), n,
+             thr_idx, val_idx);
+    }
+  }
+
+  // A labels
+  for (int m = 0, k = -1; m < size<0>(A); ++m) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), m);
+  }
+  for (int k = 0, m = -1; k < size<1>(A); ++k) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", m, k-1-size<1>(A), k);
+  }
+  // B labels
+  for (int n = 0, k = -1; n < size<0>(B); ++n) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, n);
+  }
+  for (int k = 0, n = -1; k < size<1>(B); ++k) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", k-1-size<1>(B), n, k);
+  }
+
+  // Footer
+  printf(latex_footer);
+}
+
+} // namespace mute
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include <mute/atom/mma_traits_mp22.hpp>
+#include <mute/atom/mma_traits_mp31_sqmma.hpp>
+#include <mute/atom/mma_traits_mp31.hpp>
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits.hpp
new file mode 100644
index 000000000..f47f420a2
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits.hpp
@@ -0,0 +1,229 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/arch/mma.hpp>
+
+#include <mute/tensor.hpp>
+
+namespace mute
+{
+
+namespace detail {
+
+template <class X, class = void>
+struct supports_output_scaling { static constexpr bool value = false; };
+
+template <class X>
+struct supports_output_scaling<X, void_t<decltype(declval<X>().accumulate_)>> { static constexpr bool value = true; };
+
+} // end namespace detail
+
+/**
+ * concept MMA_Traits
+ * {
+ *   using ValTypeD =  // Logical A-value type
+ *   using ValTypeA =  // Logical B-value type
+ *   using ValTypeB =  // Logical C-value type
+ *   using ValTypeC =  // Logical D-value type    (NOTE: Not used? Assumed == ValTypeD)
+ *
+ *   using FrgTypeA =  // A-type consumed by MMA  (if ommitted, same as ValTypeA)
+ *   using FrgTypeB =  // B_type consumed by MMA  (if ommitted, same as ValTypeB)
+ *   using FrgTypeC =  // C_type consumed by MMA  (if ommitted, same as ValTypeC)
+ *
+ *   using Shape_MNK =    // Logical MxNxK shape of the MMA
+ *
+ *   using ThrID     =    // Logical thread id (tid) -> tidx
+ *
+ *   using ALayout =      // (Logical thread id (tid), Logical value id (vid)) -> Flat MK-coord
+ *   using BLayout =      // (Logical thread id (tid), Logical value id (vid)) -> Flat NK-coord
+ *   using CLayout =      // (Logical thread id (tid), Logical value id (vid)) -> Flat MN-coord
+ * };
+ */
+
+template <class MMAOperation, class... MMAOpArgs>
+struct MMA_Traits
+{
+  static_assert(sizeof(MMAOperation) == 0, "MMA_Traits not implemented for this MMA_Operation.");
+};
+
+template <class D, class A, class B, class C>
+struct MMA_Traits<UniversalFMA<D,A,B,C>>
+{
+  using ValTypeD = D;
+  using ValTypeA = A;
+  using ValTypeB = B;
+  using ValTypeC = C;
+
+  // Logical shape of the MMA
+  using Shape_MNK = Shape<_1,_1,_1>;
+
+  // Logical thread id (tid) -> tidx
+  using ThrID   = Layout<_1>;
+
+  // (Logical thread id (tid), Logical value id (vid)) -> coord
+
+  // (tid,vid) -> (m,k)
+  using ALayout = Layout<Shape<_1,_1>>;
+  // (tid,vid) -> (n,k)
+  using BLayout = Layout<Shape<_1,_1>>;
+  // (tid,vid) -> (m,n)
+  using CLayout = Layout<Shape<_1,_1>>;
+};
+
+//
+// Generic mma_unpack for any MMA_Traits
+//
+template <class MMA_Op, class... MMA_Args,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE constexpr
+void
+mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
+           Tensor<TD, DLayout>      & D,
+           Tensor<TA, ALayout> const& A,
+           Tensor<TB, BLayout> const& B,
+           Tensor<TC, CLayout> const& C)
+{
+  static_assert(is_rmem<TD>::value, "Expected registers in MMA_Atom::call");
+  static_assert(is_rmem<TA>::value, "Expected registers in MMA_Atom::call");
+  static_assert(is_rmem<TB>::value, "Expected registers in MMA_Atom::call");
+  static_assert(is_rmem<TC>::value, "Expected registers in MMA_Atom::call");
+
+  // Register value types from the MMA_Operation register arrays
+  using RegTypeD = typename remove_extent<typename MMA_Op::DRegisters>::type;
+  using RegTypeA = typename remove_extent<typename MMA_Op::ARegisters>::type;
+  using RegTypeB = typename remove_extent<typename MMA_Op::BRegisters>::type;
+  using RegTypeC = typename remove_extent<typename MMA_Op::CRegisters>::type;
+  using MMATraits = MMA_Traits<MMA_Op, MMA_Args...>;
+
+  using UniversalFMAOp = UniversalFMA<
+                            typename TD::value_type,
+                            typename TA::value_type,
+                            typename TB::value_type,
+                            typename TC::value_type>;
+
+  [[maybe_unused]] constexpr int RegNumD = extent<typename MMA_Op::DRegisters>::value;
+  constexpr int RegNumA = extent<typename MMA_Op::ARegisters>::value;
+  constexpr int RegNumB = extent<typename MMA_Op::BRegisters>::value;
+  constexpr int RegNumC = extent<typename MMA_Op::CRegisters>::value;
+
+  Tensor rA = recast<RegTypeA>(A);
+  Tensor rB = recast<RegTypeB>(B);
+
+  MUTE_STATIC_ASSERT_V(size(rA) == Int<RegNumA>{});
+  MUTE_STATIC_ASSERT_V(size(rB) == Int<RegNumB>{});
+
+  if constexpr (is_same<RegTypeD, void>::value)
+  {
+    static_assert(is_same<typename TD::value_type, typename TC::value_type>::value, "SQMMA C and D value_type must match.");
+    static_assert(is_same<DLayout, CLayout>::value, "SQMMA C and D layouts must match.");
+    // assert((void*)&C == (void*)&D);
+
+    Tensor rC = recast<RegTypeC>(D);  // NOTE: D and C are same, so use mutable D
+
+    //MUTE_STATIC_ASSERT_V(size(rC) == Int<RegNumC>{});
+
+    if constexpr (detail::supports_output_scaling<MMATraits>::value) {
+      MMA_Op::fma(rA.data(), rB.data(), rC.data(), traits.accumulate_);
+    }
+    else {
+      detail::explode(MMA_Op::fma,
+                      rA, make_int_sequence<RegNumA>{},
+                      rB, make_int_sequence<RegNumB>{},
+                      rC, make_int_sequence<RegNumC>{});
+    }
+  }
+  else {
+    Tensor rD = recast<RegTypeD>(D);
+    Tensor rC = recast<RegTypeC>(C);
+
+    MUTE_STATIC_ASSERT_V(size(rD) == Int<RegNumD>{});
+    MUTE_STATIC_ASSERT_V(size(rC) == Int<RegNumC>{});
+    if constexpr (detail::supports_output_scaling<MMATraits>::value) {
+      MMA_Op::fma(rD.data(), rA.data(), rB.data(), rC.data(), traits.accumulate_);
+    } else {
+      if constexpr (is_same<MMA_Op, UniversalFMAOp>::value) {
+        detail::explode(MMA_Op::fma,
+                  rD, make_int_sequence<RegNumD>{},
+                  rA, make_int_sequence<RegNumA>{},
+                  rB, make_int_sequence<RegNumB>{},
+                  rC, make_int_sequence<RegNumC>{});
+      } else {
+        MMA_Op::fma(rD.data(), rA.data(), rB.data(), rC.data());
+      }
+    }
+  }
+}
+
+//
+// Accept mutable temporaries
+//
+
+template <class MMA_Op, class... MMA_Args,
+          class TD, class DLayout,
+          class TA, class ALayout,
+          class TB, class BLayout,
+          class TC, class CLayout>
+MUTE_HOST_DEVICE constexpr
+void
+mma_unpack(MMA_Traits<MMA_Op, MMA_Args...> const& traits,
+           Tensor<TD, DLayout>     &&  D,
+           Tensor<TA, ALayout> const&  A,
+           Tensor<TB, BLayout> const&  B,
+           Tensor<TC, CLayout> const&  C)
+{
+  mma_unpack(traits, D, A, B, C);
+}
+
+namespace detail {
+
+template <class X, class = void>
+struct FrgTypeA_or_Default { using type = typename X::ValTypeA; };
+template <class X>
+struct FrgTypeA_or_Default<X,void_t<typename X::FrgTypeA>> { using type = typename X::FrgTypeA; };
+
+template <class X, class = void>
+struct FrgTypeB_or_Default { using type = typename X::ValTypeB; };
+template <class X>
+struct FrgTypeB_or_Default<X,void_t<typename X::FrgTypeB>> { using type = typename X::FrgTypeB; };
+
+template <class X, class = void>
+struct FrgTypeC_or_Default { using type = typename X::ValTypeC; };
+template <class X>
+struct FrgTypeC_or_Default<X,void_t<typename X::FrgTypeC>> { using type = typename X::FrgTypeC; };
+
+} // end namespace detail
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp22.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp22.hpp
new file mode 100644
index 000000000..60dbaf3ca
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp22.hpp
@@ -0,0 +1,301 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/arch/mma_mp22.hpp>
+#include <mute/atom/mma_traits.hpp>
+
+#include <mute/layout.hpp>
+
+namespace mute {
+
+namespace {
+  using MP22_32x32_32b = Layout<Shape <Shape < _8, _4,   _2, _2>, Shape <  _2, _4>>,
+                                Stride<Stride<_32, _1, _256, _4>, Stride<_512, _8>>>;
+
+  using MP22_32x16_Row = Layout<Shape <Shape < _4, _16,   _2>, Shape < _2,  _2>>,
+                                Stride<Stride<_64,  _1, _256>, Stride<_32, _16>>>;
+
+  using MP22_32x16_Col = Layout<Shape <Shape <_8, _16>, Shape <_2,  _2>>,
+                                Stride<Stride<_2, _32>, Stride<_1, _16>>>;
+
+  using MP22_32x8_Row  = Layout<Shape <Shape < _4, _16,   _2>,  _2>,
+                                Stride<Stride<_32,  _1, _128>, _16>>;
+
+  using MP22_32x8_Col  = Layout<Shape <Shape <_16,  _8>,  _2>,
+                                Stride<Stride< _1, _32>, _16>>;
+
+  using MP22_32x32_Row = Layout<Shape <Shape <  _4, _16,   _2>, Shape < _4,  _2>>,
+                                Stride<Stride<_128,  _1, _512>, Stride<_32, _16>>>;
+
+  using MP22_32x32_Col = Layout<Shape <Shape <_4, _32>, Shape <_4,  _2>>,
+                                Stride<Stride<_4, _32>, Stride<_1, _16>>>;
+}
+
+/////////////////////////////////////////////////////////////////////////////////
+////////////////////////// fp32 = fp16 * fp16 + fp32 ////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+template <>
+struct MMA_Traits<MP22_32x32x16_F32F16F16F32_TT>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_32, _32, _16>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x16_Row;
+  using BLayout = MP22_32x16_Col;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x16_F32F16F16F32_TN>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_32, _32, _16>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x16_Row;
+  using BLayout = MP22_32x16_Row;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x16_F32F16F16F32_NT>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_32, _32, _16>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x16_Col;
+  using BLayout = MP22_32x16_Col;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x16_F32F16F16F32_NN>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_32, _32, _16>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x16_Col;
+  using BLayout = MP22_32x16_Row;
+  using CLayout = MP22_32x32_32b;
+};
+
+/////////////////////////////////////////////////////////////////////////////////
+////////////////////////// fp32 = bf16 * bf16 + fp32 ////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+template <>
+struct MMA_Traits<MP22_32x32x16_F32BF16BF16F32_TT>
+     : MMA_Traits<MP22_32x32x16_F32F16F16F32_TT>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x16_F32BF16BF16F32_TN>
+     : MMA_Traits<MP22_32x32x16_F32F16F16F32_TN>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x16_F32BF16BF16F32_NT>
+     : MMA_Traits<MP22_32x32x16_F32F16F16F32_NT>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x16_F32BF16BF16F32_NN>
+     : MMA_Traits<MP22_32x32x16_F32F16F16F32_NN>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+};
+
+/////////////////////////////////////////////////////////////////////////////////
+////////////////////////// fp32 = tf32 * tf32 + fp32 ////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+template <>
+struct MMA_Traits<MP22_32x32x8_F32TF32TF32F32_TT>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_32, _32, _8>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x8_Row;
+  using BLayout = MP22_32x8_Col;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x8_F32TF32TF32F32_TN>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_32, _32, _8>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x8_Row;
+  using BLayout = MP22_32x8_Row;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x8_F32TF32TF32F32_NT>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_32, _32, _8>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x8_Col;
+  using BLayout = MP22_32x8_Col;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x8_F32TF32TF32F32_NN>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_32, _32, _8>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x8_Col;
+  using BLayout = MP22_32x8_Row;
+  using CLayout = MP22_32x32_32b;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+////////////////////////// s32 = s8 * s8 + s32 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+template <>
+struct MMA_Traits<MP22_32x32x32_S32S8S8S32_TT>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x32_Row;
+  using BLayout = MP22_32x32_Col;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x32_S32S8S8S32_TN>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x32_Row;
+  using BLayout = MP22_32x32_Row;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x32_S32S8S8S32_NT>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x32_Col;
+  using BLayout = MP22_32x32_Col;
+  using CLayout = MP22_32x32_32b;
+};
+
+template <>
+struct MMA_Traits<MP22_32x32x32_S32S8S8S32_NN>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID = Layout<_128>;
+  using ALayout = MP22_32x32_Col;
+  using BLayout = MP22_32x32_Row;
+  using CLayout = MP22_32x32_32b;
+};
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp31.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp31.hpp
new file mode 100644
index 000000000..b25c98ed5
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp31.hpp
@@ -0,0 +1,1127 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/arch/mma_mp31.hpp>
+#include <mute/atom/mma_traits.hpp>
+#include <mute/tensor.hpp>
+
+namespace mute {
+
+namespace {
+template <int Bits, int M, int K, TCE::Major tnsp>
+using MP31_MxK = mute::conditional_t<tnsp == TCE::Major::K,
+  // Row layout
+  Layout<Shape <Shape <                _4, _8>, Shape <Int<32 / Bits>, Int<K * Bits / 128>, Int<M / 8>>>,
+         Stride<Stride<Int<M * 32 / Bits>, _1>, Stride<        Int<M>, Int<128 * M / Bits>,        _8>>>,
+  // Column layout
+  Layout<Shape <Shape <Int< Bits / 4>, Int<128 / Bits>>,Shape <Int<32 / Bits>, Int<M / 8>, Int<K * Bits / 128>>>,
+         Stride<Stride<Int<32 / Bits>, Int<         M>>,Stride<            _1,         _8, Int<128 * M / Bits>>>>
+>;
+
+template <int M, int N>
+using MP31_MxN = Layout<Shape <Shape <    _8, _4>, Shape <Int<N/8>, Int<M/4>>>,
+                        Stride<Stride<Int<M>, _1>, Stride<Int<M*8>,      _4>>>;
+}
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x8_F32F16F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _8>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 8, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 8, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x8_F32BF16BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _8>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 8, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 8, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32F16F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32BF16BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32F16F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32BF16BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32F16F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32BF16BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32F16F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32BF16BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_S32S8S8S32<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_U32U8U8U32<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32E5M2E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32E4M3E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32E4M3E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32E5M2E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_S32S8S8S32<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_U32U8U8U32<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32E5M2E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32E4M3E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32E4M3E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32E5M2E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_S32S8S8S32<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_U32U8U8U32<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32E5M2E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32E4M3E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32E4M3E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32E5M2E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_S32S8S8S32<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_U32U8U8U32<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32E5M2E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32E4M3E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32E4M3E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32E5M2E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x64_S32S8S8S32<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using Shape_MNK = Shape<_16, _16, _64>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 64, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 64, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x64_U32U8U8U32<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using Shape_MNK = Shape<_16, _16, _64>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 64, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 64, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x64_F32E5M2E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _64>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 64, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 64, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x64_F32E4M3E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _64>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 64, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 64, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x64_F32E4M3E5M2F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _64>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 64, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 64, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x64_F32E5M2E4M3F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _64>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 64, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 64, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x4_F32TF32TF32F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _4>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 4, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 4, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x8_F32TF32TF32F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _8>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 8, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 8, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32TF32TF32F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32F16S4F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32S4F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32BF16S4F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32S4BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32F16S8F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32S8F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32BF16S8F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x8x16_F32S8BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _8, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>,  8, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 8>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32F16S8F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32S8F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32BF16S8F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_8x16x16_F32S8BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_8, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>,  8, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<8, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32F16S8F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32S8F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32BF16S8F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x16_F32S8BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _16>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 16, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 16, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32F16S8F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32S8F16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32BF16S8F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x16x32_F32S8BF16F32<tnspA, tnspB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using Shape_MNK = Shape<_16, _16, _32>;
+  using ThrID = Layout<_32>;
+  using ALayout = MP31_MxK<sizeof_bits_v<ValTypeA>, 16, 32, tnspA>;
+  using BLayout = MP31_MxK<sizeof_bits_v<ValTypeB>, 16, 32, tnspB>;
+  using CLayout = MP31_MxN<16, 16>;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp31_sqmma.hpp b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp31_sqmma.hpp
new file mode 100644
index 000000000..a26778cb5
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/atom/mma_traits_mp31_sqmma.hpp
@@ -0,0 +1,10316 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/arch/mma_mp31.hpp>
+#include <mute/atom/mma_traits.hpp>
+
+#include <mute/tensor.hpp>
+
+namespace mute {
+
+namespace MP31::SQMMA {
+
+////////////////////////////////////////////
+// Common layouts for SQMMA Shared Memory //
+////////////////////////////////////////////
+template <int MN, int K, class Type, TCE::Major tnsp>
+using Layout_Atom_Bits = mute::conditional_t<tnsp == TCE::Major::MN,
+                            decltype(make_layout(Shape<Int<MN * sizeof_bits_v<Type>>, Int<K>>{}, LayoutLeft{})),
+                            decltype(make_layout(Shape<Int<MN>, Int<K * sizeof_bits_v<Type>>>{}, LayoutRight{}))>;
+
+template <int MN, int K, class Type, TCE::Major tnsp>
+using Layout_SL256_SS256_NONE_Atom = decltype(upcast<sizeof_bits_v<Type>>(ComposedLayout<Swizzle<0,4,4>, smem_ptr_flag, Layout_Atom_Bits<MN,K,Type,tnsp>>{}));
+
+template <int MN, int K, class Type, TCE::Major tnsp>
+using Layout_SL256_SS256_SG16_Atom = decltype(upcast<sizeof_bits_v<Type>>(ComposedLayout<Swizzle<4,4,4>, smem_ptr_flag, Layout_Atom_Bits<MN,K,Type,tnsp>>{}));
+
+template <int MN, int K, class Type, TCE::Major tnsp>
+using Layout_SL256_SS256_SG32_Atom = decltype(upcast<sizeof_bits_v<Type>>(ComposedLayout<Swizzle<3,5,3>, smem_ptr_flag, Layout_Atom_Bits<MN,K,Type,tnsp>>{}));
+
+template <int MN, int K, class Type, TCE::Major tnsp>
+using Layout_SL256_SS256_SG64_Atom = decltype(upcast<sizeof_bits_v<Type>>(ComposedLayout<Swizzle<2,6,2>, smem_ptr_flag, Layout_Atom_Bits<MN,K,Type,tnsp>>{}));
+
+// Only used for canonical gemm
+template <class Type, TCE::Major tnsp, class AtomMN, class AtomK>
+MUTE_HOST_DEVICE constexpr
+auto
+make_canonical_gemm_smem_atom_layout() {
+  constexpr int BITS = sizeof_bits_v<Type>;
+  constexpr int ATOM_MN = size(AtomMN{});
+  constexpr int ATOM_K  = size(AtomK{});
+  if constexpr (tnsp == TCE::Major::MN) {
+    // S4/S8/U8/FP8
+    if constexpr (BITS <= 8) {
+      return Layout_SL256_SS256_SG16_Atom<ATOM_MN, ATOM_K, Type, tnsp>{};
+    }
+    // FP16/BF16
+    else if constexpr (BITS == 16) {
+      return Layout_SL256_SS256_SG32_Atom<ATOM_MN, ATOM_K, Type, tnsp>{};
+    }
+    // TF32
+    else if constexpr (BITS == 32) {
+      return Layout_SL256_SS256_SG64_Atom<ATOM_MN, ATOM_K, Type, tnsp>{};
+    }
+    else {
+      static_assert(BITS <= 32, "Unsupported Sqmma Type");
+    }
+  }
+  else if constexpr (tnsp == TCE::Major::K) {
+    return Layout_SL256_SS256_SG16_Atom<ATOM_MN, ATOM_K, Type, tnsp>{};
+  }
+  else {
+    static_assert(tnsp != TCE::Major::MN && tnsp != TCE::Major::K, "Unrecognized MajorMode!");
+  }
+}
+
+//
+// Tensor (position-dependent swizzle) to SwizzleGranularityType utility
+//
+
+template <class Engine, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+SwizzleGranularityType
+swizzle_granularity_type(Tensor<Engine, Layout<Shape, Stride>> const&)
+{
+  static_assert(is_same<uint8_t, typename Engine::value_type>::value,
+                "Expected uint8_t type in SwizzleGranularityType conversion.");
+
+  using Swizzle = get_swizzle_t<Engine>;
+  constexpr int B = Swizzle::num_bits;
+  constexpr int M = Swizzle::num_base;
+  constexpr int S = Swizzle::num_shft;
+
+  static_assert(4 <= M && M <= 6, "Unsupported swizzle granularity");
+
+  // The Swizzle Line for Sqmma is 256B(2^8)
+  static_assert((M+S) == 8, "Unsupported swizzle granularity");
+
+  if constexpr (B == 0) {
+    return SwizzleGranularityType::SWIZZLE_GRANULARITY_NONE;
+  }
+
+  // The Swizzle Stride for Mp31 Sqmma must be 256B(2^8)
+  static_assert((M+B) == 8, "Unsupported swizzle granularity");
+
+  switch (M) {
+    case 4: return SwizzleGranularityType::SWIZZLE_GRANULARITY_B16;
+    case 5: return SwizzleGranularityType::SWIZZLE_GRANULARITY_B32;
+    case 6: return SwizzleGranularityType::SWIZZLE_GRANULARITY_B64;
+  }
+  return SwizzleGranularityType::SWIZZLE_GRANULARITY_NONE;
+}
+
+
+///////////////////////////////////////////////////////////////////////////////
+// Construction method for SQMMA Descriptors
+///////////////////////////////////////////////////////////////////////////////
+template <TCE::Major MajorMode, class TEngine, class TLayout>
+MUTE_HOST_DEVICE constexpr
+SqmmaDescriptor
+make_sqmma_desc(Tensor<TEngine, TLayout> const& tensor)
+{
+  static_assert(is_smem<TEngine>::value, "SQMMA Descriptors can only be constructed on smem.");
+  static_assert(TLayout::rank == 2, "SQMMA Descriptors can obly be constructed on rank-2 tensors.");
+  using value_type = typename TEngine::value_type;
+
+  Tensor u8_tensor = recast<uint8_t const>(tensor);
+
+  // Result
+  SqmmaDescriptor desc;
+
+  // SwizzleGranularity type
+  constexpr SQMMA::SwizzleGranularityType SG_TYPE = SQMMA::swizzle_granularity_type(u8_tensor);
+  desc.bitfield.swizzle_granularity_type_ = static_cast<uint8_t>(SG_TYPE);
+
+  // Start address
+  desc.bitfield.start_address_ = cast_smem_ptr_to_uint(raw_pointer_cast(u8_tensor.data()));
+
+  if constexpr (MajorMode == TCE::Major::MN)
+  {
+    MUTE_STATIC_ASSERT_V(size<0>(u8_tensor) % Int<16>{} == Int<0>{},    // M|N  size
+                         "Not a canonical SQMMA_MN Layout: Expected MN-size multiple of 16.");
+
+    MUTE_STATIC_ASSERT_V(size<1>(u8_tensor) % Int<8>{} == Int<0>{} &&
+                         size<1>(u8_tensor) <= Int<256>{},                         // K    size
+                         "Not a canonical SQMMA_MN Layout: Expected K-size multiple of 8 (in units of uint8_t).");
+
+    constexpr uint32_t stride_1 = stride<1>(u8_tensor);
+    static_assert(8 <= stride_1 && stride_1 <= 256, "Not a canonical SQMMA_MN Layout: Expected stride failure.");
+
+    switch (stride_1) {
+      case   8: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B8);   break;
+      case  16: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B16);  break;
+      case  32: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B32);  break;
+      case  64: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B64);  break;
+      case 128: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B128); break;
+      case 256: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B256); break;
+    }
+  }
+  else if constexpr (MajorMode == TCE::Major::K)
+  {
+    MUTE_STATIC_ASSERT_V(size<0>(u8_tensor) % Int<16>{} == Int<0>{},    // M|N  size
+                         "Not a canonical SQMMA_K Layout: Expected MN-size multiple of 16.");
+    MUTE_STATIC_ASSERT_V(size<1>(u8_tensor) % Int<8>{} == Int<0>{} &&
+                         size<1>(u8_tensor) <= Int<256>{},                         // K    size
+                         "Not a canonical SQMMA_K Layout: Expected K-size multiple of 8 (in units of uint8_t).");
+
+    constexpr uint32_t stride_0 = stride<0>(u8_tensor);
+    static_assert(8 <= stride_0 && stride_0 <= 256, "Not a canonical SQMMA_K Layout: Expected stride failure.");
+
+    switch (stride_0) {
+      case   8: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B8);   break;
+      case  16: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B16);  break;
+      case  32: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B32);  break;
+      case  64: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B64);  break;
+      case 128: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B128); break;
+      case 256: desc.bitfield.leading_stride_type_ = static_cast<uint8_t>(LeadingStrideType::B256); break;
+    }
+  }
+  else {
+    static_assert(MajorMode != TCE::Major::MN && MajorMode != TCE::Major::K, "Unrecognized MajorMode!");
+  }
+  return desc;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// Higher level SQMMA Descriptor utilities
+///////////////////////////////////////////////////////////////////////////////
+
+struct DescriptorIterator
+{
+  using reference    = SqmmaDescriptor;
+  using element_type = SqmmaDescriptor;
+  using value_type   = SqmmaDescriptor;
+
+  SqmmaDescriptor desc_;
+
+  // Dereference returns the SqmmaDescriptor
+  MUTE_HOST_DEVICE constexpr
+  reference operator*() const { return desc_; }
+
+  // Advance and return a new SqmmaDescriptor
+  template <class Index>
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](Index const& i) const { return *(*this + i); }
+
+  // Return an advanced iterator
+  template <class Index>
+  MUTE_HOST_DEVICE constexpr
+  DescriptorIterator operator+(Index const& offset) const
+  {
+    return { SqmmaDescriptor{desc_ + uint64_t(offset)} };
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  operator int32_t() const { return static_cast<int32_t>(desc_); }
+
+  MUTE_HOST_DEVICE friend void
+  print(DescriptorIterator) { printf("SQMMA::DescriptorIterator"); }
+};
+
+MUTE_HOST_DEVICE constexpr
+SqmmaDescriptor
+raw_pointer_cast(DescriptorIterator const& ptr) {
+  return ptr.desc_;
+}
+
+// Recast a DescriptorIterator Tensor to uint32_t, it's RegType in mma_unpack
+template <class NewT>
+MUTE_HOST_DEVICE constexpr
+DescriptorIterator
+recast_ptr(DescriptorIterator const& iter) {
+  static_assert(is_same<NewT, int32_t>::value, "Can only cast SqmmaDescriptorIterator to int32_t.");
+  return iter;
+}
+
+// The SQMMA Traits below have custom fragment type flags for their smem desc tensors.
+// These flags specialize a MakeTensor customization point to correctly make the fragment that is desired.
+template <TCE::Major>
+struct smem_desc : DescriptorIterator {};
+
+} // namespace MP31::SQMMA
+
+// Customization point for creating a SQMMA::smem_desc Tensor
+template <TCE::Major MajorMode>
+struct MakeTensor<MP31::SQMMA::smem_desc<MajorMode>>
+{
+  template <class TEngine, class TLayout>
+  MUTE_HOST_DEVICE constexpr auto
+  operator()(Tensor<TEngine,TLayout> const& smem_tensor)
+  {
+    static_assert(is_smem<TEngine>::value, "Expected SMEM Tensor to construct a SQMMA Desc Tensor");
+    return make_tensor(MP31::SQMMA::DescriptorIterator{MP31::SQMMA::make_sqmma_desc<MajorMode>(tensor<0>(smem_tensor))},
+                       replace<0>(recast<uint8_t const>(smem_tensor).layout(), Layout<_1,_0>{}));
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+//////////////////////////// MMA_TRAITS ///////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////
+
+namespace MP31::SQMMA {
+
+// Accumulator layouts
+template <int M, int N>
+using CLayout = Layout<Shape <Shape <    _8, _4, _4>, Shape <      _2, Int<N/16>, Int<M/16>>>,
+                       Stride<Stride<Int<M>, _1, _4>, Stride<Int<M*8>, Int<M*16>,      _16>>>;
+
+// Shared memory source layouts for any value type
+template <int M, int K>
+using ABLayout  = Layout<Shape <_128, Shape <Int<M>,Int<K>>>,
+                         Stride<  _0, Stride<    _1,Int<M>>>>;
+
+} // namespace MP31::SQMMA
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  U32 +=   U8 *   U8 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x64x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x64x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x64x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x32x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x32x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x32x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x64x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x64x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x64x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x128x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x128x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x128x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x16x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x16x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x16x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x32x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x32x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x32x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x64x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x64x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x64x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x128x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x128x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x128x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x32x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x32x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x32x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x64x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x64x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x64x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x128x32_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x128x64_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x128x128_U32U8U8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = uint32_t;
+  using ValTypeA = uint8_t;
+  using ValTypeB = uint8_t;
+  using ValTypeC = uint32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  S32 +=   S8 *   S8 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x64x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x64x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_16x64x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x32x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x32x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x32x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x64x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x64x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x64x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x128x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x128x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_32x128x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x16x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x16x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x16x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x32x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x32x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x32x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x64x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x64x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x64x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x128x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x128x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_64x128x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x32x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x32x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x32x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x64x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x64x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x64x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x128x32_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x128x64_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB>
+struct MMA_Traits<MP31_128x128x128_S32S8S8_SS<tnspA, tnspB>>
+{
+  using ValTypeD = int32_t;
+  using ValTypeA = int8_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = int32_t;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 +=  F16 *  F16 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x16_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32F16F16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 += BF16 * BF16 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x16_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32BF16BF16_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 += TF32 * TF32 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x8_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _8>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,   8>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,   8>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x16_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x32_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x8_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _8>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,   8>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,   8>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x16_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x8_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _8>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,   8>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,   8>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x16_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x8_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _8>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,   8>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,   8>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x16_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x8_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _8>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,   8>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,   8>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x16_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x8_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _8>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,   8>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,   8>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x16_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x8_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _8>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,   8>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,   8>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x16_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  16>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x8_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _8>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,   8>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,   8>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x16_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _16>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  16>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32TF32TF32_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = tfloat32_t;
+  using ValTypeB = tfloat32_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 += E4M3 * E4M3 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x128_F32E4M3E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 += E5M2 * E5M2 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x128_F32E5M2E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 +=  F16 *   S4 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_16x64x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_16x64x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x32x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x32x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x64x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x64x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x128x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x128x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x32x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x32x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x64x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x64x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x128x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x128x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x32x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x32x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x64x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x64x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x128x32_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x128x64_F32F16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 +=  F16 *   S8 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_16x64x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_16x64x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x32x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x32x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x64x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x64x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x128x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x128x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x32x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x32x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x64x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x64x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x128x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x128x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x32x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x32x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x64x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x64x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x128x32_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x128x64_F32F16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = half_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 +=   S4 *  F16 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32S4F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 +=   S8 *  F16 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32S8F16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = half_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 += E4M3 * E5M2 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x128_F32E4M3E5M2_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e4m3_t;
+  using ValTypeB = float_e5m2_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 += E5M2 * E4M3 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_16x64x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x128_F32E5M2E4M3_SS<tnspA, tnspB, scaleA, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = float_e5m2_t;
+  using ValTypeB = float_e4m3_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _128>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using BLayout   = MP31::SQMMA::ABLayout<128, 128>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 += BF16 *   S4 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_16x64x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_16x64x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x32x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x32x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x64x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x64x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x128x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x128x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x32x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x32x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x64x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x64x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x128x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x128x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x32x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x32x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x64x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x64x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x128x32_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x128x64_F32BF16S4_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int4_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 += BF16 *   S8 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M16N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_16x64x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_16x64x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_16, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 16,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x32x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x32x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x64x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x64x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x128x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_32x128x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x32x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x32x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x64x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x64x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x128x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_64x128x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x32x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x32x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x64x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x64x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x128x32_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleA>
+struct MMA_Traits<MP31_128x128x64_F32BF16S8_SS<tnspA, tnspB, scaleA>>
+{
+  using ValTypeD = float;
+  using ValTypeA = bfloat16_t;
+  using ValTypeB = int8_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 +=   S4 * BF16 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32S4BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int4_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////
+//////////////////////////  F32 +=   S8 * BF16 //////////////////////////////////
+/////////////////////////////////////////////////////////////////////////////////
+
+//
+// M32N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x32x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x64x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M32N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_32x128x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_32, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 32, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N16
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x16x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _16, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 16,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  16>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x32x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x64x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M64N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_64x128x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_64, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout < 64, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N32
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x32x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _32, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 32,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  32>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N64
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x64x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _64, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout< 64,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128,  64>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+//
+// M128N128
+//
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x32_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _32>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  32>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+template <TCE::Major tnspA, TCE::Major tnspB, MP31::SQMMA::ScaleIn scaleB>
+struct MMA_Traits<MP31_128x128x64_F32S8BF16_SS<tnspA, tnspB, scaleB>>
+{
+  using ValTypeD = float;
+  using ValTypeA = int8_t;
+  using ValTypeB = bfloat16_t;
+  using ValTypeC = float;
+
+  using FrgTypeA = MP31::SQMMA::smem_desc<tnspA>;
+  using FrgTypeB = MP31::SQMMA::smem_desc<tnspB>;
+
+  using Shape_MNK = Shape<_128, _128, _64>;
+  using ThrID     = Layout<_128>;
+  using ALayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using BLayout   = MP31::SQMMA::ABLayout<128,  64>;
+  using CLayout   = MP31::SQMMA::CLayout <128, 128>;
+
+  MP31::SQMMA::ScaleOut accumulate_ = MP31::SQMMA::ScaleOut::One;
+};
+
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/config.hpp b/vllm_musa/csrc/flash_attn/include/mute/config.hpp
new file mode 100644
index 000000000..baaf14e00
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/config.hpp
@@ -0,0 +1,148 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#if defined(__MUSACC__)
+#  define MUTE_HOST_DEVICE __forceinline__ __host__ __device__
+#  define MUTE_DEVICE      __forceinline__          __device__
+#  define MUTE_HOST        __forceinline__ __host__
+#else
+#  define MUTE_HOST_DEVICE inline
+#  define MUTE_DEVICE      inline
+#  define MUTE_HOST        inline
+#endif // MUTE_HOST_DEVICE, MUTE_DEVICE
+
+#if defined(__MUSACC_RTC__)
+#  define MUTE_HOST_RTC MUTE_HOST_DEVICE
+#else
+#  define MUTE_HOST_RTC MUTE_HOST
+#endif
+
+#if defined(__MUSACC_RTC__) || defined(__clang__)
+#  define MUTE_UNROLL    _Pragma("unroll")
+#  define MUTE_NO_UNROLL _Pragma("unroll 1")
+#else
+#  define MUTE_UNROLL
+#  define MUTE_NO_UNROLL
+#endif // MUTE_UNROLL
+
+#if defined(__MUSA_ARCH__)
+#  define MUTE_INLINE_CONSTANT                 static const __device__
+#else
+#  define MUTE_INLINE_CONSTANT                 static constexpr
+#endif
+
+#define MUTE_GRID_CONSTANT
+
+// Some versions of GCC < 11 have trouble deducing that a
+// function with "auto" return type and all of its returns in an "if
+// constexpr ... else" statement must actually return.  Thus, GCC
+// emits spurious "missing return statement" build warnings.
+// Developers can suppress these warnings by using the
+// MUTE_GCC_UNREACHABLE macro, which must be followed by a semicolon.
+// It's harmless to use the macro for other GCC versions or other
+// compilers, but it has no effect.
+#if ! defined(MUTE_GCC_UNREACHABLE)
+#  if defined(__GNUC__)
+#    define MUTE_GCC_UNREACHABLE __builtin_unreachable()
+#  else
+#    define MUTE_GCC_UNREACHABLE
+#  endif
+#endif
+
+#if defined(_MSC_VER)
+// Provides support for alternative operators 'and', 'or', and 'not'
+#  include <iso646.h>
+#endif // _MSC_VER
+
+#if defined(__MUSACC_RTC__)
+#  define MUTE_STL_NAMESPACE musa::std
+#  define MUTE_STL_NAMESPACE_IS_MUSA_STD
+#else
+#  define MUTE_STL_NAMESPACE std
+#endif
+
+//
+// Assertion helpers
+//
+
+#if defined(__MUSACC_RTC__)
+#  include <musa/std/cassert>
+#else
+#  include <cassert>
+#endif
+
+#define MUTE_STATIC_V(x)            decltype(x)::value
+
+#define MUTE_STATIC_ASSERT          static_assert
+#define MUTE_STATIC_ASSERT_V(x,...) static_assert(decltype(x)::value, ##__VA_ARGS__)
+
+// Fail and print a message. Typically used for notification of a compiler misconfiguration.
+#if defined(__MUSA_ARCH__)
+#  define MUTE_INVALID_CONTROL_PATH(x) assert(0 && x); printf(x);
+#else
+#  define MUTE_INVALID_CONTROL_PATH(x) assert(0 && x); printf(x)
+#endif
+
+//
+// IO
+//
+
+#if !defined(__MUSACC_RTC__)
+#  include <cstdio>
+#  include <iostream>
+#  include <iomanip>
+#endif
+
+//
+// AddressSpace
+//
+#include <mute/address_space.hpp>
+
+//
+// Support
+//
+
+#include <mute/util/type_traits.hpp>
+
+//
+// Basic types
+//
+
+#include <mute/numeric/numeric_types.hpp>
+
+//
+// Debugging utilities
+//
+
+#include <mute/util/print.hpp>
+#include <mute/util/debug.hpp>
diff --git a/vllm_musa/csrc/flash_attn/include/mute/container/alignment.hpp b/vllm_musa/csrc/flash_attn/include/mute/container/alignment.hpp
new file mode 100644
index 000000000..b5de2e27d
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/container/alignment.hpp
@@ -0,0 +1,71 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/numeric/numeric_types.hpp>
+#include <mute/numeric/math.hpp>
+
+namespace mute
+{
+
+// Test if a pointer is aligned to N bytes
+template <int N>
+MUTE_HOST_DEVICE constexpr
+bool
+is_byte_aligned(void const* const ptr)
+{
+  static_assert(N > 0 && (N & (N - 1)) == 0, "N must be a power of 2 in alignment check");
+  return (reinterpret_cast<uintptr_t>(ptr) & (N-1)) == 0;
+}
+
+#if defined(__MUSACC__)
+#  define MUTE_ALIGNAS(n) __align__(n)
+#else
+#  define MUTE_ALIGNAS(n) alignas(n)
+#endif
+
+template <size_t Alignment, class Child = void>
+struct aligned_struct {};
+
+template <class Child> struct MUTE_ALIGNAS(  1) aligned_struct<  1, Child> {};
+template <class Child> struct MUTE_ALIGNAS(  2) aligned_struct<  2, Child> {};
+template <class Child> struct MUTE_ALIGNAS(  4) aligned_struct<  4, Child> {};
+template <class Child> struct MUTE_ALIGNAS(  8) aligned_struct<  8, Child> {};
+template <class Child> struct MUTE_ALIGNAS( 16) aligned_struct< 16, Child> {};
+template <class Child> struct MUTE_ALIGNAS( 32) aligned_struct< 32, Child> {};
+template <class Child> struct MUTE_ALIGNAS( 64) aligned_struct< 64, Child> {};
+template <class Child> struct MUTE_ALIGNAS(128) aligned_struct<128, Child> {};
+template <class Child> struct MUTE_ALIGNAS(256) aligned_struct<256, Child> {};
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/container/array.hpp b/vllm_musa/csrc/flash_attn/include/mute/container/array.hpp
new file mode 100644
index 000000000..9a633d244
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/container/array.hpp
@@ -0,0 +1,493 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/numeric/integral_constant.hpp>
+#include <mute/util/type_traits.hpp>
+
+namespace mute
+{
+
+template <class T, size_t N>
+struct array
+{
+  using element_type = T;
+  using value_type = remove_cv_t<T>;
+  using size_type = size_t;
+  using difference_type = ptrdiff_t;
+  using reference = element_type&;
+  using const_reference = const element_type&;
+  using pointer = element_type*;
+  using const_pointer = const element_type*;
+  using iterator = pointer;
+  using const_iterator = const_pointer;
+
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](size_type pos)
+  {
+    return begin()[pos];
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference operator[](size_type pos) const
+  {
+    return begin()[pos];
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference front()
+  {
+    return *begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference front() const
+  {
+    return *begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference back()
+  {
+    // return *rbegin();
+    return operator[](N-1);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference back() const
+  {
+    // return *rbegin();
+    return operator[](N-1);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  T* data()
+  {
+    return __elems_;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  T const* data() const
+  {
+    return __elems_;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  iterator begin()
+  {
+    return data();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator begin() const
+  {
+    return data();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cbegin()
+  {
+    return begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cbegin() const
+  {
+    return begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  iterator end()
+  {
+    return data() + size();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator end() const
+  {
+    return data() + size();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cend()
+  {
+    return end();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cend() const
+  {
+    return end();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  bool empty() const
+  {
+    return size() == 0;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  size_type size() const
+  {
+    return N;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  size_type max_size() const
+  {
+    return size();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  void fill(const T& value)
+  {
+    for (auto& e : *this) {
+      e = value;
+    }
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  void clear()
+  {
+    fill(T(0));
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  void swap(array& other)
+  {
+    using MUTE_STL_NAMESPACE::swap;
+    for (size_type i = 0; i < size(); ++i) {
+      swap((*this)[i], other[i]);
+    }
+  }
+
+  element_type __elems_[N];
+};
+
+
+template <class T>
+struct array<T, 0>
+{
+  using element_type = T;
+  using value_type = remove_cv_t<T>;
+  using size_type = size_t;
+  using difference_type = ptrdiff_t;
+  using reference = element_type&;
+  using const_reference = const element_type&;
+  using pointer = element_type*;
+  using const_pointer = const element_type*;
+  using const_iterator = const_pointer;
+  using iterator = pointer;
+
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](size_type pos)
+  {
+    return begin()[pos];
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference operator[](size_type pos) const
+  {
+    return begin()[pos];
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference front()
+  {
+    return *begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference front() const
+  {
+    return *begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference back()
+  {
+    return *begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference back() const
+  {
+    return *begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  T* data()
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  T const* data() const
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  iterator begin()
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator begin() const
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cbegin()
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cbegin() const
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  iterator end()
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator end() const
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cend()
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cend() const
+  {
+    return nullptr;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  bool empty() const
+  {
+    return true;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  size_type size() const
+  {
+    return 0;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  size_type max_size() const
+  {
+    return 0;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  void fill(const T& value)
+  {}
+
+  MUTE_HOST_DEVICE constexpr
+  void clear()
+  {}
+
+  MUTE_HOST_DEVICE constexpr
+  void swap(array& other)
+  {}
+};
+
+template <class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+bool operator==(array<T,N> const& lhs, array<T,N> const& rhs)
+{
+  for (size_t i = 0; i < N; ++i) {
+    if (lhs[i] != rhs[i]) {
+      return false;
+    }
+  }
+  return true;
+}
+
+template <class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+void clear(array<T,N>& a)
+{
+  a.fill(T(0));
+}
+
+template <class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+void fill(array<T,N>& a, T const& value)
+{
+  a.fill(value);
+}
+
+template <class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+void swap(array<T,N>& a, array<T,N>& b)
+{
+  a.swap(b);
+}
+
+/// @return A mute::array of the elements of @c t in reverse order.
+template <class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+mute::array<T,N> reverse(mute::array<T,N> const& t)
+{
+  if constexpr (N == 0u) {
+    return t;
+  } else {
+    mute::array<T,N> t_r{};
+    for (size_t k = 0; k < N; ++k) {
+      t_r[k] = t[N - k - 1];
+    }
+    return t_r;
+  }
+}
+
+} // end mute
+
+
+//
+// Specialize tuple-related functionality for mute::array
+//
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/tuple>
+#else
+#include <tuple>
+#endif
+
+namespace mute
+{
+
+template <size_t I, class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+T& get(array<T,N>& a)
+{
+  static_assert(I < N, "Index out of range");
+  return a[I];
+}
+
+template <size_t I, class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+T const& get(array<T,N> const& a)
+{
+  static_assert(I < N, "Index out of range");
+  return a[I];
+}
+
+template <size_t I, class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+T&& get(array<T,N>&& a)
+{
+  static_assert(I < N, "Index out of range");
+  return mute::move(a[I]);
+}
+
+} // end namespace mute
+
+namespace MUTE_STL_NAMESPACE
+{
+
+template <class T, size_t N>
+struct tuple_size<mute::array<T,N>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, N>
+{};
+
+template <size_t I, class T, size_t N>
+struct tuple_element<I, mute::array<T,N>>
+{
+  using type = T;
+};
+
+template <class T, size_t N>
+struct tuple_size<mute::array<T,N> const>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, N>
+{};
+
+template <size_t I, class T, size_t N>
+struct tuple_element<I, mute::array<T,N> const>
+{
+  using type = T;
+};
+
+} // end namespace MUTE_STL_NAMESPACE
+
+#ifdef MUTE_STL_NAMESPACE_IS_MUSA_STD
+namespace std
+{
+
+#if defined(__MUSACC_RTC__)
+template <class... _Tp>
+struct tuple_size;
+
+template <size_t _Ip, class... _Tp>
+struct tuple_element;
+#endif
+
+template <class T, size_t N>
+struct tuple_size<mute::array<T,N>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, N>
+{};
+
+template <size_t I, class T, size_t N>
+struct tuple_element<I, mute::array<T,N>>
+{
+  using type = T;
+};
+
+template <class T, size_t N>
+struct tuple_size<mute::array<T,N> const>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, N>
+{};
+
+template <size_t I, class T, size_t N>
+struct tuple_element<I, mute::array<T,N> const>
+{
+  using type = T;
+};
+
+} // end namespace std
+#endif // MUTE_STL_NAMESPACE_IS_MUSA_STD
diff --git a/vllm_musa/csrc/flash_attn/include/mute/container/array_aligned.hpp b/vllm_musa/csrc/flash_attn/include/mute/container/array_aligned.hpp
new file mode 100644
index 000000000..b812f5920
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/container/array_aligned.hpp
@@ -0,0 +1,43 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/container/array.hpp>
+#include <mute/container/alignment.hpp>
+
+namespace mute
+{
+
+template <class T, size_t N, size_t Alignment = 16>
+struct MUTE_ALIGNAS(Alignment) array_aligned : mute::array<T,N> {};
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/container/array_subbyte.hpp b/vllm_musa/csrc/flash_attn/include/mute/container/array_subbyte.hpp
new file mode 100644
index 000000000..3defe7c85
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/container/array_subbyte.hpp
@@ -0,0 +1,663 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Statically sized array of elements that accommodates subbyte trivial types
+           in a packed storage.
+*/
+
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/numeric/numeric_types.hpp>
+#include <mute/numeric/integral_constant.hpp>
+
+namespace mute
+{
+//
+// Underlying subbyte storage type
+//
+template <class T>
+using subbyte_storage_type_t = conditional_t<(mute::sizeof_bits_v<T> <=   8), uint8_t,
+                               conditional_t<(mute::sizeof_bits_v<T> <=  16), uint16_t,
+                               conditional_t<(mute::sizeof_bits_v<T> <=  32), uint32_t,
+                               conditional_t<(mute::sizeof_bits_v<T> <=  64), uint64_t,
+                               conditional_t<(mute::sizeof_bits_v<T> <= 128), uint_bit_t<128>,
+                               T>>>>>;
+
+template <class T> struct subbyte_iterator;
+template <class, class> struct swizzle_ptr;
+
+//
+// subbyte_reference
+//   Proxy object for sub-byte element references
+//
+template <class T>
+struct subbyte_reference
+{
+  // Iterator Element type (const or non-const)
+  using element_type = T;
+  // Iterator Value type without type qualifier.
+  using value_type   = remove_cv_t<T>;
+  // Storage type (const or non-const)
+  using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
+
+  static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
+
+  static_assert(sizeof_bits_v<element_type> <= sizeof_bits_v<storage_type>,
+                "Size of Element must not be greater than Storage.");
+
+private:
+
+  // Bitmask for covering one item
+  static constexpr storage_type BitMask = storage_type(storage_type(-1) >> (sizeof_bits_v<storage_type> - sizeof_bits_v<element_type>));
+  // Flag for fast branching on straddled elements
+  static constexpr bool is_storage_unaligned = ((sizeof_bits_v<storage_type> % sizeof_bits_v<element_type>) != 0);
+
+  friend struct subbyte_iterator<T>;
+
+  // Pointer to storage element
+  storage_type* ptr_ = nullptr;
+
+  // Bit index of value_type starting position within storage_type element.
+  // RI: 0 <= idx_ < sizeof_bit<storage_type>
+  uint8_t idx_ = 0;
+
+  // Ctor
+  template <class PointerType>
+  MUTE_HOST_DEVICE constexpr
+  subbyte_reference(PointerType* ptr, uint8_t idx = 0) : ptr_(reinterpret_cast<storage_type*>(ptr)), idx_(idx) {}
+
+public:
+
+  // Copy Ctor
+  MUTE_HOST_DEVICE constexpr
+  subbyte_reference(subbyte_reference<value_type> const& other) {
+    *this = other.get();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  subbyte_reference(subbyte_reference<value_type const> const& other) {
+    *this = other.get();
+  }
+
+  // Copy Assignment
+  MUTE_HOST_DEVICE constexpr
+  subbyte_reference& operator=(subbyte_reference<value_type> const& other) {
+    return *this = other.get();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  subbyte_reference& operator=(subbyte_reference<value_type const> const& other) {
+    return *this = other.get();
+  }
+
+  // Assignment
+  template <class T_ = element_type>
+  MUTE_HOST_DEVICE constexpr
+  enable_if_t<!is_const_v<T_>, subbyte_reference&> operator=(value_type x)
+  {
+    static_assert(is_same_v<T_, element_type>, "Do not specify template arguments!");
+    storage_type item = (reinterpret_cast<storage_type const&>(x) & BitMask);
+
+    // Update the current storage element
+    storage_type bit_mask_0 = storage_type(BitMask << idx_);
+    ptr_[0] = storage_type((ptr_[0] & ~bit_mask_0) | (item << idx_));
+
+    // If value_type is unaligned with storage_type (static) and this is a straddled value (dynamic)
+    if (is_storage_unaligned && idx_ + sizeof_bits_v<value_type> > sizeof_bits_v<storage_type>) {
+      uint8_t straddle_bits = uint8_t(sizeof_bits_v<storage_type> - idx_);
+      storage_type bit_mask_1 = storage_type(BitMask >> straddle_bits);
+      // Update the next storage element
+      ptr_[1] = storage_type((ptr_[1] & ~bit_mask_1) | (item >> straddle_bits));
+    }
+
+    return *this;
+  }
+
+  // Comparison of referenced values
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator==(subbyte_reference const& x, subbyte_reference const& y) { return x.get() == y.get(); }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator!=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() != y.get(); }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator< (subbyte_reference const& x, subbyte_reference const& y) { return x.get() <  y.get(); }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator> (subbyte_reference const& x, subbyte_reference const& y) { return x.get() >  y.get(); }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator<=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() <= y.get(); }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator>=(subbyte_reference const& x, subbyte_reference const& y) { return x.get() >= y.get(); }
+
+  // Value
+  MUTE_HOST_DEVICE
+  value_type get() const
+  {
+    if constexpr (is_same_v<bool, value_type>) {      // Extract to bool -- potentially faster impl
+      return bool((*ptr_) & (BitMask << idx_));
+    } else {                                          // Extract to value_type
+      // Extract from the current storage element
+      auto item = storage_type((ptr_[0] >> idx_) & BitMask);
+
+      // If value_type is unaligned with storage_type (static) and this is a straddled value (dynamic)
+      if (is_storage_unaligned && idx_ + sizeof_bits_v<value_type> > sizeof_bits_v<storage_type>) {
+        uint8_t straddle_bits = uint8_t(sizeof_bits_v<storage_type> - idx_);
+        storage_type bit_mask_1 = storage_type(BitMask >> straddle_bits);
+        // Extract from the next storage element
+        item |= storage_type((ptr_[1] & bit_mask_1) << straddle_bits);
+      }
+
+      return reinterpret_cast<value_type&>(item);
+    }
+  }
+
+  // Extract to type value_type
+  MUTE_HOST_DEVICE constexpr
+  operator value_type() const {
+    return get();
+  }
+
+  // Address
+  MUTE_HOST_DEVICE
+  subbyte_iterator<T> operator&() const {
+    return {ptr_, idx_};
+  }
+};
+
+template <class T>
+MUTE_HOST_DEVICE
+void
+print(subbyte_reference<T> ref) {
+  mute::print(ref.get());
+}
+
+template <class T>
+MUTE_HOST_DEVICE
+void
+pretty_print(subbyte_reference<T> ref) {
+  mute::pretty_print(ref.get());
+}
+
+//
+// subbyte_iterator
+//   Random-access iterator over subbyte references
+//
+template <class T>
+struct subbyte_iterator
+{
+  // Iterator Element type (const or non-const)
+  using element_type = T;
+  // Iterator Value type without type qualifier.
+  using value_type   = remove_cv_t<T>;
+  // Storage type (const or non-const)
+  using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
+  // Reference proxy type
+  using reference = subbyte_reference<element_type>;
+
+  static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
+
+  static_assert(sizeof_bits_v<element_type> <= sizeof_bits_v<storage_type>,
+                "Size of Element must not be greater than Storage.");
+
+private:
+
+  template <class, class> friend struct swizzle_ptr;
+  template <class U> friend MUTE_HOST_DEVICE constexpr U* raw_pointer_cast(subbyte_iterator<U> const&);
+  template <class N, class U> friend MUTE_HOST_DEVICE constexpr auto recast_ptr(subbyte_iterator<U> const&);
+  template <class U> friend MUTE_HOST_DEVICE void print(subbyte_iterator<U> const&);
+
+  // Pointer to storage element
+  storage_type* ptr_;
+
+  // Bit index of value_type starting position within storage_type element.
+  // RI: 0 <= idx_ < sizeof_bit<storage_type>
+  uint8_t idx_;
+
+public:
+
+  // Default Ctor
+  MUTE_HOST_DEVICE constexpr
+  subbyte_iterator() : ptr_{nullptr}, idx_{0} {};
+
+  // Ctor
+  template <class PointerType>
+  MUTE_HOST_DEVICE constexpr
+  subbyte_iterator(PointerType* ptr, uint8_t idx = 0) : ptr_(reinterpret_cast<storage_type*>(ptr)), idx_(idx) { }
+
+  MUTE_HOST_DEVICE constexpr
+  reference operator*() const {
+    return reference(ptr_, idx_);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  subbyte_iterator& operator+=(uint64_t k) {
+    k = sizeof_bits_v<value_type> * k + idx_;
+    ptr_ += k / sizeof_bits_v<storage_type>;
+    idx_  = k % sizeof_bits_v<storage_type>;
+    return *this;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  subbyte_iterator operator+(uint64_t k) const {
+    return subbyte_iterator(ptr_, idx_) += k;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](uint64_t k) const {
+    return *(*this + k);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  subbyte_iterator& operator++() {
+    idx_ += sizeof_bits_v<value_type>;
+    if (idx_ >= sizeof_bits_v<storage_type>) {
+      ++ptr_;
+      idx_ -= sizeof_bits_v<storage_type>;
+    }
+    return *this;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  subbyte_iterator operator++(int) {
+    subbyte_iterator ret(*this);
+    ++(*this);
+    return ret;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  subbyte_iterator& operator--() {
+    if (idx_ >= sizeof_bits_v<value_type>) {
+      idx_ -= sizeof_bits_v<value_type>;
+    } else {
+      --ptr_;
+      idx_ += sizeof_bits_v<storage_type> - sizeof_bits_v<value_type>;
+    }
+    return *this;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  subbyte_iterator operator--(int) {
+    subbyte_iterator ret(*this);
+    --(*this);
+    return ret;
+  }
+
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator==(subbyte_iterator const& x, subbyte_iterator const& y) {
+    return x.ptr_ == y.ptr_ && x.idx_ == y.idx_;
+  }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator!=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x == y); }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator< (subbyte_iterator const& x, subbyte_iterator const& y) {
+    return x.ptr_ < y.ptr_ || (x.ptr_ == y.ptr_ && x.idx_ < y.idx_);
+  }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator<=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(y <  x); }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator> (subbyte_iterator const& x, subbyte_iterator const& y) { return  (y <  x); }
+  MUTE_HOST_DEVICE constexpr friend
+  bool operator>=(subbyte_iterator const& x, subbyte_iterator const& y) { return !(x <  y); }
+};
+
+// Conversion to raw pointer with loss of subbyte index
+template <class T>
+MUTE_HOST_DEVICE constexpr
+T*
+raw_pointer_cast(subbyte_iterator<T> const& x) {
+  assert(x.idx_ == 0);
+  return reinterpret_cast<T*>(x.ptr_);
+}
+
+// Conversion to NewT_ with possible loss of subbyte index
+template <class NewT_, class T>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(subbyte_iterator<T> const& x) {
+  using NewT = conditional_t<(is_const_v<T>), NewT_ const, NewT_>;
+  if constexpr (mute::is_subbyte_v<NewT>) {       // Making subbyte_iter, preserve the subbyte idx
+    return subbyte_iterator<NewT>(x.ptr_, x.idx_);
+  } else {                                       // Not subbyte, assume/assert subbyte idx 0
+    return reinterpret_cast<NewT*>(raw_pointer_cast(x));
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Dynamic pointers have unknown static alignment
+template <class T>
+MUTE_HOST_DEVICE constexpr
+Int<0>
+max_alignment(subbyte_iterator<T> const& x) {
+  return {};
+}
+
+template <class T>
+MUTE_HOST_DEVICE void
+print(subbyte_iterator<T> const& x) {
+  printf("subptr[%db](%p.%u)", int(sizeof_bits_v<T>), x.ptr_, x.idx_);
+}
+
+template <class T>
+MUTE_HOST_DEVICE void
+print(subbyte_reference<T> const& x) {
+  print(x.get());
+}
+
+//
+// array_subbyte
+//   Statically sized array for non-byte-aligned data types
+//
+template <class T, size_t N>
+struct array_subbyte
+{
+  using element_type    = T;
+  using value_type      = remove_cv_t<T>;
+  using pointer         = element_type*;
+  using const_pointer   = element_type const*;
+
+  using size_type       = size_t;
+  using difference_type = ptrdiff_t;
+
+  //
+  // References
+  //
+  using reference       = subbyte_reference<element_type>;
+  using const_reference = subbyte_reference<element_type const>;
+
+  //
+  // Iterators
+  //
+  using iterator        = subbyte_iterator<element_type>;
+  using const_iterator  = subbyte_iterator<element_type const>;
+
+  // Storage type (const or non-const)
+  using storage_type = conditional_t<(is_const_v<T>), subbyte_storage_type_t<T> const, subbyte_storage_type_t<T>>;
+
+  static_assert(sizeof_bits_v<storage_type> % 8 == 0, "Storage type is not supported");
+
+private:
+
+  // Number of storage elements, ceil_div
+  static constexpr size_type StorageElements = (N * sizeof_bits_v<value_type> + sizeof_bits_v<storage_type> - 1) / sizeof_bits_v<storage_type>;
+
+  // Internal storage
+  storage_type storage[StorageElements];
+
+public:
+
+  MUTE_HOST_DEVICE constexpr
+  size_type size() const {
+    return N;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  size_type max_size() const {
+    return N;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  bool empty() const {
+    return !N;
+  }
+
+  // Efficient clear method
+  MUTE_HOST_DEVICE constexpr
+  void clear() {
+    MUTE_UNROLL
+    for (size_type i = 0; i < StorageElements; ++i) {
+      storage[i] = storage_type(0);
+    }
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  void fill(T const& value) {
+    MUTE_UNROLL
+    for (size_type i = 0; i < N; ++i) {
+      at(i) = value;
+    }
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference at(size_type pos) {
+    return iterator(storage)[pos];
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference at(size_type pos) const {
+    return const_iterator(storage)[pos];
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](size_type pos) {
+    return at(pos);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference operator[](size_type pos) const {
+    return at(pos);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference front() {
+    return at(0);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference front() const {
+    return at(0);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference back() {
+    return at(N-1);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_reference back() const {
+    return at(N-1);
+  }
+
+  // In analogy to std::vector<bool>::data(), these functions are deleted to prevent bugs.
+  // Instead, prefer
+  //   auto* data = raw_pointer_cast(my_subbyte_array.begin());
+  // where the type of auto* is implementation-defined and
+  // with the knowledge that [data, data + my_subbyte_array.size()) may not be a valid range.
+  MUTE_HOST_DEVICE constexpr
+  pointer data() = delete;
+
+  MUTE_HOST_DEVICE constexpr
+  const_pointer data() const = delete;
+
+  MUTE_HOST_DEVICE constexpr
+  iterator begin() {
+    return iterator(storage);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator begin() const {
+    return const_iterator(storage);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cbegin() const {
+    return begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  iterator end() {
+    return iterator(storage) + N;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator end() const {
+    return const_iterator(storage) + N;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  const_iterator cend() const {
+    return end();
+  }
+
+  //
+  // Comparison operators
+  //
+
+};
+
+//
+// Operators
+//
+
+template <class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+void clear(array_subbyte<T,N>& a)
+{
+  a.clear();
+}
+
+template <class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+void fill(array_subbyte<T,N>& a, T const& value)
+{
+  a.fill(value);
+}
+
+} // namespace mute
+
+//
+// Specialize tuple-related functionality for mute::array_subbyte
+//
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/tuple>
+#else
+#include <tuple>
+#endif
+
+namespace mute
+{
+
+template <size_t I, class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+T& get(array_subbyte<T,N>& a)
+{
+  static_assert(I < N, "Index out of range");
+  return a[I];
+}
+
+template <size_t I, class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+T const& get(array_subbyte<T,N> const& a)
+{
+  static_assert(I < N, "Index out of range");
+  return a[I];
+}
+
+template <size_t I, class T, size_t N>
+MUTE_HOST_DEVICE constexpr
+T&& get(array_subbyte<T,N>&& a)
+{
+  static_assert(I < N, "Index out of range");
+  return mute::move(a[I]);
+}
+
+} // end namespace mute
+
+namespace MUTE_STL_NAMESPACE
+{
+
+template <class T>
+struct is_reference<mute::subbyte_reference<T>>
+    : MUTE_STL_NAMESPACE::true_type
+{};
+
+
+template <class T, size_t N>
+struct tuple_size<mute::array_subbyte<T,N>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, N>
+{};
+
+template <size_t I, class T, size_t N>
+struct tuple_element<I, mute::array_subbyte<T,N>>
+{
+  using type = T;
+};
+
+template <class T, size_t N>
+struct tuple_size<const mute::array_subbyte<T,N>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, N>
+{};
+
+template <size_t I, class T, size_t N>
+struct tuple_element<I, const mute::array_subbyte<T,N>>
+{
+  using type = T;
+};
+
+} // end namespace MUTE_STL_NAMESPACE
+
+#ifdef MUTE_STL_NAMESPACE_IS_MUSA_STD
+namespace std
+{
+
+#if defined(__MUSACC_RTC__)
+template <class... _Tp>
+struct tuple_size;
+
+template <size_t _Ip, class... _Tp>
+struct tuple_element;
+#endif
+
+template <class T, size_t N>
+struct tuple_size<mute::array_subbyte<T,N>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, N>
+{};
+
+template <size_t I, class T, size_t N>
+struct tuple_element<I, mute::array_subbyte<T,N>>
+{
+  using type = T;
+};
+
+template <class T, size_t N>
+struct tuple_size<const mute::array_subbyte<T,N>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, N>
+{};
+
+template <size_t I, class T, size_t N>
+struct tuple_element<I, const mute::array_subbyte<T,N>>
+{
+  using type = T;
+};
+
+} // end namespace std
+#endif // MUTE_STL_NAMESPACE_IS_MUSA_STD
diff --git a/vllm_musa/csrc/flash_attn/include/mute/container/bit_field.hpp b/vllm_musa/csrc/flash_attn/include/mute/container/bit_field.hpp
new file mode 100644
index 000000000..d9a3069c8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/container/bit_field.hpp
@@ -0,0 +1,133 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Portable bit field that supports byte and word straddling that can
+           be used in unions to bit-wise define parameters.
+*/
+
+#pragma once
+
+#include <mute/config.hpp>                  // MUTE_HOST_DEVICE
+#include <mute/numeric/numeric_types.hpp>   // uint_bit_t
+#include <mute/util/type_traits.hpp>        // mute::is_same
+namespace mute
+{
+
+class dummy_type {};
+
+template <uint32_t BitStart, uint32_t NumBits, class OtherValueType = dummy_type>
+struct bit_field
+{
+  static_assert(0 < NumBits && NumBits <= 64, "bit_fields with more than 64 bits are not supported.");
+
+  // value_type: Use the smallest value type that fits NumBits
+  static constexpr uint32_t value_type_bits = (NumBits <=  8) ?  8 :
+                                              (NumBits <= 16) ? 16 :
+                                              (NumBits <= 32) ? 32 : 64;
+  using value_type   = mute::uint_bit_t<value_type_bits>;
+  // storage_type: Use the smallest storage_type that avoids boundary crossing
+  static constexpr uint32_t storage_type_bits = (BitStart /  8 == (BitStart + NumBits - 1) /  8) ?  8 :
+                                                (BitStart / 16 == (BitStart + NumBits - 1) / 16) ? 16 :
+                                                (BitStart / 32 == (BitStart + NumBits - 1) / 32) ? 32 : 64;
+  using storage_type = mute::uint_bit_t<storage_type_bits>;
+
+  static_assert(sizeof(OtherValueType) == sizeof(value_type) || is_same<OtherValueType,dummy_type>::value,
+                "sizeof(OtherValueType) must be same as sizeof(value_type).");
+
+  // Number of storage values needed: ceil_div(BitStart + NumBits, storage_type_bits)
+  static constexpr uint32_t N      = (BitStart + NumBits + storage_type_bits - 1) / storage_type_bits;
+  // Index of storage value for BitStart
+  static constexpr uint32_t idx    = BitStart / storage_type_bits;
+  // Bit of data_[idx] for BitStart
+  static constexpr uint32_t bit_lo = BitStart % storage_type_bits;
+  // Number of bits in data_[idx] used for NumBits if straddling, else 0
+  static constexpr uint32_t bit_hi = (idx + 1 < N) ? (storage_type_bits - bit_lo) : 0;
+
+public:
+
+  // NumBits mask
+  static constexpr value_type   mask    = value_type(uint64_t(-1) >> (64u - NumBits));
+  // NumBits mask for BitStart
+  static constexpr storage_type mask_lo = storage_type(mask) << bit_lo;
+  // NumBits mask for leftover bits in data_[idx+1] if straddling, else 0
+  static constexpr storage_type mask_hi = (idx + 1 < N) ? (storage_type(mask) >> bit_hi) : 0;
+
+  storage_type data_[N];
+
+  // Get value
+  MUTE_HOST_DEVICE constexpr
+  value_type get() const {
+    storage_type result = (data_[idx] & mask_lo) >> bit_lo;
+    if constexpr (bit_hi != 0) {
+      result |= (data_[idx+1] & mask_hi) << bit_hi;
+    }
+    return static_cast<value_type>(result);
+  }
+
+  // Set value
+  MUTE_HOST_DEVICE constexpr
+  void set(value_type x) {
+    storage_type item = static_cast<storage_type>(x & mask);
+    data_[idx] = static_cast<storage_type>((data_[idx] & ~mask_lo) | (item << bit_lo));
+    if constexpr (bit_hi != 0) {
+      data_[idx+1] = static_cast<storage_type>((data_[idx+1] & ~mask_hi) | (item >> bit_hi));
+    }
+  }
+
+  // Assign value
+  MUTE_HOST_DEVICE constexpr
+  bit_field& operator=(value_type x) {
+    set(x);
+    return *this;
+  }
+
+  // Cast to value
+  MUTE_HOST_DEVICE constexpr
+  operator value_type () const {
+    return get();
+  }
+
+  // Assign OtherValueType
+  MUTE_HOST_DEVICE constexpr
+  bit_field& operator=(OtherValueType x) {
+    return *this = *reinterpret_cast<value_type*>(&x);
+  }
+
+  // Cast to OtherValueType
+  MUTE_HOST_DEVICE constexpr
+  operator OtherValueType () const {
+    value_type x = get();
+    return *reinterpret_cast<OtherValueType*>(&x);
+  }
+};
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/container/musa_types.hpp b/vllm_musa/csrc/flash_attn/include/mute/container/musa_types.hpp
new file mode 100644
index 000000000..7a1ed1777
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/container/musa_types.hpp
@@ -0,0 +1,183 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE, MUTE_GCC_UNREACHABLE
+#include <mute/numeric/integral_constant.hpp>  // mute::integral_constant
+namespace mute
+{
+
+//
+// dim3
+//
+
+using dim3 = ::dim3;
+
+// MSVC doesn't define its C++ version macro to match
+// its C++ language version.  This means that when
+// building with MSVC, dim3 isn't constexpr-friendly.
+template <size_t I>
+MUTE_HOST_DEVICE
+#if ! defined(_MSC_VER)
+constexpr
+#endif
+uint32_t& get(dim3& a)
+{
+  static_assert(I < 3, "Index out of range");
+  if constexpr (I == 0) {
+    return a.x;
+  } else if constexpr (I == 1) {
+    return a.y;
+  } else if constexpr (I == 2) {
+    return a.z;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <size_t I>
+MUTE_HOST_DEVICE
+#if ! defined(_MSC_VER)
+constexpr
+#endif
+uint32_t const& get(dim3 const& a)
+{
+  static_assert(I < 3, "Index out of range");
+  if constexpr (I == 0) {
+    return a.x;
+  } else if constexpr (I == 1) {
+    return a.y;
+  } else if constexpr (I == 2) {
+    return a.z;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <size_t I>
+MUTE_HOST_DEVICE
+#if ! defined(_MSC_VER)
+constexpr
+#endif
+uint32_t&& get(dim3&& a)
+{
+  static_assert(I < 3, "Index out of range");
+  if constexpr (I == 0) {
+    return mute::move(a.x);
+  } else if constexpr (I == 1) {
+    return mute::move(a.y);
+  } else if constexpr (I == 2) {
+    return mute::move(a.z);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Specialize mute::tuple-traits for external types
+template <>
+struct tuple_size<dim3>
+    : integral_constant<size_t, 3>
+{};
+
+template <size_t I>
+struct tuple_element<I, dim3>
+{
+  using type = uint32_t;
+};
+
+//
+// uint3
+//
+
+using uint3 = ::uint3;
+
+template <size_t I>
+MUTE_HOST_DEVICE constexpr
+uint32_t& get(uint3& a)
+{
+  static_assert(I < 3, "Index out of range");
+  if constexpr (I == 0) {
+    return a.x;
+  } else if constexpr (I == 1) {
+    return a.y;
+  } else if constexpr (I == 2) {
+    return a.z;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <size_t I>
+MUTE_HOST_DEVICE constexpr
+uint32_t const& get(uint3 const& a)
+{
+  static_assert(I < 3, "Index out of range");
+  if constexpr (I == 0) {
+    return a.x;
+  } else if constexpr (I == 1) {
+    return a.y;
+  } else if constexpr (I == 2) {
+    return a.z;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <size_t I>
+MUTE_HOST_DEVICE constexpr
+uint32_t&& get(uint3&& a)
+{
+  static_assert(I < 3, "Index out of range");
+  if constexpr (I == 0) {
+    return mute::move(a.x);
+  } else if constexpr (I == 1) {
+    return mute::move(a.y);
+  } else if constexpr (I == 2) {
+    return mute::move(a.z);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Specialize mute::tuple-traits for external types
+template <>
+struct tuple_size<uint3>
+    : integral_constant<size_t, 3>
+{};
+
+template <size_t I>
+struct tuple_element<I, uint3>
+{
+  using type = uint32_t;
+};
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/container/tuple.hpp b/vllm_musa/csrc/flash_attn/include/mute/container/tuple.hpp
new file mode 100644
index 000000000..50cfc66f2
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/container/tuple.hpp
@@ -0,0 +1,730 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+#include <mute/util/type_traits.hpp>
+#include <mute/numeric/integral_constant.hpp>  // mute::true_type, mute::false_type
+#include <mute/numeric/integer_sequence.hpp>
+
+#include <mute/container/musa_types.hpp>
+
+//#include <mute/container/array.hpp>            // Advanced optimizations
+
+//
+// mute::tuple is like std::tuple, with two differences.
+//
+// 1. It works on both host and device.
+// 2. Its template arguments must be semiregular types.
+//
+// Semiregular types are default constructible and copyable.
+// They include "value types" like int or float,
+// but do _not_ include references like int& or float&.
+// (See std::tie for an example of a tuple of references.)
+//
+// This is simplified over the implementations in std::, musa::std::, and thrust:: by ignoring much of
+//    the conversion SFINAE, special overloading, and avoiding cvref template types.
+//    Furthermore, the empty base optimization (EBO) is MORE aggressive by avoiding
+//    construction calls, and ignoring any need for unique element addresses.
+//
+// Over standard-conforming tuple implementations, this appears to accelerate compilation times by over 3x.
+
+namespace mute
+{
+
+namespace detail
+{
+
+// EBO stands for "empty base optimization."
+// We use this technique to ensure that mute::tuple
+// doesn't need to waste space storing any template arguments
+// of mute::tuple that have no data (like integral_constant).
+// Otherwise, mute::tuple would need to spend at least 1 byte
+// for each of its template arguments.
+//
+// EBO always "holds" a single value of type T.
+// N is like an array index that TupleBase uses
+// to access the desired tuple element.
+template <size_t N, class T, bool IsEmpty = is_empty<T>::value>
+struct EBO;
+
+template <class T, size_t N, bool B>
+MUTE_HOST_DEVICE constexpr C<N> findt(EBO<N, T, B> const&)
+{ return {}; }
+
+// Specialization for types T that have no data;
+// the "static tuple leaf."  Valid T here include
+// integral_constant<U, Value>, Int<Value>,
+// and any other semiregular type
+// for which std::is_empty_v<T> is true.
+template <size_t N, class T>
+struct EBO<N, T, true>
+{
+  MUTE_HOST_DEVICE constexpr
+  EBO() {}
+
+  MUTE_HOST_DEVICE constexpr
+  EBO(T const&) {}
+};
+
+template <size_t N, class T>
+MUTE_HOST_DEVICE constexpr T getv(EBO<N, T, true> const&)
+{ return {}; }
+
+// Specialization for types T that are not empty;
+// the "dynamic tuple leaf."  Valid T here include int,
+// any other integral or floating-point type,
+// or any semiregular type for which std::is_empty_v<T> is false.
+template <size_t N, class T>
+struct EBO<N, T, false>
+{
+  MUTE_HOST_DEVICE constexpr
+  EBO() : t_{} {}
+
+  template <class U>
+  MUTE_HOST_DEVICE constexpr
+  EBO(U const& u) : t_{u} {}
+
+  T t_;
+};
+
+template <size_t N, class T>
+MUTE_HOST_DEVICE constexpr T const& getv(EBO<N, T, false> const& x)
+{ return x.t_; }
+
+template <size_t N, class T>
+MUTE_HOST_DEVICE constexpr T& getv(EBO<N, T, false>& x)
+{ return x.t_; }
+
+template <size_t N, class T>
+MUTE_HOST_DEVICE constexpr T&& getv(EBO<N, T, false>&& x)
+{ return mute::move(x.t_); }
+
+template <class IdxSeq, class... T>
+struct TupleBase;
+
+// Base class of mute::tuple binds each element to an index
+// by inheriting from EBO<i, t> for each (i, t) in (I..., T...).
+// The storage (for nonempty t) lives in the base classes.
+template <size_t... I, class... T>
+struct TupleBase<index_sequence<I...>, T...>
+    : EBO<I,T>...
+{
+  MUTE_HOST_DEVICE constexpr
+  TupleBase() {}
+
+  template <class... U>
+  MUTE_HOST_DEVICE constexpr explicit
+  TupleBase(U const&... u)
+      : EBO<I,T>(u)... {}
+
+  template <class... U>
+  MUTE_HOST_DEVICE constexpr
+  TupleBase(TupleBase<index_sequence<I...>, U...> const& u)
+      : EBO<I,T>(getv(static_cast<EBO<I,U> const&>(u)))... {}
+};
+
+} // end namespace detail
+
+// Attempting to use the following commented-out alias
+// in the declaration of `struct tuple` causes MSVC 2022 build errors.
+//
+//template <class... T>
+//using TupleBase = detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>;
+
+// This is the actual mute::tuple class.
+// The storage (if any) lives in TupleBase's EBO base classes.
+//
+// Inheriting from the above alias TupleBase
+// causes MSVC 2022 build errors when assigning one tuple to another:
+// In summary: this is verbose as a work-around for MSVC build errors.
+template <class... T>
+struct tuple : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>
+{
+  MUTE_HOST_DEVICE constexpr
+  tuple() {}
+
+  template <class... U>
+  MUTE_HOST_DEVICE constexpr
+  tuple(U const&... u) : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(u...) {}
+
+  template <class... U>
+  MUTE_HOST_DEVICE constexpr
+  tuple(tuple<U...> const& u)
+      : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(static_cast<detail::TupleBase<make_index_sequence<sizeof...(U)>, U...> const&>(u)) {}
+};
+
+//
+// get for mute::tuple (just like std::get for std::tuple)
+//
+
+template <size_t I, class... T>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+get(tuple<T...> const& t) noexcept
+{
+  static_assert(I < sizeof...(T), "Index out of range");
+  return detail::getv<I>(t);
+}
+
+template <size_t I, class... T>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+get(tuple<T...>& t) noexcept
+{
+  static_assert(I < sizeof...(T), "Index out of range");
+  return detail::getv<I>(t);
+}
+
+template <size_t I, class... T>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+get(tuple<T...>&& t) noexcept
+{
+  static_assert(I < sizeof...(T), "Index out of range");
+  return detail::getv<I>(static_cast<tuple<T...>&&>(t));
+}
+
+//
+// find a type X within a mute::tuple
+//   Requires X to be unique in tuple
+//   Returns a static integer
+//
+
+template <class X, class... T>
+MUTE_HOST_DEVICE constexpr
+auto
+find(tuple<T...> const& t) noexcept
+{
+  return detail::findt<X>(t);
+}
+
+//
+// Custom is_tuple trait simply checks the existence of tuple_size
+//      and assumes std::get<I>(.), std::tuple_element<I,.>
+//
+namespace detail {
+
+template <class T>
+auto has_tuple_size( T*) -> bool_constant<(0 <= tuple_size<T>::value)>;
+auto has_tuple_size(...) -> false_type;
+
+} // end namespace detail
+
+template <class T>
+struct is_tuple : decltype(detail::has_tuple_size((T*)0)) {};
+
+//
+// make_tuple (value-based implementation)
+//
+
+template <class... T>
+MUTE_HOST_DEVICE constexpr
+tuple<T...>
+make_tuple(T const&... t)
+{
+  return {t...};
+}
+
+//
+// tuple_cat concatenates multiple mute::tuple into a single mute::tuple,
+// just like std::tuple_cat for std::tuple.
+//
+
+#if 0
+// Original implementation
+
+namespace detail {
+
+template <class T0, class T1,
+          size_t... I0, size_t... I1>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1,
+          index_sequence<I0...>, index_sequence<I1...>)
+{
+  return mute::make_tuple(get<I0>(t0)..., get<I1>(t1)...);
+}
+
+} // end namespace detail
+
+MUTE_HOST_DEVICE constexpr
+tuple<>
+tuple_cat()
+{
+  return {};
+}
+
+template <class Tuple,
+          __MUTE_REQUIRES(is_tuple<Tuple>::value)>
+MUTE_HOST_DEVICE constexpr
+Tuple const&
+tuple_cat(Tuple const& t)
+{
+  return t;
+}
+
+template <class T0, class T1>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1)
+{
+  return detail::tuple_cat(t0, t1,
+                           make_index_sequence<tuple_size<T0>::value>{},
+                           make_index_sequence<tuple_size<T1>::value>{});
+}
+
+template <class T0, class T1, class T2, class... Ts>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, Ts const&... ts)
+{
+  return mute::tuple_cat(mute::tuple_cat(t0,t1),t2,ts...);
+}
+#endif
+
+#if 1
+// Extended implementation
+
+namespace detail {
+
+template <class T0, class T1,
+          size_t... I0, size_t... I1>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1,
+          index_sequence<I0...>, index_sequence<I1...>)
+{
+  return mute::make_tuple(get<I0>(t0)..., get<I1>(t1)...);
+}
+
+template <class T0, class T1, class T2,
+          size_t... I0, size_t... I1, size_t... I2>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2,
+          index_sequence<I0...>, index_sequence<I1...>, index_sequence<I2...>)
+{
+  return mute::make_tuple(get<I0>(t0)..., get<I1>(t1)..., get<I2>(t2)...);
+}
+
+template <class T0, class T1, class T2, class T3,
+          size_t... I0, size_t... I1, size_t... I2, size_t... I3>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3,
+          index_sequence<I0...>, index_sequence<I1...>, index_sequence<I2...>, index_sequence<I3...>)
+{
+  return mute::make_tuple(get<I0>(t0)..., get<I1>(t1)..., get<I2>(t2)..., get<I3>(t3)...);
+}
+
+template <class T0, class T1, class T2, class T3, class T4,
+          size_t... I0, size_t... I1, size_t... I2, size_t... I3, size_t... I4>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4,
+          index_sequence<I0...>, index_sequence<I1...>, index_sequence<I2...>, index_sequence<I3...>, index_sequence<I4...>)
+{
+  return mute::make_tuple(get<I0>(t0)..., get<I1>(t1)..., get<I2>(t2)..., get<I3>(t3)..., get<I4>(t4)...);
+}
+
+template <class T0, class T1>
+struct tuple_cat_static;
+
+template <class... T0s, class... T1s>
+struct tuple_cat_static<tuple<T0s...>, tuple<T1s...>> {
+  using type = tuple<T0s..., T1s...>;
+};
+
+} // end namespace detail
+
+MUTE_HOST_DEVICE constexpr
+tuple<>
+tuple_cat()
+{
+  return {};
+}
+
+template <class Tuple,
+          __MUTE_REQUIRES(is_tuple<Tuple>::value)>
+MUTE_HOST_DEVICE constexpr
+Tuple const&
+tuple_cat(Tuple const& t)
+{
+  return t;
+}
+
+template <class T0, class T1>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1)
+{
+  if constexpr (is_static<T0>::value && is_static<T1>::value &&
+		is_tuple<T0>::value && is_tuple<T1>::value) {
+    return typename detail::tuple_cat_static<T0, T1>::type{};
+  } else {
+    return detail::tuple_cat(t0, t1,
+                           make_index_sequence<tuple_size<T0>::value>{},
+                           make_index_sequence<tuple_size<T1>::value>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T0, class T1, class T2>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2)
+{
+  return detail::tuple_cat(t0, t1, t2,
+                           make_index_sequence<tuple_size<T0>::value>{},
+                           make_index_sequence<tuple_size<T1>::value>{},
+                           make_index_sequence<tuple_size<T2>::value>{});
+}
+
+template <class T0, class T1, class T2, class T3>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3)
+{
+  return detail::tuple_cat(t0, t1, t2, t3,
+                           make_index_sequence<tuple_size<T0>::value>{},
+                           make_index_sequence<tuple_size<T1>::value>{},
+                           make_index_sequence<tuple_size<T2>::value>{},
+                           make_index_sequence<tuple_size<T3>::value>{});
+}
+
+template <class T0, class T1, class T2, class T3, class T4>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4)
+{
+  return detail::tuple_cat(t0, t1, t2, t3, t4,
+                           make_index_sequence<tuple_size<T0>::value>{},
+                           make_index_sequence<tuple_size<T1>::value>{},
+                           make_index_sequence<tuple_size<T2>::value>{},
+                           make_index_sequence<tuple_size<T3>::value>{},
+                           make_index_sequence<tuple_size<T4>::value>{});
+}
+
+template <class T0, class T1, class T2, class T3, class T4, class T5, class... Ts>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1, T2 const& t2, T3 const& t3, T4 const& t4, T5 const& t5, Ts const&... ts)
+{
+  return mute::tuple_cat(mute::tuple_cat(t0,t1,t2,t3,t4), mute::tuple_cat(t5, ts...));
+}
+#endif
+
+#if 0
+// Outer-Inner indexing trick to concat all tuples at once
+
+namespace detail {
+
+template <size_t... Ns>
+struct tuple_cat_helper
+{
+  static constexpr mute::array<size_t,sizeof...(Ns)> ns = {Ns...};
+
+  static constexpr size_t total_size() {
+    size_t sum = 0;
+    for (size_t n : ns) sum += n;
+    return sum;
+  }
+  static constexpr size_t total_size_ = total_size();
+
+  static constexpr auto values() {
+    mute::array<size_t[2],total_size_> outer_inner = {};
+
+    size_t idx = 0;
+    for (size_t i = 0; i < ns.size(); ++i) {
+      for (size_t j = 0; j < ns[i]; ++j, ++idx) {
+        outer_inner[idx][0] = i;
+        outer_inner[idx][1] = j;
+      }
+    }
+    return outer_inner;
+  }
+  static constexpr auto outer_inner_ = values();
+
+  using total_sequence = make_index_sequence<total_size_>;
+};
+
+template <class Helper, class Tuple, size_t... I>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(Tuple const& t, index_sequence<I...>)
+{
+  return mute::make_tuple(get<Helper::outer_inner_[I][1]>(get<Helper::outer_inner_[I][0]>(t))...);
+}
+
+template <class T0, class T1,
+          size_t... I0, size_t... I1>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1,
+          index_sequence<I0...>, index_sequence<I1...>)
+{
+  return mute::make_tuple(get<I0>(t0)..., get<I1>(t1)...);
+}
+
+} // end namespace detail
+
+MUTE_HOST_DEVICE constexpr
+tuple<>
+tuple_cat()
+{
+  return {};
+}
+
+template <class Tuple,
+          __MUTE_REQUIRES(is_tuple<Tuple>::value)>
+MUTE_HOST_DEVICE constexpr
+Tuple const&
+tuple_cat(Tuple const& t)
+{
+  return t;
+}
+
+template <class T0, class T1>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(T0 const& t0, T1 const& t1)
+{
+  return detail::tuple_cat(t0, t1,
+                           make_index_sequence<tuple_size<T0>::value>{},
+                           make_index_sequence<tuple_size<T1>::value>{});
+}
+
+template <class... Tuples>
+MUTE_HOST_DEVICE constexpr
+auto
+tuple_cat(Tuples const&... ts)
+{
+  using Helper = detail::tuple_cat_helper<tuple_size<Tuples>::value...>;
+  return detail::tuple_cat<Helper>(mute::make_tuple(ts...), typename Helper::total_sequence{});
+}
+#endif
+
+//
+// Equality operators
+//
+
+namespace detail {
+
+template <size_t I, class TupleA, class TupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+equal_impl(TupleA const& a, TupleB const& b)
+{
+  if constexpr (I == tuple_size<TupleA>::value) {
+    return mute::true_type{};   // Terminal: TupleA is exhausted
+  } else if constexpr (I == tuple_size<TupleB>::value) {
+    return mute::false_type{};  // Terminal: TupleA is not exhausted, TupleB is exhausted
+  } else {
+    return (get<I>(a) == get<I>(b)) && equal_impl<I+1>(a,b);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+template <class TupleT, class TupleU,
+          __MUTE_REQUIRES(is_tuple<TupleT>::value && is_tuple<TupleU>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+operator==(TupleT const& t, TupleU const& u)
+{
+  return detail::equal_impl<0>(t, u);
+}
+
+template <class TupleT, class TupleU,
+          __MUTE_REQUIRES(is_tuple<TupleT>::value ^ is_tuple<TupleU>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+operator==(TupleT const& t, TupleU const& u)
+{
+  return mute::false_type{};
+}
+
+template <class TupleT, class TupleU,
+          __MUTE_REQUIRES(is_tuple<TupleT>::value && is_tuple<TupleU>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+operator!=(TupleT const& t, TupleU const& u)
+{
+  return !(t == u);
+}
+
+template <class TupleT, class TupleU,
+          __MUTE_REQUIRES(is_tuple<TupleT>::value ^ is_tuple<TupleU>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+operator!=(TupleT const& t, TupleU const& u)
+{
+  return mute::true_type{};
+}
+
+//
+// Comparison operators
+//
+
+//
+// There are many ways to compare tuple of elements and because MuTe is built
+//   on parameterizing layouts of coordinates, some comparisons are appropriate
+//   only in certain cases.
+//  -- lexicographical comparison [reverse, reflected, revref]
+//  -- colexicographical comparison [reverse, reflected, revref]
+//  -- element-wise comparison [any,all]
+// This can be very confusing. To avoid errors in selecting the appropriate
+//   comparison, op<|op<=|op>|op>= are *not* implemented for mute::tuple.
+//
+// That said, see int_tuple for more explicitly named common comparison ops.
+//
+
+//
+// Display utilities
+//
+
+namespace detail {
+
+template <class Tuple, size_t... Is>
+MUTE_HOST_DEVICE void print_tuple(Tuple const& t,
+                                  index_sequence<Is...>, char s = '(', char e = ')')
+{
+  using mute::print;
+  if (sizeof...(Is) == 0) {
+    print(s);
+  } else {
+    ((void(print(Is == 0 ? s : ',')), void(print(get<Is>(t)))), ...);
+  }
+  print(e);
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class Tuple, std::size_t... Is>
+MUTE_HOST std::ostream& print_tuple_os(std::ostream& os, Tuple const& t,
+                                       index_sequence<Is...>, char s = '(', char e = ')')
+{
+  if (sizeof...(Is) == 0) {
+    os << s;
+  } else {
+    (void(os << (Is == 0 ? s : ',') << get<Is>(t)), ...);
+  }
+  return os << e;
+}
+#endif // !defined(__MUSACC_RTC__)
+
+} // end namespace detail
+
+template <class Tuple,
+          __MUTE_REQUIRES(is_tuple<Tuple>::value)>
+MUTE_HOST_DEVICE void print(Tuple const& t)
+{
+  return detail::print_tuple(t, make_index_sequence<tuple_size<Tuple>::value>{});
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class Tuple,
+          __MUTE_REQUIRES(is_tuple<Tuple>::value)>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, Tuple const& t)
+{
+  return detail::print_tuple_os(os, t, make_index_sequence<tuple_size<Tuple>::value>{});
+}
+#endif // !defined(__MUSACC_RTC__)
+
+} // end namespace mute
+
+namespace MUTE_STL_NAMESPACE
+{
+
+template <class... T>
+struct tuple_size<mute::tuple<T...>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, mute::tuple<T...>>
+    : MUTE_STL_NAMESPACE::tuple_element<I, MUTE_STL_NAMESPACE::tuple<T...>>
+{};
+
+template <class... T>
+struct tuple_size<const mute::tuple<T...>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, const mute::tuple<T...>>
+    : MUTE_STL_NAMESPACE::tuple_element<I, const MUTE_STL_NAMESPACE::tuple<T...>>
+{};
+
+} // end namespace MUTE_STL_NAMESPACE
+
+//
+// std compatibility
+//
+
+#ifdef MUTE_STL_NAMESPACE_IS_MUSA_STD
+namespace std
+{
+
+#if defined(__MUSACC_RTC__)
+template <class... _Tp>
+struct tuple_size;
+
+template <size_t _Ip, class... _Tp>
+struct tuple_element;
+#endif
+
+template <class... T>
+struct tuple_size<mute::tuple<T...>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, mute::tuple<T...>>
+    : MUTE_STL_NAMESPACE::tuple_element<I, MUTE_STL_NAMESPACE::tuple<T...>>
+{};
+
+template <class... T>
+struct tuple_size<const mute::tuple<T...>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, const mute::tuple<T...>>
+    : MUTE_STL_NAMESPACE::tuple_element<I, const MUTE_STL_NAMESPACE::tuple<T...>>
+{};
+
+} // end namepsace std
+#endif // MUTE_STL_NAMESPACE_IS_MUSA_STD
diff --git a/vllm_musa/csrc/flash_attn/include/mute/container/type_list.hpp b/vllm_musa/csrc/flash_attn/include/mute/container/type_list.hpp
new file mode 100644
index 000000000..6841fd251
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/container/type_list.hpp
@@ -0,0 +1,137 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/numeric/integral_constant.hpp>
+
+namespace mute
+{
+
+template <class T>
+struct type_c {
+  using type = T;
+};
+
+template <class... T>
+struct type_list {};
+
+} // end namespace mute
+
+//
+// Specialize tuple-related functionality for mute::type_list
+//
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/tuple>
+#else
+#include <tuple>
+#endif
+
+#include <mute/container/tuple.hpp>
+
+namespace mute
+{
+
+template <int I, class... T>
+MUTE_HOST_DEVICE constexpr
+MUTE_STL_NAMESPACE::tuple_element_t<I, type_list<T...>>
+get(type_list<T...>&) noexcept {
+  return {};
+}
+template <int I, class... T>
+MUTE_HOST_DEVICE constexpr
+MUTE_STL_NAMESPACE::tuple_element_t<I, type_list<T...>>
+get(type_list<T...> const& t) noexcept {
+  return {};
+}
+
+} // end namespace mute
+
+namespace MUTE_STL_NAMESPACE
+{
+
+template <class... T>
+struct tuple_size<mute::type_list<T...>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, mute::type_list<T...>>
+    : mute::type_c<typename MUTE_STL_NAMESPACE::tuple_element<I, MUTE_STL_NAMESPACE::tuple<T...>>::type>
+{};
+
+template <class... T>
+struct tuple_size<const mute::type_list<T...>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, const mute::type_list<T...>>
+    : mute::type_c<typename MUTE_STL_NAMESPACE::tuple_element<I, MUTE_STL_NAMESPACE::tuple<T...>>::type>
+{};
+
+} // end namespace std
+
+#ifdef MUTE_STL_NAMESPACE_IS_MUSA_STD
+namespace std
+{
+
+#if defined(__MUSACC_RTC__)
+template <class... _Tp>
+struct tuple_size;
+
+template <size_t _Ip, class... _Tp>
+struct tuple_element;
+#endif
+
+template <class... T>
+struct tuple_size<mute::type_list<T...>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, mute::type_list<T...>>
+    : mute::type_c<typename MUTE_STL_NAMESPACE::tuple_element<I, MUTE_STL_NAMESPACE::tuple<T...>>::type>
+{};
+
+template <class... T>
+struct tuple_size<const mute::type_list<T...>>
+    : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, const mute::type_list<T...>>
+    : mute::type_c<typename MUTE_STL_NAMESPACE::tuple_element<I, MUTE_STL_NAMESPACE::tuple<T...>>::type>
+{};
+
+} // end namespace std
+#endif // MUTE_STL_NAMESPACE_IS_MUSA_STD
diff --git a/vllm_musa/csrc/flash_attn/include/mute/int_tuple.hpp b/vllm_musa/csrc/flash_attn/include/mute/int_tuple.hpp
new file mode 100644
index 000000000..91eb7083e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/int_tuple.hpp
@@ -0,0 +1,865 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                      // MUTE_HOST_DEVICE
+#include <mute/container/array.hpp>             // mute::array
+#include <mute/container/tuple.hpp>             // mute::is_tuple
+#include <mute/numeric/integral_constant.hpp>   // mute::Int
+#include <mute/algorithm/tuple_algorithms.hpp>  // mute::transform
+
+/** IntTuple is an integer or a tuple of IntTuples.
+ * This file holds utilities for working with IntTuples,
+ * but does not hold a concrete concept or class of IntTuple.
+ */
+
+namespace mute
+{
+
+// Implementation of get<0>(Integral).
+//   Even though is_tuple<Integral> is false and tuple_size<Integral> doesn't compile,
+//   MuTe defines rank(Integral) as 1, so it's useful for get<0>(Integral) to return its input
+template <size_t I, class T, __MUTE_REQUIRES(mute::is_integral<mute::remove_cvref_t<T>>::value)>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+get(T&& t) noexcept
+{
+  static_assert(I == 0, "Index out of range");
+  return static_cast<T&&>(t);
+}
+
+// Custom recursive get for anything that implements get<I>(.) (for a single integer I).
+template <size_t I0, size_t I1, size_t... Is, class T>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+get(T&& t) noexcept
+{
+  return get<I1, Is...>(get<I0>(static_cast<T&&>(t)));
+}
+
+//
+// rank
+//
+
+template <int... Is, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+rank(IntTuple const& t)
+{
+  if constexpr (sizeof...(Is) == 0) {
+    if constexpr (is_tuple<IntTuple>::value) {
+      return Int<tuple_size<IntTuple>::value>{};
+    } else {
+      return Int<1>{};
+    }
+  } else {
+    return rank(get<Is...>(t));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class IntTuple>
+using rank_t = decltype(rank(declval<IntTuple>()));
+
+template <class IntTuple>
+static constexpr auto rank_v = rank_t<IntTuple>::value;
+
+//
+// shape
+//
+
+template <class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+shape(IntTuple const& s)
+{
+  if constexpr (is_tuple<IntTuple>::value) {
+    return transform(s, [](auto const& a) { return shape(a); });
+  } else {
+    return s;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <int I, int... Is, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+shape(IntTuple const& s)
+{
+  if constexpr (is_tuple<IntTuple>::value) {
+    return shape<Is...>(get<I>(s));
+  } else {
+    return get<I,Is...>(shape(s));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// max
+//
+
+template <class T0, class... Ts>
+MUTE_HOST_DEVICE constexpr
+auto
+max(T0 const& t0, Ts const&... ts)
+{
+  if constexpr (is_tuple<T0>::value) {
+    return mute::max(mute::apply(t0, [](auto const&... a){ return mute::max(a...); }), ts...);
+  } else if constexpr (sizeof...(Ts) == 0) {
+    return t0;
+  } else {
+    return mute::max(t0, mute::max(ts...));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// min
+//
+
+template <class T0, class... Ts>
+MUTE_HOST_DEVICE constexpr
+auto
+min(T0 const& t0, Ts const&... ts)
+{
+  if constexpr (is_tuple<T0>::value) {
+    return mute::min(mute::apply(t0, [](auto const&... a){ return mute::min(a...); }), ts...);
+  } else if constexpr (sizeof...(Ts) == 0) {
+    return t0;
+  } else {
+    return mute::min(t0, mute::min(ts...));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// gcd
+//
+
+template <class T0, class... Ts>
+MUTE_HOST_DEVICE constexpr
+auto
+gcd(T0 const& t0, Ts const&... ts)
+{
+  if constexpr (is_tuple<T0>::value) {
+    return mute::gcd(mute::apply(t0, [](auto const&... a){ return mute::gcd(a...); }), ts...);
+  } else if constexpr (sizeof...(Ts) == 0) {
+    return t0;
+  } else {
+    return mute::gcd(t0, mute::gcd(ts...));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// depth
+//
+
+template <int... Is, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+depth(IntTuple const& t)
+{
+  if constexpr (sizeof...(Is) == 0) {
+    if constexpr (is_tuple<IntTuple>::value) {
+      return Int<1>{} + mute::apply(t, [](auto const&... v){ return mute::max(depth(v)...); });
+    } else {
+      return Int<0>{};
+    }
+  } else {
+    return depth(get<Is...>(t));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class Tuple>
+using depth_t = decltype(depth(declval<Tuple>()));
+
+template <class Tuple>
+static constexpr auto depth_v = depth_t<Tuple>::value;
+
+//
+// product
+//
+
+// Implementation of product as a function object
+struct Product
+{
+  template <class IntTuple>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator()(IntTuple const& a) const
+  {
+    if constexpr (is_tuple<IntTuple>::value) {
+      if constexpr (tuple_size<IntTuple>::value == 0) {
+        return Int<1>{};
+      } else {
+        return mute::transform_apply(a, Product{}, multiplies_unary_lfold{});
+      }
+    } else if constexpr (mute::is_integral<IntTuple>::value) {
+      return a;
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+};
+// Callable product function object
+MUTE_INLINE_CONSTANT Product product;
+
+// Return a rank(t) tuple @a result such that get<i>(@a result) = product(get<i>(@a t))
+template <class Tuple>
+MUTE_HOST_DEVICE constexpr
+auto
+product_each(Tuple const& t)
+{
+  return transform(wrap(t), product);
+}
+
+// Take the product of Tuple at the leaves of TupleG
+template <class Tuple, class TupleG>
+MUTE_HOST_DEVICE constexpr
+auto
+product_like(Tuple const& tuple, TupleG const& guide)
+{
+  return transform_leaf(guide, tuple, [](auto const& g, auto const& t) { return product(t); });
+}
+
+// Return the product of elements in a mode
+template <int... Is, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+size(IntTuple const& a)
+{
+  if constexpr (sizeof...(Is) == 0) {
+    return product(a);
+  } else {
+    return size(get<Is...>(a));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class IntTuple>
+static constexpr auto size_v = decltype(size(declval<IntTuple>()))::value;
+
+//
+// sum
+//
+
+template <class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+sum(IntTuple const& a)
+{
+  if constexpr (is_tuple<IntTuple>::value) {
+    return mute::apply(a, [](auto const&... v){ return (Int<0>{} + ... + sum(v)); });
+  } else {
+    return a;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// inner_product
+//
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+inner_product(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
+    static_assert(tuple_size<IntTupleA>::value == tuple_size<IntTupleB>::value, "Mismatched ranks");
+    return transform_apply(a, b, [](auto const& x, auto const& y) { return inner_product(x,y); },
+                                 [](auto const&... v) { return (Int<0>{} + ... + v); });
+  } else {
+    return a * b;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// ceil_div
+//
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+ceil_div(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value) {
+    if constexpr (is_tuple<IntTupleB>::value) {  // tuple tuple
+      static_assert(tuple_size<IntTupleA>::value >= tuple_size<IntTupleB>::value, "Mismatched ranks");
+      constexpr int R = tuple_size<IntTupleA>::value;        // Missing ranks in TupleB are implicitly 1
+      return transform(a, append<R>(b,Int<1>{}), [](auto const& x, auto const& y) { return ceil_div(x,y); });
+    } else {                                     // tuple int
+      auto const [result, rest] = fold(a, mute::make_tuple(mute::make_tuple(), b),
+        [] (auto const& init, auto const& ai) {
+          return mute::make_tuple(append(get<0>(init), ceil_div(ai, get<1>(init))), ceil_div(get<1>(init), ai));
+        });
+      return result;
+    }
+  } else
+  if constexpr (is_tuple<IntTupleB>::value) {    // int tuple
+    return ceil_div(a, product(b));
+  } else {
+    return (a + b - Int<1>{}) / b;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// round_up
+//   Round @a a up to the nearest multiple of @a b.
+//   For negative numbers, rounds away from zero.
+//
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+round_up(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
+    static_assert(tuple_size<IntTupleA>::value >= tuple_size<IntTupleB>::value, "Mismatched ranks");
+    constexpr int R = tuple_size<IntTupleA>::value;        // Missing ranks in TupleB are implicitly 1
+    return transform(a, append<R>(b,Int<1>{}), [](auto const& x, auto const& y) { return round_up(x,y); });
+  } else {
+    return ((a + b - Int<1>{}) / b) * b;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/** Division for Shapes
+ * Case Tuple Tuple:
+ *   Perform shape_div element-wise
+ * Case Tuple Int:
+ *   Fold the division of b across each element of a
+ *   Example: shape_div((4,5,6),40) -> shape_div((1,5,6),10) -> shape_div((1,1,6),2) -> (1,1,3)
+ * Case Int Tuple:
+ *   Return shape_div(a, product(b))
+ * Case Int Int:
+ *   Enforce the divisibility condition a % b == 0 || b % a == 0 when possible
+ *   Return a / b with rounding away from 0 (that is, 1 or -1 when a < b)
+ */
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+shape_div(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value) {
+    if constexpr (is_tuple<IntTupleB>::value) {  // tuple tuple
+      static_assert(tuple_size<IntTupleA>::value == tuple_size<IntTupleB>::value, "Mismatched ranks");
+      return transform(a, b, [](auto const& x, auto const& y) { return shape_div(x,y); });
+    } else {                                     // tuple int
+      auto const [result, rest] = fold(a, mute::make_tuple(mute::make_tuple(), b),
+        [] (auto const& init, auto const& ai) {
+          return mute::make_tuple(append(get<0>(init), shape_div(ai, get<1>(init))), shape_div(get<1>(init), ai));
+        });
+      return result;
+    }
+  } else
+  if constexpr (is_tuple<IntTupleB>::value) {    // int tuple
+    return shape_div(a, product(b));
+  } else
+  if constexpr (is_static<IntTupleA>::value && is_static<IntTupleB>::value) {
+    static_assert(IntTupleA::value % IntTupleB::value == 0 || IntTupleB::value % IntTupleA::value == 0, "Static shape_div failure");
+    return C<shape_div(IntTupleA::value, IntTupleB::value)>{};
+  } else {                                       // int int
+    //assert(a % b == 0 || b % a == 0);          // Waive dynamic assertion
+    return a / b != 0 ? a / b : signum(a) * signum(b);  // Division with rounding away from zero
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/** Minimum for Shapes
+ */
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+shape_min(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value || is_tuple<IntTupleB>::value) {
+    static_assert(dependent_false<IntTupleA>, "Not implemented.");
+  } else
+  if constexpr (is_constant<1, IntTupleA>::value || is_constant<1, IntTupleB>::value) {
+    return Int<1>{};            // _1 is less than all other shapes, preserve static
+  } else {
+    return mute::min(a, b);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/** Return a tuple the same profile as A scaled by corresponding elements in B
+ */
+template <class A, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+elem_scale(A const& a, B const& b)
+{
+  if constexpr (is_tuple<A>::value) {
+    return transform(a, b, [](auto const& x, auto const& y) { return elem_scale(x,y); });
+  } else {
+    return a * product(b);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/** Test if two IntTuple have the same profile (hierarchical rank division)
+ */
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+congruent(IntTupleA const& a, IntTupleB const& b)
+{
+  return bool_constant<is_same<decltype(repeat_like(shape(a),_0{})),
+                               decltype(repeat_like(shape(b),_0{}))>::value>{};
+}
+
+template <class A, class B>
+using is_congruent = decltype(congruent(declval<A>(), declval<B>()));
+
+/** Test if two IntTuple have the similar profiles up to Shape A (hierarchical rank division)
+ * weakly_congruent is a partial order on A and B: A <= B
+ */
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+weakly_congruent(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
+    if constexpr (tuple_size<IntTupleA>::value != tuple_size<IntTupleB>::value) {
+      return false_type{};
+    } else {
+      return transform_apply(a, b, [](auto const& x, auto const& y) { return weakly_congruent(x,y); },
+                                   [](auto const&... z) { return (true_type{} && ... && z); });
+    }
+  } else if constexpr (is_integral<IntTupleA>::value) {
+    return true_type{};
+  } else if constexpr (is_integral<IntTupleB>::value) {
+    return false_type{};
+  } else {
+    return weakly_congruent(shape(a), shape(b));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class A, class B>
+using is_weakly_congruent = decltype(weakly_congruent(declval<A>(), declval<B>()));
+
+/** Test if Shape A is compatible with Shape B:
+ *    the size of A and B are the same, and
+ *    any coordinate into A can also be used as a coordinate into B
+ * Equivalently, the size of Shape B is the same as Shape A at each terminal of Shape A.
+ * compatible is a partial order on A and B: A <= B
+ */
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+compatible(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
+    if constexpr (tuple_size<IntTupleA>::value != tuple_size<IntTupleB>::value) {
+      return false_type{};
+    } else {
+      return transform_apply(a, b, [](auto const& x, auto const& y) { return compatible(x,y); },
+                                   [](auto const&... z) { return (true_type{} && ... && z); });
+    }
+  } else if constexpr (is_integral<IntTupleA>::value) {
+    return a == size(b);
+  } else if constexpr (is_integral<IntTupleB>::value) {
+    return false_type{};
+  } else {
+    return compatible(shape(a), shape(b));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class A, class B>
+using is_compatible = decltype(compatible(declval<A>(), declval<B>()));
+
+/** Test if Shape A is evenly divided by Tiler B
+ * @returns Static or dynamic boolean
+ * @post if result is true_type, then
+ *       size(a) == logical_divide(make_layout(shape(a)),b) will always compile
+ *       and result in true_type.
+ */
+template <class Shape, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+evenly_divides(Shape const& a, Tiler const& b)
+{
+  if constexpr (is_tuple<Tiler>::value) {
+    if constexpr (rank_v<Tiler> > rank_v<Shape>) {
+      return false_type{};
+    } else {
+      return transform_apply(b, a, [](auto const& x, auto const& y) { return evenly_divides(y,x); },
+                                   [](auto const&... z) { return (true_type{} && ... && z); });
+    }
+  } else {
+    return size(a) == size(b) * size(ceil_div(shape(a), b));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/** Replace the elements of Tuple B that are paired with an Int<0> with an Int<1>
+ */
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+filter_zeros(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value) {
+    return transform(a, b, [](auto const& x, auto const& y) { return filter_zeros(x,y); });
+  } else if constexpr (is_constant<0, IntTupleA>::value) {
+    return repeat_like(b, Int<1>{});
+  } else {
+    return b;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class Tuple>
+MUTE_HOST_DEVICE constexpr
+auto
+filter_zeros(Tuple const& t)
+{
+  return filter_zeros(t, t);
+}
+
+//
+// Converters and constructors with arrays and params
+//
+
+/** Make an IntTuple of rank N from an Indexable array.
+ * Access elements up to a dynamic index n, then use init (requires compatible types)
+ * Consider mute::take<B,E> if all indexing is known to be valid
+ * \code
+ *   std::vector<int> a = {6,3,4};
+ *   auto tup = make_int_tuple<5>(a, a.size(), 0)            // (6,3,4,0,0)
+ * \endcode
+ */
+template <int N, class Indexable, class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_int_tuple(Indexable const& t, int n, T const& init)
+{
+  static_assert(N > 0);
+  if constexpr (N == 1) {
+    return 0 < n ? t[0] : init;
+  } else {
+    return transform(make_seq<N>{}, [&](auto i) { return i < n ? t[i] : init; });
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/** Fill the dynamic values of a Tuple with values from another Tuple
+ * \code
+ *   auto params = make_tuple(6,3,4);
+ *   mute::tuple<Int<1>, mute::tuple<int, int, Int<3>>, int, Int<2>> result;
+ *   fill_int_tuple_from(result, params);                    // (_1,(6,3,_3),4,_2)
+ * \endcode
+ */
+template <class Tuple, class TupleV>
+MUTE_HOST_DEVICE constexpr
+auto
+fill_int_tuple_from(Tuple& result, TupleV const& vals)
+{
+  return fold(result, vals, [](auto const& init, auto&& r) {
+    if constexpr (is_static<remove_cvref_t<decltype(r)>>::value) {       // Skip static elements of result
+      return init;
+    } else if constexpr (is_tuple<remove_cvref_t<decltype(r)>>::value) { // Recurse into tuples
+      return fill_int_tuple_from(r, init);
+    } else {                                                             // Assign and consume arg
+      static_assert(tuple_size<remove_cvref_t<decltype(init)>>::value > 0, "Not enough values to fill with!");
+      r = get<0>(init);
+      return remove<0>(init);
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  });
+}
+
+/** Make a "Tuple" by filling in the dynamic values in order from the arguments
+ * \code
+ *   using result_t = mute::tuple<Int<1>, mute::tuple<int, int, Int<3>>, int, Int<2>>;
+ *   auto result = make_int_tuple_from<result_t>(6,3,4);     // (_1,(6,3,_3),4,_2)
+ * \endcode
+ */
+template <class Tuple, class... Ts>
+MUTE_HOST_DEVICE constexpr
+Tuple
+make_int_tuple_from(Ts const&... ts)
+{
+  Tuple result = Tuple{};
+  fill_int_tuple_from(result, mute::make_tuple(ts...));
+  return result;
+}
+
+/** Convert a tuple to a flat homogeneous array of type T
+ * \code
+ *   auto tup = mute::make_tuple(Int<1>{}, mute::make_tuple(6,3,Int<3>{}),4,Int<2>{});
+ *   mute::array<uint64_t,6> result = to_array<uint64_t>(tup);   // [1,6,3,3,4,2]
+ * \endcode
+ */
+template <class T = int64_t, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+to_array(IntTuple const& t)
+{
+  auto flat_t = flatten_to_tuple(t);
+  constexpr int N = tuple_size<decltype(flat_t)>::value;
+  mute::array<T,N> result;
+  for_each(make_seq<N>{}, [&] (auto i) { result[i] = get<i>(flat_t); });
+  return result;
+}
+
+//
+// Comparison operators
+//
+
+//
+// There are many ways to compare tuple of elements and because MuTe is built
+//   on parameterizing layouts of coordinates, some comparisons are appropriate
+//   only in certain cases.
+//  -- lexicographical comparison [reverse, reflected, revref]   : Correct for coords in RowMajor Layout
+//  -- colexicographical comparison [reverse, reflected, revref] : Correct for coords in ColMajor Layout
+//  -- element-wise comparison [any,all]                         :
+// This can be very confusing. To avoid errors in selecting the appropriate
+//   comparison, op<|op<=|op>|op>= are *not* implemented for mute::tuple.
+//
+// When actually desiring to order coordinates, the user should map them to
+//   their indices within the Layout they came from:
+//      e.g.  layoutX(coordA) < layoutX(coordB)
+// That said, we implement the three most common ways to compare tuples below.
+//   These are implemented with slighly more explicit names than op<.
+//
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+lex_less(IntTupleA const& a, IntTupleB const& b);
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+colex_less(IntTupleA const& a, IntTupleB const& b);
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+elem_less(IntTupleA const& a, IntTupleB const& b);
+
+namespace detail {
+
+template <size_t I, class TupleA, class TupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+lex_less_impl(TupleA const& a, TupleB const& b)
+{
+  if constexpr (I == tuple_size<TupleB>::value) {
+    return mute::false_type{};    // Terminal: TupleB is exhausted
+  } else if constexpr (I == tuple_size<TupleA>::value) {
+    return mute::true_type{};     // Terminal: TupleA is exhausted, TupleB is not exhausted
+  } else {
+    return lex_less(get<I>(a), get<I>(b)) || (get<I>(a) == get<I>(b) && lex_less_impl<I+1>(a,b));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <size_t I, class TupleA, class TupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+colex_less_impl(TupleA const& a, TupleB const& b)
+{
+  if constexpr (I == tuple_size<TupleB>::value) {
+    return mute::false_type{};    // Terminal: TupleB is exhausted
+  } else if constexpr (I == tuple_size<TupleA>::value) {
+    return mute::true_type{};     // Terminal: TupleA is exhausted, TupleB is not exhausted
+  } else {
+    constexpr size_t A = tuple_size<TupleA>::value - 1 - I;
+    constexpr size_t B = tuple_size<TupleB>::value - 1 - I;
+    return colex_less(get<A>(a), get<B>(b)) || (get<A>(a) == get<B>(b) && colex_less_impl<I+1>(a,b));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <size_t I, class TupleA, class TupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+elem_less_impl(TupleA const& a, TupleB const& b)
+{
+  if constexpr (I == tuple_size<TupleA>::value) {
+    return mute::true_type{};     // Terminal: TupleA is exhausted
+  } else if constexpr (I == tuple_size<TupleB>::value) {
+    return mute::false_type{};    // Terminal: TupleA is not exhausted, TupleB is exhausted
+  } else {
+    return elem_less(get<I>(a), get<I>(b)) && elem_less_impl<I+1>(a,b);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+// Lexicographical comparison
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+lex_less(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
+    return detail::lex_less_impl<0>(a, b);
+  } else {
+    return a < b;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+lex_leq(T const& t, U const& u) {
+  return !lex_less(u, t);
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+lex_gtr(T const& t, U const& u) {
+  return lex_less(u, t);
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+lex_geq(T const& t, U const& u) {
+  return !lex_less(t, u);
+}
+
+// Colexicographical comparison
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+colex_less(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
+    return detail::colex_less_impl<0>(a, b);
+  } else {
+    return a < b;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+colex_leq(T const& t, U const& u) {
+  return !colex_less(u, t);
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+colex_gtr(T const& t, U const& u) {
+  return colex_less(u, t);
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+colex_geq(T const& t, U const& u) {
+  return !colex_less(t, u);
+}
+
+// Elementwise [all] comparison
+
+template <class IntTupleA, class IntTupleB>
+MUTE_HOST_DEVICE constexpr
+auto
+elem_less(IntTupleA const& a, IntTupleB const& b)
+{
+  if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
+    return detail::elem_less_impl<0>(a, b);
+  } else {
+    return a < b;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+elem_leq(T const& t, U const& u) {
+  return !elem_less(u, t);
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+elem_gtr(T const& t, U const& u) {
+  return elem_less(u, t);
+}
+
+template <class T, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+elem_geq(T const& t, U const& u) {
+  return !elem_less(t, u);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/layout.hpp b/vllm_musa/csrc/flash_attn/include/mute/layout.hpp
new file mode 100644
index 000000000..feec71fec
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/layout.hpp
@@ -0,0 +1,2065 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+#include <mute/int_tuple.hpp>
+#include <mute/stride.hpp>
+#include <mute/underscore.hpp>
+#include <mute/numeric/arithmetic_tuple.hpp>
+#include <mute/numeric/integral_constant.hpp>
+#include <mute/numeric/integral_ratio.hpp>
+#include <mute/numeric/numeric_types.hpp>  // mute::sizeof_bits
+
+namespace mute
+{
+
+// Aliases
+
+template <class... Shapes>
+using Shape = mute::tuple<Shapes...>;
+
+template <class... Strides>
+using Stride = mute::tuple<Strides...>;
+
+template <class... Strides>
+using Step = mute::tuple<Strides...>;
+
+template <class... Coords>
+using Coord = mute::tuple<Coords...>;
+
+template <class... Layouts>
+using Tile = mute::tuple<Layouts...>;
+
+template <class... Ts>
+MUTE_HOST_DEVICE constexpr
+Shape<Ts...>
+make_shape(Ts const&... t) {
+  return {t...};
+}
+template <class... Ts>
+MUTE_HOST_DEVICE constexpr
+Stride<Ts...>
+make_stride(Ts const&... t) {
+  return {t...};
+}
+template <class... Ts>
+MUTE_HOST_DEVICE constexpr
+Step<Ts...>
+make_step(Ts const&... t) {
+  return {t...};
+}
+template <class... Ts>
+MUTE_HOST_DEVICE constexpr
+Coord<Ts...>
+make_coord(Ts const&... t) {
+  return {t...};
+}
+template <class... Ts>
+MUTE_HOST_DEVICE constexpr
+Tile<Ts...>
+make_tile(Ts const&... t)
+{
+  return {t...};
+}
+
+//
+// Layout
+//
+
+template <class Shape, class Stride = LayoutLeft::Apply<Shape> >
+struct Layout
+    : private mute::tuple<Shape, Stride>   // EBO for static layouts
+{
+  // Expensive in compilation time...
+  //static_assert(is_congruent<Shape, Stride>::value, "Shape and Stride must be congruent");
+
+  // NOTE: This defaults static Shapes/Strides correctly, but not dynamic
+  MUTE_HOST_DEVICE constexpr
+  Layout(Shape  const& shape  = {}, Stride const& stride = {})
+      : mute::tuple<Shape, Stride>(shape, stride)
+  {}
+
+  //
+  // Accessors
+  //
+
+  static constexpr int rank  = rank_v<Shape>;
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  layout() {
+    return *this;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  layout() const {
+    return *this;
+  }
+
+  template <int... I>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  shape() {
+    return get<0,I...>(static_cast<mute::tuple<Shape, Stride>&>(*this));
+  }
+
+  template <int... I>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  shape() const {
+    return get<0,I...>(static_cast<mute::tuple<Shape, Stride> const&>(*this));
+  }
+
+  template <int... I>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  stride() {
+    return get<1,I...>(static_cast<mute::tuple<Shape, Stride>&>(*this));
+  }
+
+  template <int... I>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  stride() const {
+    return get<1,I...>(static_cast<mute::tuple<Shape, Stride> const&>(*this));
+  }
+
+  //
+  // Mappings
+  //
+
+  // Map a logical coordinate to a linear index (Coord has no Underscore slice operators)
+  // OR
+  // Slice the layout and return the sublayout (Coord has an Underscore slice op)
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator()(Coord const& coord) const {
+    if constexpr (has_underscore<Coord>::value) {
+      return slice(coord, *this);
+    } else {
+      return crd2idx(coord, shape(), stride());
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+
+  // Convenience function for multi-dimensional coordinates
+  template <class Coord0, class Coord1, class... Coords>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const {
+    return operator()(make_coord(c0,c1,cs...));
+  }
+
+  //
+  // Compose
+  //
+
+  template <class OtherLayout>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  compose(OtherLayout const& other) const {
+    return composition(*this, other);
+  }
+
+  template <class... Layouts>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  compose(Layouts const&... layouts) const {
+    return composition(*this, make_tile(layouts...));
+  }
+
+  template <class OtherShape>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  with_shape(OtherShape const& shape) const {
+    return composition(*this, make_layout(shape));
+  }
+
+  template <class... Shapes>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  with_shape(Shapes const&... shapes) const {
+    return composition(*this, make_layout(make_shape(shapes...)));
+  }
+
+  //
+  // Tile
+  //
+
+  template <class OtherLayout>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  tile(OtherLayout const& other) const {
+    return tiled_divide(*this, other);
+  }
+
+  template <class... Layouts>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  tile(Layouts const&... layouts) const {
+    return tiled_divide(*this, make_tile(layouts...));
+  }
+
+  //
+  // Utility
+  //
+
+  //
+  // Index to Coordinate
+  //
+
+  // NOTE: Only valid for compact layouts
+
+  // Return the (hierarchical) ND logical coordinate corresponding to the linear index
+  // @post crd2idx(@a result, shape(), stride()) == idx
+  // @post congruent(@a result, shape())
+  template <class IInt,
+            __MUTE_REQUIRES(is_integral<IInt>::value)>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_hier_coord(IInt const& idx) const {
+    return mute::idx2crd(idx, shape(), stride());
+  }
+
+  // Return the (flat) ND logical coordinate corresponding to the linear index
+  // @post crd2idx(@a result, shape(), stride()) == idx
+  // @post rank(@a result) == rank(shape()) && depth(@a result) == 1
+  template <class IInt,
+            __MUTE_REQUIRES(is_integral<IInt>::value)>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_flat_coord(IInt const& idx) const {
+    return mute::crd2crd(this->get_hier_coord(idx), shape(), repeat<rank>(Int<1>{}));
+  }
+
+  // Return the generalized column-major 1D logical coordinate corresponding to the linear index
+  // @post crd2idx(@a result, shape(), stride()) == idx
+  // @post is_integral<decltype(@a result)>::value
+  template <class IInt,
+            __MUTE_REQUIRES(is_integral<IInt>::value)>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_1d_coord(IInt const& idx) const {
+    return mute::crd2idx(this->get_hier_coord(idx), shape());
+  }
+
+  //
+  // Coordinate to Coordinate
+  //
+
+#if 0
+  // Return the (hierarchical) ND logical coordinate corresponding to the linear index
+  // @post congruent(@a result, shape())
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  crd_2_hier_coord(Coord const& crd) const {
+    return mute::crd2crd(crd, shape(), shape());
+  }
+
+  // Return the (flat) ND logical coordinate corresponding to the linear index
+  // @post rank(@a result) == rank(shape()) && depth(@a result) == 1
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  crd_2_flat_coord(Coord const& crd) const {
+    return mute::crd2crd(crd, shape(), product_each(shape()));
+  }
+
+  // Return the generalized column-major 1D logical coordinate corresponding to the linear index
+  // @post is_integral<decltype(@a result)>::value
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  crd_2_1d_coord(Coord const& crd) const {
+    //return mute::crd2crd(crd, shape(), product(shape()));
+    return mute::crd2idx(crd, shape());
+  }
+#endif
+};
+
+// Equality, return a static or dynamic boolean
+template <class ShapeA, class StrideA,
+          class ShapeB, class StrideB>
+MUTE_HOST_DEVICE constexpr
+auto
+operator==(Layout<ShapeA,StrideA> const& layoutA, Layout<ShapeB,StrideB> const& layoutB)
+{
+  return layoutA.shape() == layoutB.shape() && layoutA.stride() == layoutB.stride();
+}
+
+template <class Layout>
+struct is_layout : false_type {};
+template <class Shape, class Stride>
+struct is_layout<Layout<Shape,Stride>> : true_type {};
+
+//
+// Layout construction
+//
+
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout(Shape const& shape, Stride const& stride)
+{
+  static_assert(is_tuple<Shape >::value || is_integral<Shape >::value);
+  static_assert(is_tuple<Stride>::value || is_integral<Stride>::value);
+  return Layout<Shape,Stride>(shape, stride);
+}
+
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout(Shape const& shape)
+{
+  static_assert(is_tuple<Shape >::value || is_integral<Shape >::value);
+  return make_layout(shape, compact_major<LayoutLeft>(shape));
+}
+
+//
+// Convenience tags for common layouts
+//
+
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout(Shape const& shape, LayoutLeft)
+{
+  return make_layout(shape, compact_major<LayoutLeft>(shape));
+}
+
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout(Shape const& shape, LayoutRight)
+{
+  return make_layout(shape, compact_major<LayoutRight>(shape));
+}
+
+//
+// Construct a layout from multiple layouts by concatenation
+//
+
+// One argument overload
+template <class Shape0, class Stride0>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout(Layout<Shape0,Stride0> const& layout0)
+{
+  return make_layout(make_shape (layout0.shape() ),
+                     make_stride(layout0.stride()));
+}
+
+// Two argument overload
+template <class Shape0, class Stride0,
+          class Shape1, class Stride1>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout(Layout<Shape0,Stride0> const& layout0,
+            Layout<Shape1,Stride1> const& layout1)
+{
+  return make_layout(make_shape (layout0.shape() , layout1.shape() ),
+                     make_stride(layout0.stride(), layout1.stride()));
+}
+
+// Var argument overload
+template <class Shape0, class Stride0,
+          class Shape1, class Stride1,
+          class... Shapes, class... Strides>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout(Layout<Shape0,Stride0> const& layout0,
+            Layout<Shape1,Stride1> const& layout1,
+            Layout<Shapes,Strides> const&... layouts)
+{
+  return make_layout(make_shape (layout0.shape() , layout1.shape() , layouts.shape()... ),
+                     make_stride(layout0.stride(), layout1.stride(), layouts.stride()...));
+}
+
+//
+// Advanced Layout constructions
+//
+
+// Make a compact layout with shape @a shape and strides following the order induced by @a order.
+// Dynamic values in @a order are ignored, considered large, and considered ordered from left to right.
+// Example:
+//   make_ordered_layout(Shape<_2,_2,_2,_2>{}, Step<_0,_2,_3,_1>{})
+//     ->  (_2,_2,_2,_2):(_1,_4,_8,_2)
+//   make_ordered_layout(make_shape(2,3,4,5), make_step(Int<2>{}, 67, 42, Int<50>{}))
+//     -> (2,3,4,5):(_1,10,30,2)
+template <class Shape, class Order>
+MUTE_HOST_DEVICE constexpr
+auto
+make_ordered_layout(Shape const& shape, Order const& order)
+{
+  return make_layout(shape, compact_order(shape, order));
+}
+
+// Make a compact layout with the same shape as @a layout
+//   and strides following the order induced by @a layout.stride().
+// Static-0 strides in the input @a layout are preserved in the output.
+// Example:
+//   make_layout_like(Layout<Shape<_2,_2,_2,_2>, Stride<_0,_2,_4,_1>>{})
+//     ->  (_2,_2,_2,_2):(_0,_2,_4,_1)
+//   make_layout_like(make_layout(make_shape(2,3,4,5), make_stride(Int<0>{},42,Int<1>{},Int<0>{})))
+//     -> (2,3,4,5):(_0,4,_1,_0)
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout_like(Layout<Shape,Stride> const& layout)
+{
+  return make_layout(layout.shape(),
+                     compact_order(filter_zeros(layout.stride(), layout.shape()), layout.stride()));
+}
+
+// Make a compact layout with the same shape as @a layout
+//   and strides following the order induced by @a layout.stride(),
+//   except mode-0 is always stride-1 and generated column-major.
+// The 0th mode is commonly used for MMA_Atoms or Copy_Atoms so this
+//   generates the 0th mode with LayoutLeft (preserving stride-0s) regardless of the reference layout
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+make_fragment_like(Layout<Shape,Stride> const& layout)
+{
+  constexpr int R = Layout<Shape,Stride>::rank;
+  if constexpr (R > 1 && is_static<Shape>::value) {
+    return tiled_product(make_layout(get<0>(layout.shape()),
+                                     compact_major<LayoutLeft>(filter_zeros(get<0>(layout.stride()), get<0>(layout.shape())))),
+                         make_ordered_layout(take<1,R>(layout.shape()), take<1,R>(layout.stride())));
+  } else {
+    return make_layout(layout.shape());
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class Shape,
+          __MUTE_REQUIRES(is_tuple<Shape>::value || is_integral<Shape>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+make_fragment_like(Shape const& shape)
+{
+  return make_layout(shape);
+}
+
+//
+// Make an identity layout that maps a coordinate to itself
+//
+
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+make_identity_layout(Shape const& shape)
+{
+  return make_layout(shape, make_basis_like(shape));
+}
+
+//
+// Operations to manipulate Layouts like a tuple of pairs
+//
+
+// Return the Is...th sublayout.
+// For Is... = <I0,I1,...,IN>, equivalent to get<IN>(...get<I1>(get<I0>(layout)))
+template <size_t... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+get(Layout<Shape,Stride> const& layout)
+{
+  return make_layout(get<Is...>(layout.shape()),
+                     get<Is...>(layout.stride()));
+}
+
+// Return a new layout with only the modes in the range [B,E)
+template <int B, int E, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+take(Layout<Shape,Stride> const& layout)
+{
+  static_assert(B < E, "take: empty range error");
+  static_assert(0 <= B && E <= Layout<Shape,Stride>::rank, "take: range out of bounds");
+  return make_layout(take<B,E>(layout.shape()),
+                     take<B,E>(layout.stride()));
+}
+
+// Return a new layout with only the modes Is... = <I0,I1,...,IN>
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+select(Layout<Shape,Stride> const& layout)
+{
+  return make_layout(select<Is...>(layout.shape()),
+                     select<Is...>(layout.stride()));
+}
+
+// Return a layout with depth at most 1
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+flatten(Layout<Shape,Stride> const& layout)
+{
+  return make_layout(flatten(layout.shape()),
+                     flatten(layout.stride()));
+}
+
+// Return a layout whose profile is congruent to TargetProfile
+// @pre Input layout is flat, flatten(@a layout) == @a layout
+// @pre Input layout can be folded to profile, rank(@a layout) == rank(flatten(@a target_profile))
+// @post congruent(@a result, @a target_profile)
+template <class Shape, class Stride, class TargetProfile>
+MUTE_HOST_DEVICE constexpr
+auto
+unflatten(Layout<Shape,Stride> const& layout, TargetProfile const& target_profile)
+{
+  return make_layout(unflatten(layout.shape(),  target_profile),
+                     unflatten(layout.stride(), target_profile));
+}
+
+//
+// Utilities
+//
+
+// Return the sublayout of mode I...
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+layout(Layout<Shape,Stride> const& layout)
+{
+  if constexpr (sizeof...(Is) == 0) {
+    return layout;
+  } else {
+    return get<Is...>(layout);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Return the shape of a mode
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+shape(Layout<Shape,Stride>& layout)
+{
+  return layout.template shape<Is...>();
+}
+
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+shape(Layout<Shape,Stride> const& layout)
+{
+  return layout.template shape<Is...>();
+}
+
+// Return the stride of a mode
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+stride(Layout<Shape,Stride>& layout)
+{
+  return layout.template stride<Is...>();
+}
+
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+stride(Layout<Shape,Stride> const& layout)
+{
+  return layout.template stride<Is...>();
+}
+
+// Return the number of elements in a mode
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+size(Layout<Shape,Stride> const& layout)
+{
+  return size(shape<Is...>(layout));
+}
+
+// Return the number of modes
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+rank(Layout<Shape,Stride> const& layout)
+{
+  return rank(shape<Is...>(layout));
+}
+
+// Return the depth of the layout
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+depth(Layout<Shape,Stride> const& layout)
+{
+  return depth(shape<Is...>(layout));
+}
+
+// Return the codomain shape of a mode
+// @post size(coshape(@a a)) == cosize(@a a)
+// @return C Coordinate with smallest elements such that
+//           @a elem_less(sub_layout(c), C) for all c < size(@a sub_layout)
+//           where sub_layout = get<Is...>(layout).
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+coshape(Layout<Shape,Stride> const& layout)
+{
+  // Protect against negative strides
+  auto abs_sub_layout = make_layout(shape<Is...>(layout),
+                                    transform_leaf(stride<Is...>(layout), abs_fn{}));
+  auto co_coord = as_arithmetic_tuple(abs_sub_layout(size(abs_sub_layout) - Int<1>{}));
+  return co_coord + repeat_like(co_coord, Int<1>{});
+}
+
+// Return the codomain size of a mode
+// @return M smallest integer such that
+//           @a sub_layout(c) < M for all c < size(@a sub_layout)
+//           where sub_layout = get<Is...>(layout).
+template <int... Is, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+cosize(Layout<Shape,Stride> const& layout)
+{
+  return size(coshape<Is...>(layout));
+}
+
+template <class Layout>
+using cosize_t = decltype(cosize(declval<Layout>()));
+
+template <class Layout>
+static constexpr auto cosize_v = cosize_t<Layout>::value;
+
+// With crd2idx(coord, shape), makes sense to have crd2idx(coord, Layout) as well
+template <class Coord, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+crd2idx(Coord const& c, Layout<Shape,Stride> const& layout)
+{
+  return crd2idx(c, layout.shape(), layout.stride());
+}
+
+//
+// Slice and Dice a layout
+//
+
+template <class Coord, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+slice(Coord const& c, Layout<Shape,Stride> const& layout)
+{
+  return make_layout(slice(c, layout.shape()),
+                     slice(c, layout.stride()));
+}
+
+template <class Coord, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+slice_and_offset(Coord const& c, Layout<Shape,Stride> const& layout)
+{
+  return mute::make_tuple(slice(c, layout), crd2idx(c, layout));
+}
+
+template <class Coord, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+dice(Coord const& c, Layout<Shape,Stride> const& layout)
+{
+  return make_layout(dice(c, layout.shape()),
+                     dice(c, layout.stride()));
+}
+
+// Compute a pointer offset and (potentially modified) layout from a coordinate
+// This exists so it can be overloaded for ComposedLayout
+template <class Coord, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+domain_offset(Coord const& coord, Layout<Shape,Stride> const& layout)
+{
+  return mute::make_tuple(layout, layout(coord));
+}
+
+//
+// Transform the modes of a layout
+//
+
+namespace detail {
+
+template <class Tuple, class F, int... I>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_layout(Tuple const& t, F&& f, seq<I...>)
+{
+  return make_layout(f(get<I>(t))...);
+}
+
+template <class Tuple0, class Tuple1, class F, int... I, int... I0, int... I1>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_layout(Tuple0 const& t0, Tuple1 const& t1, F&& f, seq<I...>, seq<I0...>, seq<I1...>)
+{
+  return make_layout(f(get<I>(t0),get<I>(t1))..., get<I0>(t0)..., get<I1>(t1)...);
+}
+
+} // end namespace detail
+
+template <class Tuple, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_layout(Tuple const& t, F&& f)
+{
+  return detail::transform_layout(t, f, make_seq<decltype(rank(t))::value>{});
+}
+
+template <class Tuple0, class Tuple1, class F>
+MUTE_HOST_DEVICE constexpr
+auto
+transform_layout(Tuple0 const& t0, Tuple1 const& t1, F&& f)
+{
+  constexpr int R0 = decltype(rank(t0))::value;
+  constexpr int R1 = decltype(rank(t1))::value;
+  constexpr int R  = (R0 < R1) ? R0 : R1;
+  return detail::transform_layout(t0, t1, f, make_seq<R>{}, make_range<R,R0>{}, make_range<R,R1>{});
+}
+
+//
+// Coalesce and Filter
+//
+
+namespace detail {
+
+// Look at each element and the front of the stack (in order of priority)
+// front(NewLayout)  get<I>(Layout)
+//      s0:d0           _1:d1     =>  continue
+//      _1:d0           s1:d1     =>  replace_front    s1:d1
+//      s0:s1*d1        s1:d1     =>  replace_front s0*s1:d1
+//      s0:d0           s1:d1     =>  prepend          s1:d1
+//
+// @pre OldShape and OldStride are flat
+template <int I, class OldShape, class OldStride, class NewShape, class NewStride>
+MUTE_HOST_DEVICE constexpr
+auto
+bw_coalesce(OldShape const& old_shape, OldStride const& old_stride,
+            NewShape const& new_shape, NewStride const& new_stride)
+{
+  if constexpr (I == -1) {
+    // Base case, we're done
+    if constexpr (is_constant<1, NewShape>::value) {
+      return Layout<_1,_0>{};
+    } else {
+      return Layout<NewShape,NewStride>{new_shape,new_stride};
+    }
+  } else if constexpr (is_constant<1, decltype(get<I>(old_shape))>::value) {
+    // shape<I>(layout) == _1, skip it and continue
+    return bw_coalesce<I-1>(old_shape, old_stride, new_shape, new_stride);
+  } else if constexpr (is_constant<1, NewShape>::value) {
+    // Replace our shape-1 with anything (Can only happen on input new_shape/new_stride)
+    return bw_coalesce<I-1>(old_shape, old_stride, get<I>(old_shape), get<I>(old_stride));
+  } else if constexpr (is_static<decltype(get<0>(new_shape))>::value &&
+                       is_constant<true, decltype(get<I>(old_shape) * get<I>(old_stride) == get<0>(new_stride))>::value) {
+    // Merge modes because the shapes and strides match
+    return bw_coalesce<I-1>(old_shape, old_stride,
+                            replace_front(new_shape,  get<I>(old_shape) * get<0>(new_shape)),
+                            replace_front(new_stride, get<I>(old_stride)));
+  } else {
+    // Can't replace or merge, so prepend a new mode
+    return bw_coalesce<I-1>(old_shape, old_stride,
+                            prepend(new_shape,  get<I>(old_shape)),
+                            prepend(new_stride, get<I>(old_stride)));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// mute::coalesce promises to not change the Layout as a function from integers to codomain.
+// It accomplishes this inside of the Layout's domain, but not always outside of the domain.
+//   Example: (_4,_1):(_1,_0) coalesces to _4:_1.
+// detail::coalesce_x preserves the Layout function inside its domain and outside.
+//
+// @post depth(@a result) <= 1
+// @post for all i, 0 <= i, @a layout(i) == @a result(i)
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce_x(Layout<Shape,Stride> const& layout)
+{
+  auto flat_shape  = flatten(layout.shape());
+  auto flat_stride = flatten(layout.stride());
+
+  constexpr int R = decltype(rank(flat_shape))::value;
+  if constexpr (is_constant<1, decltype(get<R-1>(flat_shape))>::value) {
+    return detail::bw_coalesce<R-2>(flat_shape, flat_stride,             Int<2>{}, get<R-1>(flat_stride));
+  } else {
+    return detail::bw_coalesce<R-2>(flat_shape, flat_stride, get<R-1>(flat_shape), get<R-1>(flat_stride));
+  }
+}
+
+// Apply coalesce_x at the terminals of trg_profile
+template <class Shape, class Stride, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce_x(Layout<Shape,Stride> const& layout, IntTuple const& trg_profile)
+{
+  if constexpr (is_tuple<IntTuple>::value) {
+    static_assert(tuple_size<IntTuple>::value <= Layout<Shape,Stride>::rank);
+    return mute::transform_layout(layout, trg_profile, [](auto const& l, auto const& t) { return coalesce_x(l,t); });
+  } else {
+    return coalesce_x(layout);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+// "Simplify" the layout by combining modes that are possible to combine
+// Does not respect the shape of the layout, but does preserve total size
+// @post size(@a result) == size(@a layout)
+// @post depth(@a result) <= 1
+// @post for all i, 0 <= i < size(@a layout), @a layout(i) == @a result(i)
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce(Layout<Shape,Stride> const& layout)
+{
+  auto flat_shape  = flatten(layout.shape());
+  auto flat_stride = flatten(layout.stride());
+
+  constexpr int R = decltype(rank(flat_shape))::value;
+  return detail::bw_coalesce<R-2>(flat_shape, flat_stride, get<R-1>(flat_shape), get<R-1>(flat_stride));
+}
+
+// Apply coalesce at the terminals of trg_profile
+template <class Shape, class Stride, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce(Layout<Shape,Stride> const& layout, IntTuple const& trg_profile)
+{
+  if constexpr (is_tuple<IntTuple>::value) {
+    static_assert(tuple_size<IntTuple>::value <= Layout<Shape,Stride>::rank);
+    return transform_layout(layout, trg_profile, [](auto const& l, auto const& t) { return coalesce(l,t); });
+  } else {
+    return coalesce(layout);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Combine static and dynamic modes of a shape.
+// @post size(@a result) == size(@a shape)
+// @post depth(@a result) <= 1
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce(Shape const& shape)
+{
+  static_assert(is_integral<Shape>::value || is_tuple<Shape>::value);
+
+  return mute::fold_first(flatten(shape), [](auto const& init, auto const& a) {
+    if constexpr (is_static<decltype(back(init))>::value == is_static<decltype(a)>::value) {
+      return replace_back(init, back(init) * a);  // Both static or both dynamic, coalesce and replace
+    } else {
+      return append(init, a);                     // Can't coalesce, so append
+    }
+  });
+}
+
+// Replace the modes in layout that have a 0-stride with a 1-size
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+filter_zeros(Layout<Shape,Stride> const& layout)
+{
+  return make_layout(filter_zeros(layout.stride(), layout.shape()), layout.stride());
+}
+
+// Replace the modes in layout that correspond to a 0 at the terminals of trg_profile with a 1-size
+template <class Shape, class Stride, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+filter_zeros(Layout<Shape,Stride> const& layout, IntTuple const& trg_profile)
+{
+  return make_layout(filter_zeros(trg_profile, layout.shape()), layout.stride());
+}
+
+// Remove all of the 0-strides and 1-sizes
+// Return 1-shape if empty
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+filter(Layout<Shape,Stride> const& layout)
+{
+  return coalesce(filter_zeros(layout));
+}
+
+// Apply filter at the terminals of trg_profile
+template <class Shape, class Stride, class IntTuple>
+MUTE_HOST_DEVICE constexpr
+auto
+filter(Layout<Shape,Stride> const& layout, IntTuple const& trg_profile)
+{
+  if constexpr (is_tuple<IntTuple>::value) {
+    static_assert(tuple_size<IntTuple>::value <= Layout<Shape,Stride>::rank);
+    return transform_layout(layout, trg_profile, [](auto const& l, auto const& t) { return filter(l,t); });
+  } else {
+    return filter(layout);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Append, Prepend, Replace
+//
+
+template <int N, class ShapeA, class StrideA, class ShapeX = _1, class StrideX = _0>
+MUTE_HOST_DEVICE constexpr
+auto
+append(Layout<ShapeA,StrideA> const& layout,
+       Layout<ShapeX,StrideX> const& x = {})
+{
+  return make_layout(append<N>(layout.shape(),  x.shape()),
+                     append<N>(layout.stride(), x.stride()));
+}
+
+template <class ShapeA, class StrideA, class ShapeX = _1, class StrideX = _0>
+MUTE_HOST_DEVICE constexpr
+auto
+append(Layout<ShapeA,StrideA> const& layout,
+       Layout<ShapeX,StrideX> const& x = {})
+{
+  return make_layout(append(layout.shape(),  x.shape()),
+                     append(layout.stride(), x.stride()));
+}
+
+template <int N, class ShapeA, class StrideA, class ShapeX = _1, class StrideX = _0>
+MUTE_HOST_DEVICE constexpr
+auto
+prepend(Layout<ShapeA,StrideA> const& layout,
+        Layout<ShapeX,StrideX> const& x = {})
+{
+  return make_layout(prepend<N>(layout.shape(),  x.shape()),
+                     prepend<N>(layout.stride(), x.stride()));
+}
+
+template <class ShapeA, class StrideA, class ShapeX = _1, class StrideX = _0>
+MUTE_HOST_DEVICE constexpr
+auto
+prepend(Layout<ShapeA,StrideA> const& layout,
+        Layout<ShapeX,StrideX> const& x = {})
+{
+  return make_layout(prepend(layout.shape(),  x.shape()),
+                     prepend(layout.stride(), x.stride()));
+}
+
+template <int N, class ShapeA, class StrideA, class ShapeX, class StrideX>
+MUTE_HOST_DEVICE constexpr
+auto
+replace(Layout<ShapeA,StrideA> const& layout,
+        Layout<ShapeX,StrideX> const& x)
+{
+  return make_layout(replace<N>(layout.shape(),  x.shape()),
+                     replace<N>(layout.stride(), x.stride()));
+}
+
+template <int B, int E, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+group(Layout<Shape,Stride> const& layout)
+{
+  return make_layout(group<B,E>(layout.shape()),
+                     group<B,E>(layout.stride()));
+}
+
+//
+// Composition of two layouts: lhs o rhs
+// @post compatible(rhs, result)
+// @post result(c) = lhs(rhs(c))
+//         for all c in the domain of rhs
+//
+
+namespace detail {
+
+template <class LShape, class LStride,
+          class RShape, class RStride>
+MUTE_HOST_DEVICE constexpr
+auto
+composition_impl(LShape const& lhs_shape, LStride const& lhs_stride,
+                 RShape const& rhs_shape, RStride const& rhs_stride)
+{
+  if constexpr (is_tuple<RShape>::value) {
+    // Apply the right-distributivity of Layout composition
+    return transform_layout(rhs_shape, rhs_stride, [&](auto const& s, auto const& d) {
+      return composition_impl(lhs_shape, lhs_stride, s, d);
+    });
+  } else
+  if constexpr (is_scaled_basis<RStride>::value) {
+    // Special case for a ScaledBasis stride
+    return composition_impl(basis_get(rhs_stride, lhs_shape), basis_get(rhs_stride, lhs_stride),
+                            rhs_shape, basis_value(rhs_stride));
+  } else
+  if constexpr (is_constant<0, RStride>::value) {
+    // Special case shortcut for any static stride-0
+    return Layout<RShape, RStride>{rhs_shape, rhs_stride};
+  } else
+  if constexpr (is_integral<decltype(lhs_shape)>::value) {
+    // Special case shortcut for any integral LShape
+    return Layout{rhs_shape, rhs_stride * lhs_stride};
+  } else
+  if constexpr (is_constant<1, RStride>::value) {
+    // Special case shortcut for any static stride-1
+    constexpr int R  = rank_v<LShape>;
+    auto result_shape_0  = take<0,R-1>(lhs_shape);
+
+    // Mod out the rhs_shape from the lhs_shape
+    auto const [result_shape_1, rest_shape]  = fold(result_shape_0, mute::make_tuple(mute::make_tuple(), rhs_shape),
+      [] (auto const& init, auto const& si) {
+        return mute::make_tuple(append(get<0>(init), shape_min(abs(si), get<1>(init))), shape_div(get<1>(init), abs(si)));
+      });
+
+    // Jump into coalesce and append (rest_shape, get<R-1>(lhs_stride))
+    return detail::bw_coalesce<R-2>(result_shape_1, lhs_stride, rest_shape, get<R-1>(lhs_stride));
+  } else {
+    // General case: integral RShape and RStride, tuple LShape and LStride
+    constexpr int R  = rank_v<LShape>;
+    auto result_shape_0  = take<0,R-1>(lhs_shape);
+    auto result_stride_0 = take<0,R-1>(lhs_stride);
+
+    // Divide out the rhs_stride from the lhs_shape
+    auto const [result_shape_1, rest_stride] = fold(result_shape_0, mute::make_tuple(mute::make_tuple(), rhs_stride),
+      [] (auto const& init, auto const& di) {
+        return mute::make_tuple(append(get<0>(init), shape_div(di, get<1>(init))), shape_div(get<1>(init), di));
+      });
+
+    // Apply any lhs_shape changes to the stride
+    auto result_stride_1 = elem_scale(result_stride_0, shape_div(result_shape_0, result_shape_1));
+
+    // Mod out the rhs_shape from the lhs_shape
+    auto const [result_shape_2, rest_shape] = fold(result_shape_1, mute::make_tuple(mute::make_tuple(), rhs_shape),
+      [] (auto const& init, auto const& si) {
+        return mute::make_tuple(append(get<0>(init), shape_min(abs(si), get<1>(init))), shape_div(get<1>(init), abs(si)));
+      });
+
+    // Jump into coalesce and append (rest_shape, rest_stride * get<R-1>(lhs_stride))
+    return detail::bw_coalesce<R-2>(result_shape_2, result_stride_1, rest_shape, rest_stride * get<R-1>(lhs_stride));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+template <class LShape, class LStride,
+          class RShape, class RStride>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(Layout<LShape,LStride> const& lhs,
+            Layout<RShape,RStride> const& rhs)
+{
+  auto coprofile = repeat_like(decltype(coshape(rhs)){}, Int<0>{});
+  auto flat_lhs = detail::coalesce_x(lhs, coprofile);
+  return detail::composition_impl(flat_lhs.shape(), flat_lhs.stride(), rhs.shape(), rhs.stride());
+}
+
+template <class LShape, class LStride, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(Layout<LShape,LStride> const& lhs,
+            Tiler                  const& rhs)
+{
+  if constexpr (is_tuple<Tiler>::value) {
+    static_assert(tuple_size<Tiler>::value <= Layout<LShape,LStride>::rank);
+    // Drop any modes of lhs that aren't hit by rhs
+    return detail::transform_layout(lhs, rhs, [](auto const& l, auto const& r) { return composition(l,r); }, make_seq<tuple_size<Tiler>::value>{}, seq<>{}, seq<>{});
+  } else if constexpr (is_underscore<Tiler>::value) {
+    return lhs;
+  } else if constexpr (is_integral<Tiler>::value) {
+    auto flat_lhs = detail::coalesce_x(lhs);
+    return detail::composition_impl(flat_lhs.shape(), flat_lhs.stride(), rhs, Int<1>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Complement
+//
+// Build the complement of a layout.
+// @post size(@a result) >= @a cosize_hi / size(filter(@a layout)));
+// @post For all i in [1,size(@a result)),
+//           @a result(i) < @a result(i-1)
+//           For all j in [0, size(@a layout)),
+//               @a result(i) != @a layout(j)
+//
+
+namespace detail {
+
+// @pre @a layout has been filtered (flattened and no stride-0 or size-1 modes).
+template <class Shape, class Stride, class CoTarget>
+MUTE_HOST_DEVICE constexpr
+auto
+complement(Shape const& shape, Stride const& stride, CoTarget const& cotarget)
+{
+  if constexpr (is_constant<0, Stride>::value) {
+    // Special case for irreducible rank-1 stride-0 layout
+    return make_layout(coalesce(cotarget));
+  } else {
+    // General case
+    constexpr int R = rank_v<Shape>;
+    static_assert(R == 1 || is_static<Stride>::value,
+                  "Dynamic-stride complement only for rank-1 layouts");
+
+    // Should just be a sort and a fold...
+    // Then we could even handle dynamic strides (but they would destroy all static strides)
+    auto [shape_, stride_, result_shape_, result_stride] =
+      fold(make_seq<R-1>{},
+           mute::make_tuple(shape, stride, mute::make_tuple(), mute::make_tuple(Int<1>{})),
+           [](auto const& init, auto i)
+           {
+              auto [shape, stride, result_shape, result_stride] = init;
+              auto min_stride = mute::min(stride);
+              auto min_idx    = mute::find(stride, min_stride);
+              auto new_shape  = min_stride / get<i>(result_stride);
+              auto new_stride = min_stride * get<min_idx>(shape);
+              static_assert(not is_constant<0, decltype(new_shape)>::value, "Non-injective Layout detected in complement.");
+
+              return mute::make_tuple(remove<min_idx>(shape),              // Remove the min_idx from shape
+                                      remove<min_idx>(stride),             // Remove the min_idx from stride
+                                      append(result_shape , new_shape ),   // new shape  = min_stride / last_stride
+                                      append(result_stride, new_stride));  // new stride = min_stride * curr_shape
+            });
+
+    // Append the last shape mode
+    auto new_shape    = get<0>(stride_) / get<R-1>(result_stride);         // new shape  = min_stride / last_stride
+    static_assert(not is_constant<0, decltype(new_shape)>::value, "Non-injective Layout detected in complement.");
+    auto result_shape = append(result_shape_, new_shape);
+
+    // Compute the rest_shape and rest_stride
+    auto new_stride  = get<0>(stride_) * get<0>(shape_);                   // new stride = min_stride * curr_shape
+    auto rest_shape  = coalesce(ceil_div(cotarget, new_stride));
+    auto rest_stride = compact_major<LayoutLeft>(rest_shape, new_stride);
+
+    // Coalesce and append (rest_shape, rest_stride)
+    return coalesce(make_layout(make_shape (result_shape , rest_shape ),
+                                make_stride(result_stride, rest_stride)));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+template <class Shape, class Stride, class CoTarget>
+MUTE_HOST_DEVICE constexpr
+auto
+complement(Layout<Shape,Stride> const& layout, CoTarget const& cotarget)
+{
+  auto filter_layout = filter(layout);
+  return detail::complement(filter_layout.shape(), filter_layout.stride(), shape(cotarget));
+}
+
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+complement(Layout<Shape,Stride> const& layout)
+{
+  auto filter_layout = filter(layout);
+  return detail::complement(filter_layout.shape(), filter_layout.stride(), cosize(filter_layout));
+}
+
+//
+// Right-Inverse and Left-Inverse
+//
+
+namespace detail {
+
+template <int NextStride, class Shape, class Stride, int... Is>
+MUTE_HOST_DEVICE constexpr
+auto
+inverse_seq(Shape const& shape, Stride const& stride, seq<Is...>)
+{
+  auto next_I = mute::find_if(stride, [](auto a) { return is_constant<NextStride, decltype(a)>{}; });
+
+  if constexpr (next_I == decltype(rank(stride))::value) {
+    // If not found, return current seq
+    return seq<Is...>{};
+  } else {
+    // auto next_stride = get<next_I>(shape) * get<next_I>(stride);
+    // NOTE: Needed for g++-7
+    using next_stride = decltype(get<next_I>(shape) * get<next_I>(stride));
+
+    if constexpr (is_static<next_stride>::value && !is_constant<NextStride, next_stride>::value) {
+      // If next_stride is static and unique, then continue
+      return inverse_seq<next_stride::value>(shape, stride, seq<Is..., next_I>{});
+    } else {
+      // Else return current seq + next_I
+      return seq<Is..., next_I>{};
+    }
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+//
+// Build the right-inverse of a layout
+// @pre is_static<Layout>
+// @result A layout @a result such that
+//    @a layout(@a result(i)) == i for all i < size(@a result)
+// @result A layout @a result such that
+//    composition(@a layout, @a result) is identical to make_layout(shape(result))
+//
+
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+right_inverse(Layout<Shape,Stride> const& layout)
+{
+  auto flat_layout = coalesce(layout);
+  auto astride = transform_leaf(flat_layout.stride(), abs_fn{});
+
+  // Find Int<1>{}, the starting stride, and follow the strides to gen inverse_seq
+  [[maybe_unused]] auto iseq = detail::inverse_seq<1>(flat_layout.shape(), astride, seq<>{});
+
+  if constexpr (iseq.size() == 0) {
+    return Layout<_1,_0>{};     // Empty case, nothing found
+  } else {
+    // Generate the corresponding new strides and construct
+    auto rstride = compact_major<LayoutLeft>(flat_layout.shape());
+    return make_layout(unwrap(transform(iseq, [&](auto i) { return shape<i>(flat_layout); })),
+                       unwrap(transform(iseq, [&](auto i) { return signum(stride<i>(flat_layout)) * get<i>(rstride); })));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+MUTE_HOST_DEVICE constexpr
+auto
+right_inverse(Underscore const& _)
+{
+  return _;
+}
+
+//
+// Build the left-inverse of a layout
+// @pre is_static<Layout>
+// @pre @a layout is an injective function
+// @result A layout @a result such that
+//    @a result(@a layout(i)) == i for all i < size(@a layout)
+// @result A layout @a result such that
+//    composition(@a result, @a layout) is identical to make_layout(shape(layout))
+//
+
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+left_inverse(Layout<Shape,Stride> const& layout)
+{
+  return right_inverse(make_layout(layout, complement(layout)));
+}
+
+MUTE_HOST_DEVICE constexpr
+auto
+left_inverse(Underscore const& _)
+{
+  return _;
+}
+
+//
+// Max Common Layout
+//
+
+/* Return a layout that points to the maximum number of contiguous elements
+ * that logically correspond in the layouts of @a a and @a b.
+ *
+ * @returns Layout R
+ * @post For all 0 <= i < size(R), a(R(i)) == i and b(R(i)) == i
+ */
+template <class ShapeA, class StrideA,
+          class ShapeB, class StrideB>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_layout(Layout<ShapeA,StrideA> const& a,
+                  Layout<ShapeB,StrideB> const& b)
+{
+  Layout inv_b  = right_inverse(b);
+  Layout common = coalesce(composition(a, inv_b));
+
+  // Keep only the static identity component of the common layout
+  if constexpr (is_static<decltype(shape<0>(common))>::value &&
+                is_constant<1, decltype(stride<0>(common))>::value) {
+    // Truncate to the size of the contiguous vector (static stride-1 mode)
+    return composition(inv_b, layout<0>(common));
+  } else {
+    return Layout<_1,_0>{};
+  }
+}
+
+/* Return Int<N> such that N is the maximum number of contiguous elements
+ * that logically correspond in the layouts of @a a and @a b.
+ *
+ * @returns Int<N> with N >= 1
+ * @post For all 0 <= n < N, a(b.get_1d_coord(n)) == n
+ *       (NOTE: Problems with negative strides/coords in this post-condition)
+ */
+template <class ShapeA, class StrideA,
+          class ShapeB, class StrideB>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_vector(Layout<ShapeA,StrideA> const& a,
+                  Layout<ShapeB,StrideB> const& b)
+{
+  Layout common = coalesce(composition(a, right_inverse(b)));
+
+  // Keep only the static identity component of the common layout
+  if constexpr (is_static<decltype(shape<0>(common))>::value &&
+                is_constant<1, decltype(stride<0>(common))>::value) {
+    // Truncate to the size of the contiguous vector (static stride-1 mode)
+    return shape<0>(common);
+  } else {
+    return Int<1>{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/* Return a layout that distributes ShapeB over ShapeA.
+ *
+ * @returns Layout result
+ * @post evenly_divides(@a b, size(@a result))
+ * @post evenly_divides(@a a, @a result)
+ * @post For all i,j in [0,size(@a result)) with i < j, @a result(i) < @a result(j). Surjective and Ordered.
+ * @post composition(make_layout(shape(@a a)), @a result) is admissible
+ * \code
+ *   // Note that 6 does not divide this shape
+ *   Layout layoutA = Layout<Shape<Int<15>,Int<14>>>{};
+ *
+ *   // Want to tile any 6 elements and don't care where they come from
+ *   Layout dist = domain_distribute(layoutA, Int<6>{});   // (_3,_2):(_1,_15)
+ *
+ *   // Not guaranteed to find all 6 though...
+ *   MUTE_STATIC_ASSERT_V(Int<6>{} == size(dist));
+ *
+ *   Layout result = zipped_divide(layoutA, dist);         // (_6,Rest)
+ * \endcode
+ */
+template <class ShapeA, class ShapeB>
+MUTE_HOST_DEVICE constexpr
+auto
+domain_distribute(ShapeA const& a, ShapeB const& b)
+{
+  static_assert(is_integral<ShapeB>::value);
+  static_assert(is_static<ShapeB>::value);
+
+  auto flat_shape_a = flatten(shape(a));
+
+  static_assert(is_static<decltype(flat_shape_a)>::value);
+
+  // Compute the shape of the result
+  auto [result_shape, b_rest] = mute::fold(flat_shape_a, mute::make_tuple(mute::tuple<>{}, size(b)), [](auto init, auto a_) {
+    auto [result, b_] = init;
+    auto gcd_ = gcd(a_, b_);
+    return mute::make_tuple(append(result, gcd_), b_ / gcd_);
+  });
+
+  // Compute the stride of the result
+  auto result_stride = compact_major<LayoutLeft>(flat_shape_a);
+
+  return coalesce(make_layout(result_shape, result_stride));
+}
+
+//
+// Kernel (Nullspace) of a Layout
+//
+
+namespace detail {
+
+template <int NextI, class Stride, int... Is>
+MUTE_HOST_DEVICE constexpr
+auto
+nullspace_seq(Stride const& stride, seq<Is...>)
+{
+  if constexpr (NextI == rank_v<Stride>) {
+    return seq<Is...>{};
+  } else
+  if constexpr (is_constant<0, decltype(get<NextI>(stride))>::value) {
+    return detail::nullspace_seq<NextI+1>(stride, seq<Is..., NextI>{});
+  } else {
+    return detail::nullspace_seq<NextI+1>(stride, seq<Is...>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+//
+// Build the nullspace of a layout
+// @result A layout @a result such that
+//    size(@a result) == size(@a layout) / size(filter(@a layout))
+//    @a layout(@a result(i)) == 0 for all i < size(@a result)
+//
+
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+nullspace(Layout<Shape,Stride> const& layout)
+{
+  auto flat_layout = flatten(layout);
+
+  auto iseq = detail::nullspace_seq<0>(flat_layout.stride(), seq<>{});
+
+  if constexpr (iseq.size() == 0) {
+    return Layout<_1,_0>{};     // Empty case, nothing found
+  } else {
+    // Generate the corresponding new strides and construct
+    auto rstride = compact_major<LayoutLeft>(flat_layout.shape());
+    return make_layout(unwrap(transform(iseq, [&](auto i) { return shape<i>(flat_layout); })),
+                       unwrap(transform(iseq, [&](auto i) { return get<i>(rstride); })));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Zip
+//
+
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+zip(Layout<Shape,Stride> const& layout)
+{
+  return make_layout(zip(layout.shape()),
+                     zip(layout.stride()));
+}
+
+template <class TShape, class TStride,
+          class UShape, class UStride>
+MUTE_HOST_DEVICE constexpr
+auto
+zip(Layout<TShape,TStride> const& layoutA,
+    Layout<UShape,UStride> const& layoutB)
+{
+  return make_layout(zip(layoutA.shape(),  layoutB.shape()),
+                     zip(layoutA.stride(), layoutB.stride()));
+}
+
+//
+// Tile unzip
+//   Logical product and logical divide (on layouts) produce rank-2 results by design.
+//   Follow the profile of @a tile and zip the rank-2 modes located at the terminals into
+//   their own mode.
+//
+
+template <class LShape, class LStride, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+tile_unzip(Layout<LShape,LStride> const& layout,
+           Tiler                  const& tiler)
+{
+  return make_layout(zip2_by(layout.shape(),  tiler),
+                     zip2_by(layout.stride(), tiler));
+}
+
+//
+// Logical divide
+//
+
+template <class LShape, class LStride,
+          class TShape, class TStride>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_divide(Layout<LShape,LStride> const& layout,
+               Layout<TShape,TStride> const& tiler)
+{
+  return composition(layout, make_layout(tiler, complement(tiler, shape(layout))));
+}
+
+template <class LShape, class LStride, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_divide(Layout<LShape,LStride> const& layout,
+               Tiler                  const& tiler)
+{
+  if constexpr (is_tuple<Tiler>::value) {
+    static_assert(tuple_size<Tiler>::value <= Layout<LShape,LStride>::rank, "logical_divide: Too many modes in tiler.");
+    return transform_layout(layout, tiler, [](auto const& l, auto const& t) { return logical_divide(l,t); });
+  } else if constexpr (is_underscore<Tiler>::value) {
+    return layout;
+  } else if constexpr (is_integral<Tiler>::value) {
+    return logical_divide(layout, make_layout(tiler));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Generalization of ceil_div for Layout lhs
+//   is effectively the "rest mode" of logical_divide.
+// Occurs in the calculation of gridDim, for example, for generalized tilers
+// Example:
+//   dim3 gridDim(size(ceil_div(problem_shape_M, cta_tiler_M)),
+//                size(ceil_div(problem_shape_N, cta_tiler_N)));
+// This does not consider compositional acceptance, so it may be the case that
+//   ceil_div produces a result while logical_divide (and friends) do not.
+template <class Target, class TShape, class TStride>
+MUTE_HOST_DEVICE constexpr
+auto
+ceil_div(Target                 const& target,
+         Layout<TShape,TStride> const& tiler)
+{
+  return shape(complement(tiler, shape(target)));
+}
+
+//
+// Convenience operator
+//   that produces layouts like ((BLK_A,BLK_B,...),(a,b,...,x,y))
+//   by gathering the tile modes and residuals into a rank-2 result.
+//
+
+template <class LShape, class LStride,
+          class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+zipped_divide(Layout<LShape,LStride> const& layout,
+              Tiler                  const& tiler)
+{
+  return tile_unzip(logical_divide(layout, tiler), tiler);
+}
+
+// Same as zipped_divide, but unpacks the second mode: ((BLK_A,BLK_B,...),a,b,...,x,y)
+template <class LShape, class LStride,
+          class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+tiled_divide(Layout<LShape,LStride> const& layout,
+             Tiler                  const& tiler)
+{
+  auto result = zipped_divide(layout, tiler);
+
+  auto R1 = rank<1>(result);
+  return result(_, repeat<R1>(_));
+}
+
+// Same as zipped_divide, but unpacks both modes: (BLK_A,BLK_B,...,a,b,...,x,y)
+template <class LShape, class LStride,
+          class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+flat_divide(Layout<LShape,LStride> const& layout,
+            Tiler                  const& tiler)
+{
+  auto result = zipped_divide(layout, tiler);
+
+  auto R0 = rank<0>(result);
+  auto R1 = rank<1>(result);
+  return result(repeat<R0>(_), repeat<R1>(_));
+}
+
+//
+// Logical product
+//
+
+template <class LShape, class LStride,
+          class TShape, class TStride>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_product(Layout<LShape,LStride> const& block,
+                Layout<TShape,TStride> const& tiler)
+{
+  return make_layout(block, composition(complement(block, size(block)*cosize(tiler)), tiler));
+}
+
+template <class LShape, class LStride, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_product(Layout<LShape,LStride> const& block,
+                Tiler                  const& tiler)
+{
+  if constexpr (is_tuple<Tiler>::value) {
+    static_assert(tuple_size<Tiler>::value <= Layout<LShape,LStride>::rank, "logical_product: Too many modes in tiler.");
+    return transform_layout(block, tiler, [](auto const& l, auto const& t) { return logical_product(l,t); });
+  } else if constexpr (is_underscore<Tiler>::value) {
+    return block;
+  } else if constexpr (is_integral<Tiler>::value) {
+    return logical_product(block, make_layout(tiler));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Convenience operator
+//   that produces layouts like ((BLK_A,BLK_B,...),(a,b,...,x,y))
+//   by gathering the block modes and products into a rank-2 result.
+//
+
+template <class LShape, class LStride,
+          class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+zipped_product(Layout<LShape,LStride> const& block,
+               Tiler                  const& tiler)
+{
+  return tile_unzip(logical_product(block, tiler), tiler);
+}
+
+// Same as zipped_product, but unpacks the second mode: ((BLK_A,BLK_B,...),a,b,...,x,y)
+template <class LShape, class LStride,
+          class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+tiled_product(Layout<LShape,LStride> const& block,
+              Tiler                  const& tiler)
+{
+  auto result = zipped_product(block, tiler);
+
+  auto R1 = rank<1>(result);
+  return result(_, repeat<R1>(_));
+}
+
+// Same as zipped_product, but unpacks both modes: (BLK_A,BLK_B,...,a,b,...,x,y)
+template <class LShape, class LStride,
+          class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+flat_product(Layout<LShape,LStride> const& block,
+             Tiler                  const& tiler)
+{
+  auto result = zipped_product(block, tiler);
+
+  auto R0 = rank<0>(result);
+  auto R1 = rank<1>(result);
+  return result(repeat<R0>(_), repeat<R1>(_));
+}
+
+//
+// Rank-sensitive products
+//
+
+// blocked_product -- Reproduce a block over a tiler.
+// Think of every element of "tiler" as a "block"
+//   and return the layout of the resulting structure.
+// @post rank(@a result) == mute::max(rank(@a block), rank(@a tiler))
+template <class TShape, class TStride,
+          class UShape, class UStride>
+MUTE_HOST_DEVICE constexpr
+auto
+blocked_product(Layout<TShape,TStride> const& block,
+                Layout<UShape,UStride> const& tiler)
+{
+  constexpr int R = mute::max(rank_v<TShape>, rank_v<UShape>);
+
+  auto result = logical_product(append<R>(block), append<R>(tiler));
+
+  return coalesce(zip(get<0>(result), get<1>(result)), tuple_repeat<R>(Int<1>{}));
+}
+
+// raked_product -- Reproduce a block over a tiler with block-interleaving.
+// Think of every element of "tiler" as a "block", interleave those blocks,
+//   and return the layout of the resulting structure.
+// @post rank(@a result) == mute::max(rank(@a block), rank(@a tiler))
+template <class TShape, class TStride,
+          class UShape, class UStride>
+MUTE_HOST_DEVICE constexpr
+auto
+raked_product(Layout<TShape,TStride> const& block,
+              Layout<UShape,UStride> const& tiler)
+{
+  constexpr int R = mute::max(rank_v<TShape>, rank_v<UShape>);
+
+  auto result = logical_product(append<R>(block), append<R>(tiler));
+
+  return coalesce(zip(get<1>(result), get<0>(result)), tuple_repeat<R>(Int<1>{}));
+}
+
+// tile_to_shape -- Perform a product of a layout so that the result matches a target shape.
+// This is similar to blocked_product, but specifies the result shape instead of the
+//   product shape, which is more convenient in certain circumstances.
+// @param block The layout to repeat
+// @param trg_shape The target shape of the result
+// @param ord_shape The order of the modes of @a trg_shape to tile @a layout with.
+//                  Defaults to GenColMajor, so @a layout will repeat
+//                    across the first mode first, the second mode second, etc
+//                  E.g. Step<_2,_1,_3> will cause @a layout to repeat
+//                    across the second mode first, the first mode second, and the third mode last.
+// @pre rank(@a block) <= rank(@a trg_shape)
+// @post compatible(@a trg_shape, shape(@a result))
+template <class Shape, class Stride,
+          class TrgShape, class ModeOrder = LayoutLeft>
+MUTE_HOST_DEVICE constexpr
+auto
+tile_to_shape(Layout<Shape,Stride> const& block,
+              TrgShape             const& trg_shape,
+              ModeOrder            const& ord_shape = {})
+{
+  MUTE_STATIC_ASSERT_V(rank(block) <= rank(trg_shape), "Rank of layout must be <= rank of target shape.");
+  constexpr int R = rank_v<TrgShape>;
+
+  auto padded_block = append<R>(block);
+
+  auto block_shape  = product_each(shape(padded_block));
+  auto target_shape = product_each(shape(trg_shape));
+
+  // Assert proper division
+  if constexpr (is_static<decltype(target_shape)>::value) {
+    MUTE_STATIC_ASSERT_V(evenly_divides(target_shape, block_shape),
+                         "tile_to_shape: block shape does not divide the target shape.");
+  }
+
+  auto product_shape = ceil_div(target_shape, block_shape);
+
+  return coalesce(blocked_product(padded_block, make_ordered_layout(product_shape, ord_shape)), product_shape);
+}
+
+//
+// Upcast
+//   For stride-1 mode, divide size by N. Divide all other strides by N.
+//
+
+template <int N, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+upcast(Shape const& shape, Stride const& stride)
+{
+  if constexpr (is_tuple<Shape>::value) {                  // tuple stride
+    return transform_layout(shape, stride, [](auto const& s, auto const& d) { return upcast<N>(s,d); });
+  } else if constexpr (is_constant<0, Stride>::value) {    // static-0 stride
+    return Layout<Shape,Stride>{shape,stride};
+  } else if constexpr (is_static<Stride>::value) {         // static stride
+    return make_layout(shape_div(shape,  shape_div(Int<N>{}, abs(stride))),
+                       shape_div(stride, Int<N>{}));
+  } else {                                                 // dynamic stride
+    // assume dynamic strides are larger than N and divisible
+    // assert(stride % N == 0);
+    return make_layout(shape, safe_div(stride, Int<N>{}));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <int N, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+upcast(Layout<Shape,Stride> const& layout)
+{
+  return upcast<N>(layout.shape(), layout.stride());
+}
+
+//
+// Downcast
+//   For stride-1 mode, multiply size by N. Multiply all other strides by N.
+//
+
+template <int N, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+downcast(Shape const& shape, Stride const& stride)
+{
+  if constexpr (is_tuple<Shape>::value) {
+    return transform_layout(shape, stride, [](auto const& s, auto const& d) { return downcast<N>(s,d); });
+  } else if constexpr (is_constant<1, Stride>::value || is_constant<-1, Stride>::value) {
+    return make_layout(shape * Int<N>{}, stride);
+  } else {
+    return make_layout(shape, stride * Int<N>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <int N, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+downcast(Layout<Shape,Stride> const& layout)
+{
+  MUTE_STATIC_ASSERT(has_int1<Stride>::value, "Downcast requires adjacent elements");
+  return downcast<N>(layout.shape(), layout.stride());
+}
+
+//
+// Recast
+//
+
+template <class OldType, class NewType,
+          class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_layout(Layout<Shape,Stride> const& layout)
+{
+  using scale = decltype(trait_ratio(sizeof_bits<NewType>{}, sizeof_bits<OldType>{}));
+  if constexpr (scale::num == 1 && scale::den == 1) {
+    return layout;
+  }
+  else if constexpr (scale::num == 1) {
+    return downcast<scale::den>(layout);
+  }
+  else if constexpr (scale::den == 1) {
+    return upcast<scale::num>(layout);
+  }
+  else {
+    return downcast<scale::den>(upcast<scale::num>(layout));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Determine the maximum alignment of a Layout.
+// The maximum alignment is the largest N for which upcast<N>(layout) will compile.
+//   upcast<N>(layout) compiles when the static shapes and strides pass divisibility checks.
+//   Therefore, upcast<M>(layout) will also compile for all divisors M of N.
+// Note that this only considers the static shapes and strides of the Layout
+//   in symmetry with upcast<N> only checking against static shapes and strides and assuming all
+//   dynamic shapes and strides are large and multiples of N.
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+max_alignment(Layout<Shape,Stride> const& layout)
+{
+  auto flat_layout   = coalesce(layout);
+  auto static_shape  = transform( shape(flat_layout), [](auto s){ return conditional_return<is_static<decltype(s)>::value>(s, Int<1>{}); });
+  auto static_stride = transform(stride(flat_layout), [](auto d){ return conditional_return<is_static<decltype(d)>::value>(d, Int<0>{}); });
+  auto filter_layout = make_layout(static_shape, static_stride);
+  auto permuted = logical_divide(filter_layout, right_inverse(filter_layout));
+  return gcd(size<0>(permuted), stride<1>(permuted));
+}
+
+//
+// Display utilities
+//
+
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE void print(Layout<Shape,Stride> const& layout)
+{
+  print(layout.shape()); print(":"); print(layout.stride());
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class Shape, class Stride>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, Layout<Shape,Stride> const& layout)
+{
+  return os << shape(layout) << ":" << stride(layout);
+}
+#endif
+
+// Generic 2D Layout to console table
+template <class Layout>
+MUTE_HOST_DEVICE
+void
+print_layout(Layout const& layout)  // (m,n) -> idx
+{
+  MUTE_STATIC_ASSERT_V(rank(layout) == Int<2>{});
+
+  int idx_width = num_digits(cosize(layout)) + 2;
+  const char* delim = "+-----------------------";
+
+  print(layout); print("\n");
+
+  // Column indices
+  print("    ");
+  for (int n = 0; n < size<1>(layout); ++n) { printf("  %*d ", idx_width-2, n); }
+  printf("\n");
+
+  // Print out A m-by-n
+  for (int m = 0; m < size<0>(layout); ++m) {
+    // Header
+    print("    ");
+    for (int n = 0; n < size<1>(layout); ++n) { printf("%.*s", idx_width+1, delim); }
+    printf("+\n");
+    // Values
+    printf("%2d  ", m);  // Row indices
+    for (int n = 0; n < size<1>(layout); ++n) { printf("| %*d ", idx_width-2, int(layout(m,n))); }
+    printf("|\n");
+  }
+  // Footer
+  print("    ");
+  for (int n = 0; n < size<1>(layout); ++n) { printf("%.*s", idx_width+1, delim); }
+  printf("+\n");
+}
+
+// Generic ThrVal 2D Layout to console table
+template <class Layout, class ThrID>
+MUTE_HOST_DEVICE
+void
+print_layout(Layout const& layout, ThrID const& thrid)  // (m,n) -> (tid,vid)  and  tid -> thr_idx
+{
+  MUTE_STATIC_ASSERT_V(rank(layout) == Int<2>{});
+
+  print(layout); print("\n");
+  print(thrid);  print("\n");
+
+  // Print out m-by-n
+  for (int m = 0; m < size<0>(layout); ++m) {
+    // Header
+    for (int n = 0; n < size<1>(layout); ++n) printf("+------");
+    printf("+\n");
+    // Values
+    for (int n = 0; n < size<1>(layout); ++n) printf("|%03d-%02d", int(thrid(layout(m,n) % size(thrid))), int(layout(m,n) / size(thrid)));
+    printf("|\n");
+  }
+  // Footer
+  for (int n = 0; n < size<1>(layout); ++n) printf("+------");
+  printf("+\n");
+}
+
+struct TikzColor_White {
+  MUTE_HOST_DEVICE char const*
+  operator()(int idx) const {
+    return "white";
+  }
+};
+
+struct TikzColor_BWx8 {
+  MUTE_HOST_DEVICE char const*
+  operator()(int idx) const {
+    static char const* color_map[8] = {"black!00", "black!40", "black!20", "black!60",
+                                       "black!10", "black!50", "black!30", "black!70"};
+    return color_map[idx % 8];
+  }
+};
+
+struct TikzColor_TV {
+  MUTE_HOST_DEVICE char const*
+  operator()(int tid, int vid) const {
+    static char const* color_map[8] = {"{rgb,255:red,175;green,175;blue,255}",
+                                       "{rgb,255:red,175;green,255;blue,175}",
+                                       "{rgb,255:red,255;green,255;blue,175}",
+                                       "{rgb,255:red,255;green,175;blue,175}",
+                                       "{rgb,255:red,210;green,210;blue,255}",
+                                       "{rgb,255:red,210;green,255;blue,210}",
+                                       "{rgb,255:red,255;green,255;blue,210}",
+                                       "{rgb,255:red,255;green,210;blue,210}"};
+    return color_map[tid % 8];
+  }
+};
+
+// Generic 2D Layout to LaTeX printer
+template <class LayoutA, class TikzColorFn = TikzColor_BWx8>
+MUTE_HOST_DEVICE
+void
+print_latex(LayoutA const& layout_a,   // (m,n) -> idx
+            TikzColorFn color = {})    // lambda(idx) -> tikz color string
+{
+  MUTE_STATIC_ASSERT_V(rank(layout_a) <= Int<2>{});
+  auto layout = append<2>(layout_a, Layout<_1,_0>{});
+
+  // Commented print(layout)
+  printf("%% Layout: "); print(layout); printf("\n");
+  // Header
+  printf("\\documentclass[convert]{standalone}\n"
+         "\\usepackage{tikz}\n\n"
+         "\\begin{document}\n"
+         "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
+
+  // Layout
+  for (int i = 0; i < size<0>(layout); ++i) {
+    for (int j = 0; j < size<1>(layout); ++j) {
+      int idx = layout(i,j);
+      printf("\\node[fill=%s] at (%d,%d) {%d};\n",
+             color(idx), i, j, idx);
+    }
+  }
+  // Grid
+  printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n",
+         int(size<0>(layout)), int(size<1>(layout)));
+  // Labels
+  for (int i =  0, j = -1; i < size<0>(layout); ++i) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
+  }
+  for (int i = -1, j =  0; j < size<1>(layout); ++j) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j);
+  }
+
+  // Footer
+  printf("\\end{tikzpicture}\n"
+         "\\end{document}\n");
+}
+
+// Generic ThrVal 2D Layout to LaTeX TikZ
+template <class Layout, class ThrID, class TikzColorFn = TikzColor_TV>
+MUTE_HOST_DEVICE
+void
+print_latex(Layout const& layout,    // (m,n) -> (tid,vid)
+            ThrID  const& thr,       // tid -> thr_idx
+            TikzColorFn color = {})  // lambda(thr_idx,val_idx) -> tikz color string
+{
+  MUTE_STATIC_ASSERT_V(rank(layout) == Int<2>{});
+
+  // Commented prints
+  printf("%% Layout: "); print(layout); printf("\n");
+  printf("%% ThrID : "); print(thr);  printf("\n");
+  // Header
+  printf("\\documentclass[convert]{standalone}\n"
+         "\\usepackage{tikz}\n\n"
+         "\\begin{document}\n"
+         "\\begin{tikzpicture}[x={(0cm,-1cm)},y={(1cm,0cm)},every node/.style={minimum size=1cm, outer sep=0pt}]\n\n");
+
+  // Layout
+  for (int i = 0; i < size<0>(layout); ++i) {
+    for (int j = 0; j < size<1>(layout); ++j) {
+      int thrid   = layout(i,j) % size(thr);
+      int val_idx = layout(i,j) / size(thr);
+      int thr_idx = thr(thrid);
+
+      printf("\\node[fill=%s] at (%d,%d) {\\shortstack{T%d \\\\ V%d}};\n",
+             color(thr_idx, val_idx),
+             i, j,
+             thr_idx, val_idx);
+    }
+  }
+  // Grid
+  printf("\\draw[color=black,thick,shift={(-0.5,-0.5)}] (0,0) grid (%d,%d);\n\n",
+         int(size<0>(layout)), int(size<1>(layout)));
+  // Labels
+  for (int i = 0, j = -1; i < size<0>(layout); ++i) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, i);
+  }
+  for (int j = 0, i = -1; j < size<1>(layout); ++j) {
+    printf("\\node at (%d,%d) {\\Large{\\texttt{%d}}};\n", i, j, j);
+  }
+
+  // Footer
+  printf("\\end{tikzpicture}\n"
+         "\\end{document}\n");
+}
+
+} // end namespace mute
+
+//
+// Extended Layouts
+//
+
+#include <mute/swizzle_layout.hpp>
diff --git a/vllm_musa/csrc/flash_attn/include/mute/layout_composed.hpp b/vllm_musa/csrc/flash_attn/include/mute/layout_composed.hpp
new file mode 100644
index 000000000..3e9fa26c5
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/layout_composed.hpp
@@ -0,0 +1,660 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE, MUTE_GCC_UNREACHABLE
+#include <mute/layout.hpp>                     // mute::tuple
+#include <mute/numeric/integral_constant.hpp>  // mute::true_type, mute::false_type, mute::Int
+
+/* This implements a ComposedLayout of the form
+ *   LayoutA o Offset o LayoutB
+ * and is useful in cases where composition() does not or cannot apply to LayoutA and LayoutB.
+ * For example, when the "divisibility condition" in shape_div is violated in composition(LayoutA, LayoutB).
+ *
+ * This ComposedLayout provides similar functionality to Layout including tiling, partitioning,
+ * coordinate-to-index mapping and layout manipulations, but is not considered a "normal" layout.
+ * For example, this layout provides shape() and size() functions, but does not provide stride() functions.
+ * Mostly, the similar functionality is accomplished by applying each operation to LayoutB only
+ * as LayoutB defines the domain.
+ */
+
+namespace mute
+{
+
+// A Layout of non-trivially composable functions: F o I o L
+template <class LayoutA, class Offset, class LayoutB>
+struct ComposedLayout : private mute::tuple<LayoutA, Offset, LayoutB>  // EBO for static layouts
+{
+  MUTE_HOST_DEVICE constexpr
+  ComposedLayout(LayoutA const& layoutA = {},
+                 Offset  const& offset  = {},
+                 LayoutB const& layoutB = {})
+      : mute::tuple<LayoutA, Offset, LayoutB>(layoutA, offset, layoutB)
+  {}
+
+  //
+  // Accessors
+  //
+
+  static constexpr int rank  = LayoutB::rank;
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  layout_a() const {
+    return get<0>(static_cast<mute::tuple<LayoutA, Offset, LayoutB> const&>(*this));
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  offset() const {
+    return get<1>(static_cast<mute::tuple<LayoutA, Offset, LayoutB> const&>(*this));
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  layout_b() const {
+    return get<2>(static_cast<mute::tuple<LayoutA, Offset, LayoutB> const&>(*this));
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  layout() const {
+    return *this;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  shape() const {
+    return layout_b().shape();
+  }
+
+  // Doesn't really make sense to ask for the strides of this "layout"
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  stride() const = delete;
+
+  //
+  // Mappings
+  //
+
+  // Map a logical coordinate to a linear index (Coord has no Underscore slice operators)
+  // OR
+  // Slice the layout and return the sublayout (Coord has an Underscore slice op)
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator()(Coord const& coord) const {
+    if constexpr (has_underscore<Coord>::value) {
+      return slice(coord, *this);
+    } else {
+      return layout_a()(offset() + layout_b()(coord));    // (A o O o B)(c)
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+
+  // Convenience function for multi-dimensional coordinates
+  template <class Coord0, class Coord1, class... Coords>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const {
+    return operator()(make_coord(c0,c1,cs...));
+  }
+
+  //
+  // Compose
+  //
+
+  template <class OtherLayout>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  compose(OtherLayout const& other) const {
+    return composition(*this, other);
+  }
+
+  template <class... Layouts>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  compose(Layouts const&... layouts) const {
+    return composition(*this, make_tile(layouts...));
+  }
+
+  template <class OtherShape>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  with_shape(OtherShape const& shape) const {
+    return composition(*this, make_layout(shape));
+  }
+
+  template <class... Shapes>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  with_shape(Shapes const&... shapes) const {
+    return composition(*this, make_layout(make_shape(shapes...)));
+  }
+
+  //
+  // Tile
+  //
+
+  template <class OtherLayout>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  tile(OtherLayout const& other) const {
+    return tiled_divide(*this, other);
+  }
+
+  template <class... Layouts>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  tile(Layouts const&... layouts) const {
+    return tiled_divide(*this, make_tile(layouts...));
+  }
+
+  // Equality, return a static or dynamic boolean
+  template <class... Args>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator==(ComposedLayout<Args...> const& other) const {
+    return this->layout_a() == other.layout_a() &&
+           this->layout_b() == other.layout_b() &&
+           this->offset()   == other.offset();
+  }
+};
+
+template <class A, class O, class B>
+struct is_layout<ComposedLayout<A,O,B>> : true_type {};
+
+template <class T>
+struct is_composed_layout : false_type {};
+template <class A, class O, class B>
+struct is_composed_layout<ComposedLayout<A,O,B>> : true_type {};
+
+//
+// Constructors
+//
+
+template <class LayoutA, class Offset, class LayoutB>
+MUTE_HOST_DEVICE constexpr
+auto
+make_composed_layout(LayoutA const& layoutA,
+                     Offset  const& offset,
+                     LayoutB const& layoutB)
+{
+  return ComposedLayout<LayoutA, Offset, LayoutB>{layoutA, offset, layoutB};
+}
+
+//
+// Utilities
+//
+
+// Return the layout of a mode
+template <int... Is, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+layout(ComposedLayout<A,O,B> const& clayout)
+{
+  return composition(clayout.layout_a(), clayout.offset(), layout<Is...>(clayout.layout_b()));
+}
+
+// Return the shape of a mode
+template <int... Is, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+shape(ComposedLayout<A,O,B> const& layout)
+{
+  return shape<Is...>(layout.layout_b());
+}
+
+// Doesn't make sense to directly ask for the strides of this "layout"
+template <int... Is, class Fn, class O, class Layout>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+stride(ComposedLayout<Fn,O,Layout> const& layout) = delete;
+
+// Return the number of elements in a mode
+template <int... Is, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+size(ComposedLayout<A,O,B> const& layout)
+{
+  return size<Is...>(layout.layout_b());
+}
+
+// Return the number of modes
+template <int... Is, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+rank(ComposedLayout<A,O,B> const& layout)
+{
+  return rank<Is...>(layout.layout_b());
+}
+
+// Return the depth of the layout
+template <int... Is, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+depth(ComposedLayout<A,O,B> const& layout)
+{
+  return depth<Is...>(layout.layout_b());
+}
+
+// Return the codomain size of a mode
+template <int... Is, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+cosize(ComposedLayout<A,O,B> const& layout)
+{
+  return cosize<Is...>(layout.layout_b());
+}
+
+//
+// Operations to manipulate Layouts like a tuple of pairs
+//
+
+template <size_t I, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+get(ComposedLayout<A,O,B> const& a)
+{
+  return composition(a.layout_a(), a.offset(), get<I>(a.layout_b()));
+}
+
+template <int Begin, int End, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+take(ComposedLayout<A,O,B> const& a)
+{
+  return composition(a.layout_a(), a.offset(), take<Begin,End>(a.layout_b()));
+}
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+flatten(ComposedLayout<A,O,B> const& a)
+{
+  return composition(a.layout_a(), a.offset(), flatten(a.layout_b()));
+}
+
+template <int N, class A, class O, class B, class X>
+MUTE_HOST_DEVICE constexpr
+auto
+append(ComposedLayout<A,O,B> const& a, X const& x)
+{
+  return composition(a.layout_a(), a.offset(), append<N>(a.layout_b(), x));
+}
+
+template <int Begin, int End, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+group(ComposedLayout<A,O,B> const& a)
+{
+  return composition(a.layout_a(), a.offset(), group<Begin,End>(a.layout_b()));
+}
+
+//
+// Slice a ComposedLayout
+//
+
+template <class Coord, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+slice_and_offset(Coord const& coord, ComposedLayout<A,O,B> const& layout)
+{
+  auto [slice, offset] = slice_and_offset(coord, layout.layout_b());
+  return mute::make_tuple(ComposedLayout{layout.layout_a(), layout.offset() + offset, slice}, Int<0>{});
+}
+
+template <class Coord, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+slice(Coord const& coord, ComposedLayout<A,O,B> const& layout)
+{
+  return get<0>(slice_and_offset(coord, layout));
+}
+
+// Compute a pointer offset and (potentially modified) layout from a coordinate
+// For composed layout tensors the offset is accumulated in the layout itself while pointer is not updated
+template <class Coord, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+domain_offset(Coord const& coord, ComposedLayout<A,O,B> const& layout)
+{
+  return mute::make_tuple(ComposedLayout{layout.layout_a(), layout.offset() + layout.layout_b()(coord), layout.layout_b()}, Int<0>{});
+}
+
+//
+// composition
+//
+
+template <class LayoutA,
+          class Offset,
+          class LayoutB>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(LayoutA const& layoutA,
+            Offset  const& offset,
+            LayoutB const& layoutB)
+{
+  return ComposedLayout<LayoutA, Offset, LayoutB>{layoutA, offset, layoutB};
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(ComposedLayout<A,O,B> const& a,
+            Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), composition(a.layout_b(), b));
+}
+
+template <class ShapeA, class StrideA,
+          class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(Layout<ShapeA,StrideA> const& a,
+            ComposedLayout<A,O,B>  const& b)
+{
+  MUTE_STATIC_ASSERT_V(b.offset() == Int<0>{}, "Require offset == 0.");
+
+  return composition(composition(a, b.layout_a()), b.layout_b());
+}
+
+//
+// complement
+//
+
+template <class A, class O, class B, class CoTarget>
+MUTE_HOST_DEVICE constexpr
+auto
+complement(ComposedLayout<A,O,B> const& layout, CoTarget const& cotarget)
+{
+  return complement(layout.layout_b(), cotarget);
+}
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+complement(ComposedLayout<A,O,B> const& layout)
+{
+  return complement(layout, cosize(layout));
+}
+
+//
+// inverse
+//
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+right_inverse(ComposedLayout<A,O,B> const& layout)
+{
+  return composition(right_inverse(layout.layout_b()), right_inverse(layout.offset()), right_inverse(layout.layout_a()));
+}
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+left_inverse(ComposedLayout<A,O,B> const& layout)
+{
+  return composition(left_inverse(layout.layout_b()), left_inverse(layout.offset()), left_inverse(layout.layout_a()));
+}
+
+//
+// Other operations
+//
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+zip(ComposedLayout<A,O,B> const& a)
+{
+  return composition(a.layout_a(), a.offset(), zip(a.layout_b()));
+}
+
+// Partitions
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_divide(ComposedLayout<A,O,B> const& a,
+               Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), logical_divide(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+tile_unzip(ComposedLayout<A,O,B> const& a,
+           Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), tile_unzip(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+tiled_divide(ComposedLayout<A,O,B> const& a,
+             Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), tiled_divide(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+zipped_divide(ComposedLayout<A,O,B> const& a,
+              Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), zipped_divide(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+flat_divide(ComposedLayout<A,O,B> const& a,
+            Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), flat_divide(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_product(ComposedLayout<A,O,B> const& a,
+                Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), logical_product(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+zipped_product(ComposedLayout<A,O,B> const& a,
+               Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), zipped_product(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+tiled_product(ComposedLayout<A,O,B> const& a,
+              Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), tiled_product(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+flat_product(ComposedLayout<A,O,B> const& a,
+             Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), flat_product(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+blocked_product(ComposedLayout<A,O,B> const& a,
+                Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), blocked_product(a.layout_b(), b));
+}
+
+template <class A, class O, class B, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+raked_product(ComposedLayout<A,O,B> const& a,
+              Tiler                 const& b)
+{
+  return composition(a.layout_a(), a.offset(), raked_product(a.layout_b(), b));
+}
+
+template <class A, class O, class B,
+          class Shape, class ModeOrder = GenColMajor>
+MUTE_HOST_DEVICE constexpr
+auto
+tile_to_shape(ComposedLayout<A,O,B> const& layout,
+              Shape                 const& trg_shape,
+              ModeOrder             const& ord_shape = {})
+{
+  return composition(layout.layout_a(), layout.offset(), tile_to_shape(layout.layout_b(), trg_shape, ord_shape));
+}
+
+template <class A, class O, class B,
+          class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+filter(ComposedLayout<A,O,B> const& layout, Shape const& trg_profile)
+{
+  return composition(layout.layout_a(), layout.offset(), filter(layout.layout_b(), trg_profile));
+}
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce(ComposedLayout<A,O,B> const& layout)
+{
+  return composition(layout.layout_a(), layout.offset(), coalesce(layout.layout_b()));
+}
+
+template <class A, class O, class B, class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce(ComposedLayout<A,O,B> const& layout, Shape const& trg_profile)
+{
+  return composition(layout.layout_a(), layout.offset(), coalesce(layout.layout_b(), trg_profile));
+}
+
+//
+// Upcast and Downcast
+//
+
+template <int N, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+upcast(ComposedLayout<A,O,B> const& layout)
+{
+  return composition(upcast<N>(layout.layout_a()), upcast<N>(layout.offset()), upcast<N>(layout.layout_b()));
+}
+
+template <int N, class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+downcast(ComposedLayout<A,O,B> const& layout)
+{
+  return composition(downcast<N>(layout.layout_a()), downcast<N>(layout.offset()), downcast<N>(layout.layout_b()));
+}
+
+template <class OldType, class NewType,
+          class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_layout(ComposedLayout<A,O,B> const& layout)
+{
+  using scale = decltype(trait_ratio(sizeof_bits<NewType>{}, sizeof_bits<OldType>{}));
+  if constexpr (scale::num == 1 && scale::den == 1) {
+    return layout;
+  }
+  else if constexpr (scale::num == 1) {
+    return downcast<scale::den>(layout);
+  }
+  else if constexpr (scale::den == 1) {
+    return upcast<scale::num>(layout);
+  }
+  else {
+    return downcast<scale::den>(upcast<scale::num>(layout));
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+max_alignment(ComposedLayout<A,O,B> const& layout)
+{
+  // Do not attempt for general ComposedLayouts
+  //return gcd(max_alignment(layout.layout_a()), max_alignment(layout.offset()), max_alignment(layout.layout_b()));
+  return Int<1>{};
+}
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+nullspace(ComposedLayout<A,O,B> const& layout)
+{
+  // Do not attempt for general ComposedLayouts
+  return Layout<_1,_0>{};
+}
+
+//
+// Display utilities
+//
+
+template <class A, class O, class B>
+MUTE_HOST_DEVICE void print(ComposedLayout<A,O,B> const& layout)
+{
+  print(layout.layout_a()); print(" o "); print(layout.offset()); print(" o "); print(layout.layout_b());
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class A, class O, class B>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, ComposedLayout<A,O,B> const& layout)
+{
+  return os << layout.layout_a() << " o " << layout.offset() << " o " << layout.layout_b();
+}
+#endif
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/arithmetic_tuple.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/arithmetic_tuple.hpp
new file mode 100644
index 000000000..a1a9d69c3
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/arithmetic_tuple.hpp
@@ -0,0 +1,557 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/container/tuple.hpp>
+#include <mute/numeric/integral_constant.hpp>
+#include <mute/algorithm/functional.hpp>
+#include <mute/algorithm/tuple_algorithms.hpp>
+#include <mute/util/type_traits.hpp>
+
+namespace mute
+{
+
+template <class... T>
+struct ArithmeticTuple : tuple<T...>
+{
+  template <class... U>
+  MUTE_HOST_DEVICE constexpr
+  ArithmeticTuple(ArithmeticTuple<U...> const& u)
+    : tuple<T...>(static_cast<tuple<U...> const&>(u)) {}
+
+  template <class... U>
+  MUTE_HOST_DEVICE constexpr
+  ArithmeticTuple(tuple<U...> const& u)
+    : tuple<T...>(u) {}
+
+  template <class... U>
+  MUTE_HOST_DEVICE constexpr
+  ArithmeticTuple(U const&... u)
+    : tuple<T...>(u...) {}
+};
+
+template <class... T>
+struct is_tuple<ArithmeticTuple<T...>> : true_type {};
+
+template <class... Ts>
+struct is_flat<ArithmeticTuple<Ts...>> : is_flat<tuple<Ts...>> {};
+
+template <class... T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_arithmetic_tuple(T const&... t) {
+  return ArithmeticTuple<T...>(t...);
+}
+
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+as_arithmetic_tuple(T const& t) {
+  if constexpr (is_tuple<T>::value) {
+    return detail::tapply(t, [](auto const& x){ return as_arithmetic_tuple(x); },
+                          [](auto const&... a){ return make_arithmetic_tuple(a...); },
+                          tuple_seq<T>{});
+  } else {
+    return t;
+  }
+}
+
+//
+// Numeric operators
+//
+
+// Addition
+template <class... T, class... U>
+MUTE_HOST_DEVICE constexpr
+auto
+operator+(ArithmeticTuple<T...> const& t, ArithmeticTuple<U...> const& u) {
+  constexpr int R = mute::max(int(sizeof...(T)), int(sizeof...(U)));
+  return transform_apply(append<R>(t,Int<0>{}), append<R>(u,Int<0>{}), plus{}, [](auto const&... a){ return make_arithmetic_tuple(a...); });
+}
+
+template <class... T, class... U>
+MUTE_HOST_DEVICE constexpr
+auto
+operator+(ArithmeticTuple<T...> const& t, tuple<U...> const& u) {
+  return t + ArithmeticTuple<U...>(u);
+}
+
+template <class... T, class... U>
+MUTE_HOST_DEVICE constexpr
+auto
+operator+(tuple<T...> const& t, ArithmeticTuple<U...> const& u) {
+  return ArithmeticTuple<T...>(t) + u;
+}
+
+// Subtraction
+template <class... T, class... U>
+MUTE_HOST_DEVICE constexpr
+auto
+operator-(ArithmeticTuple<T...> const& t, ArithmeticTuple<U...> const& u) {
+  constexpr int R = mute::max(int(sizeof...(T)), int(sizeof...(U)));
+  return transform_apply(append<R>(t,Int<0>{}), append<R>(u,Int<0>{}), minus{}, [](auto const&... a){ return make_arithmetic_tuple(a...); });
+}
+
+template <class... T, class... U>
+MUTE_HOST_DEVICE constexpr
+auto
+operator-(ArithmeticTuple<T...> const& t, tuple<U...> const& u) {
+  return t - ArithmeticTuple<U...>(u);
+}
+
+template <class... T, class... U>
+MUTE_HOST_DEVICE constexpr
+auto
+operator-(tuple<T...> const& t, ArithmeticTuple<U...> const& u) {
+  return ArithmeticTuple<T...>(t) - u;
+}
+
+// Negation
+template <class... T>
+MUTE_HOST_DEVICE constexpr
+auto
+operator-(ArithmeticTuple<T...> const& t) {
+  return transform_apply(t, negate{}, [](auto const&... a){ return make_arithmetic_tuple(a...); });
+}
+
+//
+// Special cases
+//
+
+template <auto t, class... U>
+MUTE_HOST_DEVICE constexpr
+ArithmeticTuple<U...> const&
+operator+(C<t>, ArithmeticTuple<U...> const& u) {
+  static_assert(t == 0, "Arithmetic tuple op+ error!");
+  return u;
+}
+
+template <class... T, auto u>
+MUTE_HOST_DEVICE constexpr
+ArithmeticTuple<T...> const&
+operator+(ArithmeticTuple<T...> const& t, C<u>) {
+  static_assert(u == 0, "Arithmetic tuple op+ error!");
+  return t;
+}
+
+template <auto t, class... U>
+MUTE_HOST_DEVICE constexpr
+ArithmeticTuple<U...> const&
+operator-(C<t>, ArithmeticTuple<U...> const& u) {
+  static_assert(t == 0, "Arithmetic tuple op- error!");
+  return -u;
+}
+
+template <class... T, auto u>
+MUTE_HOST_DEVICE constexpr
+ArithmeticTuple<T...> const&
+operator-(ArithmeticTuple<T...> const& t, C<u>) {
+  static_assert(u == 0, "Arithmetic tuple op- error!");
+  return t;
+}
+
+//
+// ArithmeticTupleIterator
+//
+
+template <class ArithTuple>
+struct ArithmeticTupleIterator
+{
+  using value_type   = ArithTuple;
+  using element_type = ArithTuple;
+  using reference    = ArithTuple;
+
+  ArithTuple coord_;
+
+  MUTE_HOST_DEVICE constexpr
+  ArithmeticTupleIterator(ArithTuple const& coord = {}) : coord_(coord) {}
+
+  MUTE_HOST_DEVICE constexpr
+  ArithTuple operator*() const { return coord_; }
+
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  auto operator[](Coord const& c) const { return *(*this + c); }
+
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  auto operator+(Coord const& c) const {
+    return ArithmeticTupleIterator<remove_cvref_t<decltype(coord_ + c)>>(coord_ + c);
+  }
+};
+
+template <class Tuple>
+MUTE_HOST_DEVICE constexpr
+auto
+make_inttuple_iter(Tuple const& t) {
+  return ArithmeticTupleIterator(as_arithmetic_tuple(t));
+}
+
+template <class T0, class T1, class... Ts>
+MUTE_HOST_DEVICE constexpr
+auto
+make_inttuple_iter(T0 const& t0, T1 const& t1, Ts const&... ts) {
+  return make_inttuple_iter(mute::make_tuple(t0, t1, ts...));
+}
+
+//
+// ArithmeticTuple "basis" elements
+//   A ScaledBasis<T,N> is a (at least) rank-N+1 ArithmeticTuple:
+//      (_0,_0,...,T,_0,...)
+//   with value T in the Nth mode
+
+template <class T, int N>
+struct ScaledBasis : private tuple<T>
+{
+  MUTE_HOST_DEVICE constexpr
+  ScaledBasis(T const& t = {}) : tuple<T>(t) {}
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto) value()       { return get<0>(static_cast<tuple<T>      &>(*this)); }
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto) value() const { return get<0>(static_cast<tuple<T> const&>(*this)); }
+
+  MUTE_HOST_DEVICE static constexpr
+  auto mode() { return Int<N>{}; }
+};
+
+template <class T>
+struct is_scaled_basis : false_type {};
+template <class T, int N>
+struct is_scaled_basis<ScaledBasis<T,N>> : true_type {};
+
+template <class T, int N>
+struct is_integral<ScaledBasis<T,N>> : true_type {};
+
+// Get the scalar T out of a ScaledBasis
+template <class SB>
+MUTE_HOST_DEVICE constexpr auto
+basis_value(SB const& e)
+{
+  if constexpr (is_scaled_basis<SB>::value) {
+    return basis_value(e.value());
+  } else {
+    return e;
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Apply the N... pack to another Tuple
+template <class SB, class Tuple>
+MUTE_HOST_DEVICE decltype(auto)
+basis_get(SB const& e, Tuple&& t)
+{
+  if constexpr (is_scaled_basis<SB>::value) {
+    return basis_get(e.value(), get<SB::mode()>(static_cast<Tuple&&>(t)));
+  } else {
+    return static_cast<Tuple&&>(t);
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+namespace detail {
+
+template <class T, int... I>
+MUTE_HOST_DEVICE constexpr
+auto
+to_atuple_i(T const& t, seq<I...>) {
+  return make_arithmetic_tuple((void(I),Int<0>{})..., t);
+}
+
+} // end namespace detail
+
+// Turn a ScaledBases<T,N> into a rank-N+1 ArithmeticTuple
+//    with N prefix 0s:  (_0,_0,...N...,_0,T)
+template <class T, int N>
+MUTE_HOST_DEVICE constexpr
+auto
+as_arithmetic_tuple(ScaledBasis<T,N> const& t) {
+  return detail::to_atuple_i(as_arithmetic_tuple(t.value()), make_seq<N>{});
+}
+
+namespace detail {
+
+template <int... Ns>
+struct Basis;
+
+template <>
+struct Basis<> {
+  using type = Int<1>;
+};
+
+template <int N, int... Ns>
+struct Basis<N,Ns...> {
+  using type = ScaledBasis<typename Basis<Ns...>::type, N>;
+};
+
+} // end namespace detail
+
+// Shortcut for writing ScaledBasis<ScaledBasis<ScaledBasis<Int<1>, N0>, N1>, ...>
+// E<>    := _1
+// E<0>   := (_1,_0,_0,...)
+// E<1>   := (_0,_1,_0,...)
+// E<0,0> := ((_1,_0,_0,...),_0,_0,...)
+// E<0,1> := ((_0,_1,_0,...),_0,_0,...)
+// E<1,0> := (_0,(_1,_0,_0,...),_0,...)
+// E<1,1> := (_0,(_0,_1,_0,...),_0,...)
+template <int... N>
+using E = typename detail::Basis<N...>::type;
+
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+make_basis_like(Shape const& shape)
+{
+  if constexpr (is_integral<Shape>::value) {
+    return Int<1>{};
+  } else {
+    // Generate bases for each rank of shape
+    return transform(tuple_seq<Shape>{}, shape, [](auto I, auto si) {
+      // Generate bases for each rank of si and add an i on front
+      using I_type = decltype(I);
+      return transform_leaf(make_basis_like(si), [](auto e) {
+        // MSVC has trouble capturing variables as constexpr,
+        // so that they can be used as template arguments.
+        // This is exactly what the code needs to do with i, unfortunately.
+        // The work-around is to define i inside the inner lambda,
+        // by using just the type from the enclosing scope.
+        constexpr int i = I_type::value;
+        return ScaledBasis<decltype(e), i>{};
+      });
+    });
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Arithmetic
+//
+
+template <class T, int M, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+safe_div(ScaledBasis<T,M> const& b, U const& u)
+{
+  auto t = safe_div(b.value(), u);
+  return ScaledBasis<decltype(t),M>{t};
+}
+
+template <class T, int M, class U>
+MUTE_HOST_DEVICE constexpr
+auto
+shape_div(ScaledBasis<T,M> const& b, U const& u)
+{
+  auto t = shape_div(b.value(), u);
+  return ScaledBasis<decltype(t),M>{t};
+}
+
+// Equality
+template <class T, int N, class U, int M>
+MUTE_HOST_DEVICE constexpr
+auto
+operator==(ScaledBasis<T,N> const& t, ScaledBasis<U,M> const& u) {
+  return bool_constant<M == N>{} && t.value() == u.value();
+}
+
+// Not equal to anything else
+template <class T, int N, class U>
+MUTE_HOST_DEVICE constexpr
+false_type
+operator==(ScaledBasis<T,N> const&, U const&) {
+  return {};
+}
+
+template <class T, class U, int M>
+MUTE_HOST_DEVICE constexpr
+false_type
+operator==(T const&, ScaledBasis<U,M> const&) {
+  return {};
+}
+
+// Abs
+template <class T, int N>
+MUTE_HOST_DEVICE constexpr
+auto
+abs(ScaledBasis<T,N> const& e) {
+  return ScaledBasis<decltype(abs(e.value())),N>{abs(e.value())};
+}
+
+// Multiplication
+template <class A, class T, int N>
+MUTE_HOST_DEVICE constexpr
+auto
+operator*(A const& a, ScaledBasis<T,N> const& e) {
+  auto r = a * e.value();
+  return ScaledBasis<decltype(r),N>{r};
+}
+
+template <class T, int N, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+operator*(ScaledBasis<T,N> const& e, B const& b) {
+  auto r = e.value() * b;
+  return ScaledBasis<decltype(r),N>{r};
+}
+
+// Addition
+template <class T, int N, class U, int M>
+MUTE_HOST_DEVICE constexpr
+auto
+operator+(ScaledBasis<T,N> const& t, ScaledBasis<U,M> const& u) {
+  return as_arithmetic_tuple(t) + as_arithmetic_tuple(u);
+}
+
+template <class T, int N, class... U>
+MUTE_HOST_DEVICE constexpr
+auto
+operator+(ScaledBasis<T,N> const& t, ArithmeticTuple<U...> const& u) {
+  return as_arithmetic_tuple(t) + u;
+}
+
+template <class... T, class U, int M>
+MUTE_HOST_DEVICE constexpr
+auto
+operator+(ArithmeticTuple<T...> const& t, ScaledBasis<U,M> const& u) {
+  return t + as_arithmetic_tuple(u);
+}
+
+template <auto t, class U, int M>
+MUTE_HOST_DEVICE constexpr
+auto
+operator+(C<t>, ScaledBasis<U,M> const& u) {
+  static_assert(t == 0, "ScaledBasis op+ error!");
+  return u;
+}
+
+template <class T, int N, auto u>
+MUTE_HOST_DEVICE constexpr
+auto
+operator+(ScaledBasis<T,N> const& t, C<u>) {
+  static_assert(u == 0, "ScaledBasis op+ error!");
+  return t;
+}
+
+//
+// Display utilities
+//
+
+template <class ArithTuple>
+MUTE_HOST_DEVICE void print(ArithmeticTupleIterator<ArithTuple> const& iter)
+{
+  printf("ArithTuple"); print(iter.coord_);
+}
+
+template <class T, int N>
+MUTE_HOST_DEVICE void print(ScaledBasis<T,N> const& e)
+{
+  print(e.value()); printf("@%d", N);
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class ArithTuple>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, ArithmeticTupleIterator<ArithTuple> const& iter)
+{
+  return os << "ArithTuple" << iter.coord_;
+}
+
+template <class T, int N>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, ScaledBasis<T,N> const& e)
+{
+  return os << e.value() << "@" << N;
+}
+#endif
+
+} // end namespace mute
+
+
+namespace MUTE_STL_NAMESPACE
+{
+
+template <class... T>
+struct tuple_size<mute::ArithmeticTuple<T...>>
+  : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, mute::ArithmeticTuple<T...>>
+  : MUTE_STL_NAMESPACE::tuple_element<I, MUTE_STL_NAMESPACE::tuple<T...>>
+{};
+
+template <class... T>
+struct tuple_size<const mute::ArithmeticTuple<T...>>
+  : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, const mute::ArithmeticTuple<T...>>
+  : MUTE_STL_NAMESPACE::tuple_element<I, const MUTE_STL_NAMESPACE::tuple<T...>>
+{};
+
+} // end namespace MUTE_STL_NAMESPACE
+
+#ifdef MUTE_STL_NAMESPACE_IS_MUSA_STD
+namespace std
+{
+
+#if defined(__MUSACC_RTC__)
+template <class... _Tp>
+struct tuple_size;
+
+template <size_t _Ip, class... _Tp>
+struct tuple_element;
+#endif
+
+template <class... T>
+struct tuple_size<mute::ArithmeticTuple<T...>>
+  : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, mute::ArithmeticTuple<T...>>
+  : MUTE_STL_NAMESPACE::tuple_element<I, MUTE_STL_NAMESPACE::tuple<T...>>
+{};
+
+template <class... T>
+struct tuple_size<const mute::ArithmeticTuple<T...>>
+  : MUTE_STL_NAMESPACE::integral_constant<size_t, sizeof...(T)>
+{};
+
+template <size_t I, class... T>
+struct tuple_element<I, const mute::ArithmeticTuple<T...>>
+  : MUTE_STL_NAMESPACE::tuple_element<I, const MUTE_STL_NAMESPACE::tuple<T...>>
+{};
+
+} // end namespace std
+#endif // MUTE_STL_NAMESPACE_IS_MUSA_STD
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/bfloat.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/bfloat.hpp
new file mode 100644
index 000000000..f74c6dade
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/bfloat.hpp
@@ -0,0 +1,54 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <vector_types.h>
+#include <mutlass/numeric_types.h>
+
+namespace mute {
+
+using mutlass::bfloat16_t;
+
+//
+// Display utilities
+//
+
+#if !defined(__MUSACC_RTC__)
+MUTE_HOST std::ostream& operator<<(std::ostream& os, bfloat16_t const& v)
+{
+  return os << float(v);
+}
+#endif
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/complex.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/complex.hpp
new file mode 100644
index 000000000..c53521408
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/complex.hpp
@@ -0,0 +1,76 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>    // MUTE_HOST_DEVICE
+
+#include <mutlass/complex.h>  // mutlass::complexm, mutlass::real, mutlass::imag, mutlass::is_complex
+namespace mute
+{
+
+using mutlass::complex;
+using mutlass::is_complex;
+using mutlass::RealType;
+using mutlass::real;
+using mutlass::imag;
+using mutlass::conj;
+
+template <class T>
+static constexpr auto is_complex_v = is_complex<T>::value;
+
+/// Fused multiply-add for complex numbers
+template <class T>
+MUTE_HOST_DEVICE constexpr
+void
+fma(complex<T>      & d,
+    complex<T> const& a,
+    complex<T> const& b,
+    complex<T> const& c)
+{
+  fma(d.real(),  a.real(), b.real(), c.real());
+  fma(d.imag(),  a.real(), b.imag(), c.imag());
+  fma(d.real(), -a.imag(), b.imag(), d.real());
+  fma(d.imag(),  a.imag(), b.real(), d.imag());
+}
+
+/// Fused multiply-add for triplets
+template <class T>
+MUTE_HOST_DEVICE constexpr
+void
+fma(complex<T> const& a,
+    complex<T> const& b,
+    complex<T>      & c)
+{
+  return fma(c, a, b, c);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/float8.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/float8.hpp
new file mode 100644
index 000000000..5155cb96d
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/float8.hpp
@@ -0,0 +1,44 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <vector_types.h>
+#include <mutlass/numeric_types.h>
+
+namespace mute {
+
+using mutlass::float_e4m3_t;
+using mutlass::float_e5m2_t;
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/half.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/half.hpp
new file mode 100644
index 000000000..f1e7e5dc0
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/half.hpp
@@ -0,0 +1,42 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+#include <vector_types.h>
+#include <mutlass/numeric_types.h>
+
+namespace mute {
+
+using mutlass::half_t;
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/int.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/int.hpp
new file mode 100644
index 000000000..1fbaf2533
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/int.hpp
@@ -0,0 +1,136 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <cstdint>
+#endif
+
+#include <mutlass/numeric_types.h>
+
+namespace mute
+{
+
+//
+// Signed integers
+//
+
+using  v2i32_t = int32_t __attribute__((vector_size(8)));
+using  v3i32_t = int32_t __attribute__((vector_size(12)));
+using  v4i32_t = int32_t __attribute__((vector_size(16)));
+using  v5i32_t = int32_t __attribute__((vector_size(20)));
+using  v8i32_t = int32_t __attribute__((vector_size(32)));
+using v16i32_t = int32_t __attribute__((vector_size(64)));
+using v32i32_t = int32_t __attribute__((vector_size(128)));
+
+using int2_t    = mutlass::int2b_t;
+using int4_t    = mutlass::int4b_t;
+using MUTE_STL_NAMESPACE::int8_t;
+using MUTE_STL_NAMESPACE::int16_t;
+using MUTE_STL_NAMESPACE::int32_t;
+using MUTE_STL_NAMESPACE::int64_t;
+using int96_t   = mute::v3i32_t;
+using int128_t  = mute::v4i32_t;
+using int160_t  = mute::v5i32_t;
+using int256_t  = mute::v8i32_t;
+using int512_t  = mute::v16i32_t;
+using int1024_t = mute::v32i32_t;
+
+template <int N> struct int_bit;
+template <> struct int_bit<   2>  { using type = int2_t; };
+template <> struct int_bit<   4>  { using type = int4_t; };
+template <> struct int_bit<   8>  { using type = int8_t;           };
+template <> struct int_bit<  16>  { using type = int16_t;          };
+template <> struct int_bit<  32>  { using type = int32_t;          };
+template <> struct int_bit<  64>  { using type = int64_t;          };
+template <> struct int_bit<  96>  { using type = mute::int96_t;    };
+template <> struct int_bit< 128>  { using type = mute::int128_t;   };
+template <> struct int_bit< 160>  { using type = mute::int160_t;   };
+template <> struct int_bit< 256>  { using type = mute::int256_t;   };
+template <> struct int_bit< 512>  { using type = mute::int512_t;   };
+template <> struct int_bit<1024>  { using type = mute::int1024_t;  };
+
+template <int N>
+using int_bit_t = typename int_bit<N>::type;
+
+template <int N>
+using int_byte = int_bit<8*N>;
+
+template <int N>
+using int_byte_t = typename int_byte<N>::type;
+
+//
+// Unsigned integers
+//
+
+using  v4u32_t = uint32_t __attribute__((vector_size(16)));
+using  v8u32_t = uint32_t __attribute__((vector_size(32)));
+using v16u32_t = uint32_t __attribute__((vector_size(64)));
+using v32u32_t = uint32_t __attribute__((vector_size(128)));
+
+using uint1_t    = mutlass::uint1b_t;
+using uint2_t    = mutlass::uint2b_t;
+using uint4_t    = mutlass::uint4b_t;
+using MUTE_STL_NAMESPACE::uint8_t;
+using MUTE_STL_NAMESPACE::uint16_t;
+using MUTE_STL_NAMESPACE::uint32_t;
+using MUTE_STL_NAMESPACE::uint64_t;
+using mutlass::uint128_t;
+using uint256_t  = mute::v8u32_t;
+using uint512_t  = mute::v16u32_t;
+using uint1024_t = mute::v32u32_t;
+
+template <int N> struct uint_bit;
+template <> struct uint_bit<   1> { using type = uint1_t; };
+template <> struct uint_bit<   2> { using type = uint2_t; };
+template <> struct uint_bit<   4> { using type = uint4_t; };
+template <> struct uint_bit<   8> { using type = uint8_t;           };
+template <> struct uint_bit<  16> { using type = uint16_t;          };
+template <> struct uint_bit<  32> { using type = uint32_t;          };
+template <> struct uint_bit<  64> { using type = uint64_t;          };
+template <> struct uint_bit< 128> { using type = mute::v4u32_t;     };
+template <> struct uint_bit< 256> { using type = mute::uint256_t;   };
+template <> struct uint_bit< 512> { using type = mute::uint512_t;   };
+template <> struct uint_bit<1024> { using type = mute::uint1024_t;  };
+
+template <int N>
+using uint_bit_t = typename uint_bit<N>::type;
+
+template <int N>
+using uint_byte = uint_bit<8*N>;
+
+template <int N>
+using uint_byte_t = typename uint_byte<N>::type;
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/integer_sequence.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/integer_sequence.hpp
new file mode 100644
index 000000000..247c41b89
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/integer_sequence.hpp
@@ -0,0 +1,152 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+#include <mute/util/type_traits.hpp>
+#include <mute/numeric/integral_constant.hpp>
+
+namespace mute
+{
+
+using MUTE_STL_NAMESPACE::integer_sequence;
+using MUTE_STL_NAMESPACE::make_integer_sequence;
+
+namespace detail {
+
+template <class T, class S, T Begin>
+struct range_impl;
+
+template <class T, T... N, T Begin>
+struct range_impl<T, integer_sequence<T, N...>, Begin> {
+  using type = integer_sequence<T, N+Begin...>;
+};
+
+template <class S>
+struct reverse_impl;
+
+template <class T, T... N>
+struct reverse_impl<integer_sequence<T, N...>> {
+  using type = integer_sequence<T, sizeof...(N)-1-N...>;
+};
+
+} // end namespace detail
+
+template <class T, T Begin, T End>
+using make_integer_range = typename detail::range_impl<
+    T,
+    make_integer_sequence<T, (End-Begin > 0) ? (End-Begin) : 0>,
+    Begin>::type;
+
+template <class T, T N>
+using make_integer_sequence_reverse = typename detail::reverse_impl<
+    make_integer_sequence<T, N>>::type;
+
+//
+// Common aliases
+//
+
+// int_sequence
+
+template <int... Ints>
+using int_sequence = integer_sequence<int, Ints...>;
+
+template <int N>
+using make_int_sequence = make_integer_sequence<int, N>;
+
+template <int N>
+using make_int_rsequence = make_integer_sequence_reverse<int, N>;
+
+template <int Begin, int End>
+using make_int_range = make_integer_range<int, Begin, End>;
+
+// index_sequence
+
+template <size_t... Ints>
+using index_sequence = integer_sequence<size_t, Ints...>;
+
+template <size_t N>
+using make_index_sequence = make_integer_sequence<size_t, N>;
+
+template <size_t N>
+using make_index_rsequence = make_integer_sequence_reverse<size_t, N>;
+
+template <size_t Begin, size_t End>
+using make_index_range = make_integer_range<size_t, Begin, End>;
+
+//
+// Shortcuts
+//
+
+template <int... Ints>
+using seq = int_sequence<Ints...>;
+
+template <int N>
+using make_seq = make_int_sequence<N>;
+
+template <int N>
+using make_rseq = make_int_rsequence<N>;
+
+template <int Min, int Max>
+using make_range = make_int_range<Min, Max>;
+
+template <class Tuple>
+using tuple_seq = make_seq<tuple_size<remove_cvref_t<Tuple>>::value>;
+
+template <class Tuple>
+using tuple_rseq = make_rseq<tuple_size<remove_cvref_t<Tuple>>::value>;
+
+//
+// Specialize mute::tuple-traits for std::integer_sequence
+//
+
+template <class T, T... Ints>
+struct tuple_size<integer_sequence<T, Ints...>>
+    : mute::integral_constant<size_t, sizeof...(Ints)>
+{};
+
+template <size_t I, class T, T... Is>
+struct tuple_element<I, integer_sequence<T, Is...>>
+{
+  constexpr static T idx[sizeof...(Is)] = {Is...};
+  using type = mute::integral_constant<T, idx[I]>;
+};
+
+template <size_t I, class T, T... Ints>
+MUTE_HOST_DEVICE constexpr
+tuple_element_t<I, integer_sequence<T, Ints...>>
+get(integer_sequence<T, Ints...>) {
+  static_assert(I < sizeof...(Ints), "Index out of range");
+  return {};
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/integer_subbyte.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/integer_subbyte.hpp
new file mode 100644
index 000000000..d2460da9e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/integer_subbyte.hpp
@@ -0,0 +1,236 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <cstdint>
+#endif
+
+#include <mutlass/integer_subbyte.h>
+
+#include <mute/config.hpp>
+#include <mute/util/type_traits.hpp>
+
+namespace mute {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int Bits, bool Signed = true>
+struct integer_subbyte
+{
+  /// Storage type
+  using Storage = uint8_t;
+
+  /// Number of bits
+  static_assert(Bits <= 8*sizeof(Storage), "Require a subbyte of bits in integer_subbyte");
+
+  /// External type
+  using xint_t = typename conditional<Signed, int, unsigned>::type;
+
+  /// Bitmask for truncation from larger integers
+  static constexpr Storage bits_mask_ = Storage((1 << Bits) - 1);
+  /// Bitmask for the sign bit
+  static constexpr Storage sign_mask_ = Storage((Signed ? 1 : 0) << (Bits - 1));
+
+  //
+  // Data members
+  //
+
+  Storage storage;
+
+  //
+  // Methods
+  //
+
+  /// No operation
+  MUTE_HOST_DEVICE constexpr
+  integer_subbyte() {}
+
+  /// Conversion from integer type
+  MUTE_HOST_DEVICE constexpr
+  integer_subbyte(int value)   // NOTE: Sign extension?
+      : storage(reinterpret_cast<Storage const&>(value) & bits_mask_) {}
+
+  MUTE_HOST_DEVICE constexpr
+  integer_subbyte(unsigned value)
+      : storage(reinterpret_cast<Storage const&>(value) & bits_mask_) {}
+
+  /// Convert to int or unsigned
+  MUTE_HOST_DEVICE constexpr
+  operator xint_t() const {
+    if (sign_mask_ & storage) {  // Sign extend
+      return xint_t(storage) | ~xint_t(bits_mask_);
+    } else {
+      return xint_t(storage);
+    }
+  }
+
+  /// Equality
+  MUTE_HOST_DEVICE constexpr
+  bool operator==(integer_subbyte const& rhs) const {
+    return storage == rhs.storage;
+  }
+
+  /// Inequality
+  MUTE_HOST_DEVICE constexpr
+  bool operator!=(integer_subbyte const& rhs) const {
+    return storage != rhs.storage;
+  }
+
+  /// Less than or equal
+  MUTE_HOST_DEVICE constexpr
+  bool operator<=(integer_subbyte const& rhs) const {
+    if (sign_mask_ & storage) {
+      return !(rhs.storage < storage);
+    } else {
+      return storage <= rhs.storage;
+    }
+  }
+
+  /// Less than
+  MUTE_HOST_DEVICE constexpr
+  bool operator<(integer_subbyte const& rhs) const {
+    if (sign_mask_ & storage) {
+      return !(rhs.storage <= storage);
+    } else {
+      return storage < rhs.storage;
+    }
+  }
+
+  /// Greater than or equal
+  MUTE_HOST_DEVICE constexpr
+  bool operator>=(integer_subbyte const& rhs) const {
+    return !(*this < rhs);
+  }
+
+  /// Greater than
+  MUTE_HOST_DEVICE constexpr
+  bool operator>(integer_subbyte const& rhs) const {
+    return !(*this <= rhs);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 1-bit unsigned integer type
+using uint1b_t = integer_subbyte<1, false>;
+
+/// 2-bit integer type
+using int2b_t = integer_subbyte<2, true>;
+
+/// 2-bit unsigned integer type
+using uint2b_t = integer_subbyte<2, false>;
+
+/// 4-bit integer type
+using int4b_t = integer_subbyte<4, true>;
+
+/// 4-bit unsigned integer type
+using uint4b_t = integer_subbyte<4, false>;
+
+/// 1-bit binary type
+using bin1_t = bool;
+
+} // namespace mute
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if !defined(__MUSACC_RTC__)
+
+#include <limits>
+
+namespace MUTE_STL_NAMESPACE {
+
+template <>
+struct numeric_limits<mute::uint1b_t> {
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint1b_t const lowest() noexcept { return 0; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint1b_t const min() noexcept { return 0; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint1b_t const max() noexcept { return 1; }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+template <>
+struct numeric_limits<mute::int2b_t> {
+  MUTE_HOST_DEVICE static constexpr
+  mute::int2b_t lowest() noexcept { return -2; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::int2b_t min() noexcept { return -2; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::int2b_t max() noexcept { return 1; }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = true;
+};
+
+template <>
+struct numeric_limits<mute::uint2b_t> {
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint2b_t const lowest() noexcept { return 0; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint2b_t const min() noexcept { return 0; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint2b_t const max() noexcept { return 3; }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+template <>
+struct numeric_limits<mute::int4b_t> {
+  MUTE_HOST_DEVICE static constexpr
+  mute::int4b_t lowest() noexcept { return -8; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::int4b_t min() noexcept { return -8; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::int4b_t max() noexcept { return 7; }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = true;
+};
+
+template <>
+struct numeric_limits<mute::uint4b_t> {
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint4b_t const lowest() noexcept { return 0; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint4b_t const min() noexcept { return 0; }
+  MUTE_HOST_DEVICE static constexpr
+  mute::uint4b_t const max() noexcept { return 15; }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+}  // namespace std
+
+#endif // !defined(__MUSACC_RTC__)
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/integral_constant.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/integral_constant.hpp
new file mode 100644
index 000000000..09ffa44de
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/integral_constant.hpp
@@ -0,0 +1,491 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/numeric/math.hpp>      // mute::max, etc
+#include <mute/util/print.hpp>        // mute::print
+#include <mute/util/type_traits.hpp>  // __MUTE_REQUIRES, mute::is_std_integral
+
+namespace mute
+{
+
+// A constant value: short name and type-deduction for fast compilation
+template <auto v>
+struct C {
+  using type = C<v>;
+  static constexpr auto value = v;
+  using value_type = decltype(v);
+  MUTE_HOST_DEVICE constexpr operator   value_type() const noexcept { return value; }
+  MUTE_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; }
+};
+
+// Deprecate
+template <class T, T v>
+using constant = C<v>;
+
+template <bool b>
+using bool_constant = C<b>;
+
+using true_type  = bool_constant<true>;
+using false_type = bool_constant<false>;
+
+// A more std:: conforming integral_constant that enforces type but interops with C<v>
+template <class T, T v>
+struct integral_constant : C<v> {
+  using type = integral_constant<T,v>;
+  static constexpr T value = v;
+  using value_type = T;
+  // Disambiguate C<v>::operator value_type()
+  //MUTE_HOST_DEVICE constexpr operator   value_type() const noexcept { return value; }
+  MUTE_HOST_DEVICE constexpr value_type operator()() const noexcept { return value; }
+};
+
+//
+// Traits
+//
+
+// Use mute::is_std_integral<T> to match built-in integral types (int, int64_t, unsigned, etc)
+// Use mute::is_integral<T> to match both built-in integral types AND static integral types.
+
+template <class T>
+struct is_integral : bool_constant<is_std_integral<T>::value> {};
+template <auto v>
+struct is_integral<C<v>                  > : true_type {};
+template <class T, T v>
+struct is_integral<integral_constant<T,v>> : true_type {};
+
+// is_static detects if an (abstract) value is defined completely by it's type (no members)
+
+template <class T>
+struct is_static : bool_constant<is_empty<remove_cvref_t<T>>::value> {};
+
+template <class T>
+constexpr bool is_static_v = is_static<T>::value;
+
+// is_constant detects if a type is a static integral type and if v is equal to a value
+
+template <auto n, class T>
+struct is_constant : false_type {};
+template <auto n, class T>
+struct is_constant<n, T const > : is_constant<n,T> {};
+template <auto n, class T>
+struct is_constant<n, T const&> : is_constant<n,T> {};
+template <auto n, class T>
+struct is_constant<n, T      &> : is_constant<n,T> {};
+template <auto n, class T>
+struct is_constant<n, T     &&> : is_constant<n,T> {};
+template <auto n, auto v>
+struct is_constant<n, C<v>                  > : bool_constant<v == n> {};
+template <auto n, class T, T v>
+struct is_constant<n, integral_constant<T,v>> : bool_constant<v == n> {};
+
+//
+// Specializations
+//
+
+template <int v>
+using Int = C<v>;
+
+using _m32    = Int<-32>;
+using _m24    = Int<-24>;
+using _m16    = Int<-16>;
+using _m12    = Int<-12>;
+using _m10    = Int<-10>;
+using _m9     = Int<-9>;
+using _m8     = Int<-8>;
+using _m7     = Int<-7>;
+using _m6     = Int<-6>;
+using _m5     = Int<-5>;
+using _m4     = Int<-4>;
+using _m3     = Int<-3>;
+using _m2     = Int<-2>;
+using _m1     = Int<-1>;
+using _0      = Int<0>;
+using _1      = Int<1>;
+using _2      = Int<2>;
+using _3      = Int<3>;
+using _4      = Int<4>;
+using _5      = Int<5>;
+using _6      = Int<6>;
+using _7      = Int<7>;
+using _8      = Int<8>;
+using _9      = Int<9>;
+using _10     = Int<10>;
+using _12     = Int<12>;
+using _16     = Int<16>;
+using _24     = Int<24>;
+using _32     = Int<32>;
+using _64     = Int<64>;
+using _96     = Int<96>;
+using _128    = Int<128>;
+using _192    = Int<192>;
+using _256    = Int<256>;
+using _384    = Int<384>;
+using _512    = Int<512>;
+using _768    = Int<768>;
+using _1024   = Int<1024>;
+using _2048   = Int<2048>;
+using _4096   = Int<4096>;
+using _8192   = Int<8192>;
+using _16384  = Int<16384>;
+using _32768  = Int<32768>;
+using _65536  = Int<65536>;
+using _131072 = Int<131072>;
+using _262144 = Int<262144>;
+using _524288 = Int<524288>;
+
+/***************/
+/** Operators **/
+/***************/
+
+#define MUTE_LEFT_UNARY_OP(OP)                                       \
+  template <auto t>                                                  \
+  MUTE_HOST_DEVICE constexpr                                         \
+  C<(OP t)> operator OP (C<t>) {                                     \
+    return {};                                                       \
+  }
+#define MUTE_RIGHT_UNARY_OP(OP)                                      \
+  template <auto t>                                                  \
+  MUTE_HOST_DEVICE constexpr                                         \
+  C<(t OP)> operator OP (C<t>) {                                     \
+    return {};                                                       \
+  }
+#define MUTE_BINARY_OP(OP)                                           \
+  template <auto t, auto u>                                          \
+  MUTE_HOST_DEVICE constexpr                                         \
+  C<(t OP u)> operator OP (C<t>, C<u>) {                             \
+    return {};                                                       \
+  }
+
+MUTE_LEFT_UNARY_OP(+);
+MUTE_LEFT_UNARY_OP(-);
+MUTE_LEFT_UNARY_OP(~);
+MUTE_LEFT_UNARY_OP(!);
+MUTE_LEFT_UNARY_OP(*);
+
+MUTE_BINARY_OP( +);
+MUTE_BINARY_OP( -);
+MUTE_BINARY_OP( *);
+MUTE_BINARY_OP( /);
+MUTE_BINARY_OP( %);
+MUTE_BINARY_OP( &);
+MUTE_BINARY_OP( |);
+MUTE_BINARY_OP( ^);
+MUTE_BINARY_OP(<<);
+MUTE_BINARY_OP(>>);
+
+MUTE_BINARY_OP(&&);
+MUTE_BINARY_OP(||);
+
+MUTE_BINARY_OP(==);
+MUTE_BINARY_OP(!=);
+MUTE_BINARY_OP( >);
+MUTE_BINARY_OP( <);
+MUTE_BINARY_OP(>=);
+MUTE_BINARY_OP(<=);
+
+#undef MUTE_BINARY_OP
+#undef MUTE_LEFT_UNARY_OP
+#undef MUTE_RIGHT_UNARY_OP
+
+//
+// Mixed static-dynamic special cases
+//
+
+template <auto t, class U,
+          __MUTE_REQUIRES(is_std_integral<U>::value && t == 0)>
+MUTE_HOST_DEVICE constexpr
+C<0>
+operator*(C<t>, U) {
+  return {};
+}
+
+template <class U, auto t,
+          __MUTE_REQUIRES(is_std_integral<U>::value && t == 0)>
+MUTE_HOST_DEVICE constexpr
+C<0>
+operator*(U, C<t>) {
+  return {};
+}
+
+template <auto t, class U,
+          __MUTE_REQUIRES(is_std_integral<U>::value && t == 0)>
+MUTE_HOST_DEVICE constexpr
+C<0>
+operator/(C<t>, U) {
+  return {};
+}
+
+template <class U, auto t,
+          __MUTE_REQUIRES(is_std_integral<U>::value && (t == 1 || t == -1))>
+MUTE_HOST_DEVICE constexpr
+C<0>
+operator%(U, C<t>) {
+  return {};
+}
+
+template <auto t, class U,
+          __MUTE_REQUIRES(is_std_integral<U>::value && t == 0)>
+MUTE_HOST_DEVICE constexpr
+C<0>
+operator%(C<t>, U) {
+  return {};
+}
+
+template <auto t, class U,
+          __MUTE_REQUIRES(is_std_integral<U>::value && t == 0)>
+MUTE_HOST_DEVICE constexpr
+C<0>
+operator&(C<t>, U) {
+  return {};
+}
+
+template <class U, auto t,
+          __MUTE_REQUIRES(is_std_integral<U>::value && t == 0)>
+MUTE_HOST_DEVICE constexpr
+C<0>
+operator&(U, C<t>) {
+  return {};
+}
+
+template <auto t, class U,
+          __MUTE_REQUIRES(is_std_integral<U>::value && !bool(t))>
+MUTE_HOST_DEVICE constexpr
+C<false>
+operator&&(C<t>, U) {
+  return {};
+}
+
+template <auto t, class U,
+          __MUTE_REQUIRES(is_std_integral<U>::value && !bool(t))>
+MUTE_HOST_DEVICE constexpr
+C<false>
+operator&&(U, C<t>) {
+  return {};
+}
+
+template <class U, auto t,
+          __MUTE_REQUIRES(is_std_integral<U>::value && bool(t))>
+MUTE_HOST_DEVICE constexpr
+C<true>
+operator||(C<t>, U) {
+  return {};
+}
+
+template <class U, auto t,
+          __MUTE_REQUIRES(is_std_integral<U>::value && bool(t))>
+MUTE_HOST_DEVICE constexpr
+C<true>
+operator||(U, C<t>) {
+  return {};
+}
+
+//
+// Named functions from math.hpp
+//
+
+#define MUTE_NAMED_UNARY_FN(OP)                                      \
+  template <auto t>                                                  \
+  MUTE_HOST_DEVICE constexpr                                         \
+  C<OP(t)> OP (C<t>) {                                               \
+    return {};                                                       \
+  }
+#define MUTE_NAMED_BINARY_FN(OP)                                     \
+  template <auto t, auto u>                                          \
+  MUTE_HOST_DEVICE constexpr                                         \
+  C<OP(t,u)> OP (C<t>, C<u>) {                                       \
+    return {};                                                       \
+  }                                                                  \
+  template <auto t, class U,                                         \
+            __MUTE_REQUIRES(is_std_integral<U>::value)>              \
+  MUTE_HOST_DEVICE constexpr                                         \
+  auto OP (C<t>, U u) {                                              \
+    return OP(t,u);                                                  \
+  }                                                                  \
+  template <class T, auto u,                                         \
+            __MUTE_REQUIRES(is_std_integral<T>::value)>              \
+  MUTE_HOST_DEVICE constexpr                                         \
+  auto OP (T t, C<u>) {                                              \
+    return OP(t,u);                                                  \
+  }
+
+MUTE_NAMED_UNARY_FN(abs);
+MUTE_NAMED_UNARY_FN(signum);
+MUTE_NAMED_UNARY_FN(has_single_bit);
+
+MUTE_NAMED_BINARY_FN(max);
+MUTE_NAMED_BINARY_FN(min);
+MUTE_NAMED_BINARY_FN(shiftl);
+MUTE_NAMED_BINARY_FN(shiftr);
+MUTE_NAMED_BINARY_FN(gcd);
+MUTE_NAMED_BINARY_FN(lcm);
+
+#undef MUTE_NAMED_UNARY_FN
+#undef MUTE_NAMED_BINARY_FN
+
+//
+// Other functions
+//
+
+template <auto t, auto u>
+MUTE_HOST_DEVICE constexpr
+C<t / u>
+safe_div(C<t>, C<u>) {
+  static_assert(t % u == 0, "Static safe_div requires t % u == 0");
+  return {};
+}
+
+template <auto t, class U,
+          __MUTE_REQUIRES(is_std_integral<U>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+safe_div(C<t>, U u) {
+  return t / u;
+}
+
+template <class T, auto u,
+          __MUTE_REQUIRES(is_std_integral<T>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+safe_div(T t, C<u>) {
+  return t / u;
+}
+
+template <class TrueType, class FalseType>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+conditional_return(true_type, TrueType&& t, FalseType&&) {
+  return static_cast<TrueType&&>(t);
+}
+
+template <class TrueType, class FalseType>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+conditional_return(false_type, TrueType&&, FalseType&& f) {
+  return static_cast<FalseType&&>(f);
+}
+
+template <auto v>
+MUTE_HOST_DEVICE constexpr
+auto
+conditional_return(bool b, C<v> const&, C<v> const&) {
+  return C<v>{};
+}
+
+template <auto v, auto u>
+MUTE_HOST_DEVICE constexpr
+auto
+conditional_return(bool b, C<v> const&, C<u> const&) {
+  return b ? v : u;
+}
+// TrueType and FalseType must have a common type
+template <class TrueType, class FalseType>
+MUTE_HOST_DEVICE constexpr
+auto
+conditional_return(bool b, TrueType const& t, FalseType const& f) {
+  return b ? t : f;
+}
+
+// TrueType and FalseType don't require a common type
+template <bool b, class TrueType, class FalseType>
+MUTE_HOST_DEVICE constexpr
+auto
+conditional_return(TrueType const& t, FalseType const& f) {
+  if constexpr (b) {
+    return t;
+  } else {
+    return f;
+  }
+}
+
+template <class Trait>
+MUTE_HOST_DEVICE constexpr
+auto
+static_value()
+{
+  if constexpr (is_std_integral<decltype(Trait::value)>::value) {
+    return Int<Trait::value>{};
+  } else {
+    return Trait::value;
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Display utilities
+//
+
+template <auto Value>
+MUTE_HOST_DEVICE void print(C<Value>) {
+  printf("_");
+  ::mute::print(Value);
+}
+
+#if !defined(__MUSACC_RTC__)
+template <auto t>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, C<t> const&) {
+  return os << "_" << t;
+}
+#endif
+
+
+namespace detail {
+
+// parse_int_digits takes a variadic number of digits and converts them into an int
+template <class... Ts>
+constexpr uint64_t parse_int_digits(uint64_t result, int digit, Ts... digits)
+{
+  if constexpr (sizeof...(Ts) == 0) {
+    return 10 * result + digit;
+  } else {
+    return parse_int_digits(10 * result + digit, digits...);
+  }
+}
+
+} // end namespace detail
+
+
+// This user-defined literal operator allows mute::constant written as literals. For example,
+//
+//    auto var = 32_c;
+//
+//  var has type mute::constant<int,32>.
+//
+template <char... digits>
+constexpr mute::constant<int,detail::parse_int_digits(0, (digits - '0')...)> operator "" _c()
+{
+  static_assert((('0' <= digits && digits <= '9') && ...),
+                "Expected 0 <= digit <= 9 for each digit of the integer.");
+  return {};
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/integral_ratio.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/integral_ratio.hpp
new file mode 100644
index 000000000..3a0963861
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/integral_ratio.hpp
@@ -0,0 +1,295 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/numeric/integral_constant.hpp>
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE
+#include <mute/numeric/integral_constant.hpp>  // mute::false_type, mute::true_type
+#include <mute/numeric/math.hpp>               // mute::signum
+#include <mute/util/type_traits.hpp>           // __MUTE_REQUIRES
+
+namespace mute
+{
+
+/** Compile-time rational arithmetic type.
+ * Like mute::C for std::integral_constant, mute::R for std::ratio has a short name
+ *   for error messages and compile times.
+ * The static data members @a num and @a den represent the reduced numerator and denominator
+ *   of the rational value. Thus, two mute::R types with different @a n or @a d are distinct types
+ *   even if they represent the same rational value.
+ * A mute::R exposes the reduced canonical type via its ::type member.
+ *   That is, mute::R<3,6>::type is mute::R<1,2> and mute::R<6,3>::type is mute::C<2>.
+ * A mute::R<n,d>::value can be used much like any other trait::value. It can be involved in
+ *   arithmetic expressions (according to the operator-overloads for mute::C and mute::R,
+ *   though these may be incomplete) but with a potential rational value rather than an integral value.
+ */
+template <auto n, auto d>
+class R {
+  static_assert(d != 0);
+  static constexpr auto an  = abs(n);
+  static constexpr auto ad  = abs(d);
+  static constexpr auto g   = gcd(an, ad);
+
+ public:
+  static constexpr auto num = signum(n) * signum(d) * an / g;
+  static constexpr auto den =                         ad / g;
+  // RI: den >= 1 && gcd(abs(num),den) == 1
+  using type = typename conditional<num == 0 || den == 1, C<num>, R<num,den>>::type;
+};
+
+template <class T>
+struct is_ratio : false_type {};
+template <auto n, auto d>
+struct is_ratio<R<n,d>> : true_type {};
+
+template <auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+typename R<a,b>::type
+ratio(C<a>, C<b>) {
+  return {};
+}
+
+template <auto a, auto b, auto c>
+MUTE_HOST_DEVICE constexpr
+typename R<a*c,b>::type
+ratio(C<a>, R<b,c>) {
+  return {};
+}
+
+template <auto a, auto b, auto c>
+MUTE_HOST_DEVICE constexpr
+typename R<b,a*c>::type
+ratio(R<b,c>, C<a>) {
+  return {};
+}
+
+template <auto a, auto b, auto c, auto d>
+MUTE_HOST_DEVICE constexpr
+typename R<a*d,b*c>::type
+ratio(R<a,b>, R<c,d>) {
+  return {};
+}
+
+//
+// Non-reduced ratio implementations
+//
+
+template <auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+R<a,b>
+nratio(C<a>, C<b>) {
+  return {};
+}
+
+template <auto a, auto b, auto c>
+MUTE_HOST_DEVICE constexpr
+R<a*c,b>
+nratio(C<a>, R<b,c>) {
+  return {};
+}
+
+template <auto a, auto b, auto c>
+MUTE_HOST_DEVICE constexpr
+R<b,a*c>
+nratio(R<b,c>, C<a>) {
+  return {};
+}
+
+template <auto a, auto b, auto c, auto d>
+MUTE_HOST_DEVICE constexpr
+R<a*d,b*c>
+nratio(R<a,b>, R<c,d>) {
+  return {};
+}
+
+//
+// Operators
+//
+
+template <auto a, auto b, auto x, auto y>
+MUTE_HOST_DEVICE constexpr
+typename R<a*x,b*y>::type
+operator*(R<a,b>, R<x,y>) {
+  return {};
+}
+
+template <auto a, auto b, auto c>
+MUTE_HOST_DEVICE constexpr
+typename R<a*c,b>::type
+operator*(R<a,b>, C<c>) {
+  return {};
+}
+
+template <auto c, auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+typename R<a*c,b>::type
+operator*(C<c>, R<a,b>) {
+  return {};
+}
+
+// Product with dynamic type needs to produce an integer...
+template <class C, auto a, auto b,
+          __MUTE_REQUIRES(mute::is_std_integral<C>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+operator*(C const& c, R<a,b>) {
+  return c * R<a,b>::num / R<a,b>::den;
+}
+
+// Product with dynamic type needs to produce an integer...
+template <auto a, auto b, class C,
+          __MUTE_REQUIRES(mute::is_std_integral<C>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+operator*(R<a,b>, C const& c) {
+  return c * R<a,b>::num / R<a,b>::den;
+}
+
+template <class C, auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+auto
+operator/(C const& c, R<a,b>) {
+  return c * R<b,a>{};
+}
+
+template <auto a, auto b, auto x, auto y>
+MUTE_HOST_DEVICE constexpr
+typename R<a*y+b*x, b*y>::type
+operator+(R<a,b>, R<x,y>) {
+  return {};
+}
+
+template <auto a, auto b, auto c>
+MUTE_HOST_DEVICE constexpr
+typename R<a+c*b,b>::type
+operator+(R<a,b>, C<c>) {
+  return {};
+}
+
+template <auto c, auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+typename R<a+c*b,b>::type
+operator+(C<c>, R<a,b>) {
+  return {};
+}
+
+/////////////////
+// Comparisons //
+/////////////////
+
+template <auto a, auto b, auto x, auto y>
+MUTE_HOST_DEVICE constexpr
+bool_constant<R<a,b>::num == R<x,y>::num && R<a,b>::den == R<x,y>::den>
+operator==(R<a,b>, R<x,y>) {
+  return {};
+}
+
+template <auto a, auto b, auto c>
+MUTE_HOST_DEVICE constexpr
+bool_constant<R<a,b>::num == c && R<a,b>::den == 1>
+operator==(R<a,b>, C<c>) {
+  return {};
+}
+
+template <auto c, auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+bool_constant<R<a,b>::num == c && R<a,b>::den == 1>
+operator==(C<c>, R<a,b>) {
+  return {};
+}
+
+///////////////////////
+// Special functions //
+///////////////////////
+
+template <auto a, auto b, auto x, auto y>
+MUTE_HOST_DEVICE constexpr
+typename R<gcd(a*y,b*x),b*x>::type
+gcd(R<a,b>, R<x,y>) {
+  return {};
+}
+
+template <auto a, auto b, auto c>
+MUTE_HOST_DEVICE constexpr
+typename R<gcd(a,b*c),b*c>::type
+gcd(R<a,b>, C<c>) {
+  return {};
+}
+
+template <auto c, auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+typename R<gcd(a,b*c),b*c>::type
+gcd(C<c>, R<a,b>) {
+  return {};
+}
+
+template <auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+typename R<abs(a),abs(b)>::type
+abs(R<a,b>) {
+  return {};
+}
+
+template <auto a, auto b>
+MUTE_HOST_DEVICE constexpr
+int32_t
+log_2(R<a,b>) {
+  static_assert(R<a,b>::num > 0);
+  static_assert(R<a,b>::den > 0);
+  return log_2(static_cast<uint32_t>(R<a,b>::num)) - log_2(static_cast<uint32_t>(R<a,b>::den));
+}
+
+// @return A non-reduced ratio mute::R of the Trait0::value / Trait1::value
+template <class Trait0, class Trait1>
+MUTE_HOST_DEVICE constexpr
+auto
+trait_ratio(Trait0, Trait1) {
+  return nratio(static_value<Trait0>(), static_value<Trait1>());
+}
+
+//
+// Display utilities
+//
+
+template <auto a, auto b>
+MUTE_HOST_DEVICE void print(R<a,b>) {
+  print(C<a>{}); print("/"); print(C<b>{});
+}
+
+#if !defined(__MUSACC_RTC__)
+template <auto a, auto b>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, R<a,b>) {
+  return os << "_" << C<a>{} << "/" << C<b>{};
+}
+#endif
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/math.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/math.hpp
new file mode 100644
index 000000000..7a355805d
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/math.hpp
@@ -0,0 +1,357 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>            // MUTE_HOST_DEVICE
+#include <mute/util/type_traits.hpp>  // __MUTE_REQUIRES
+
+#include <mutlass/fast_math.h>
+
+namespace mute
+{
+
+//
+// Common Operations
+//
+
+template <class T, class U,
+          __MUTE_REQUIRES(is_arithmetic<T>::value &&
+                          is_arithmetic<U>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+max(T const& t, U const& u) {
+  return t < u ? u : t;
+}
+
+template <class T, class U,
+          __MUTE_REQUIRES(is_arithmetic<T>::value &&
+                          is_arithmetic<U>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+min(T const& t, U const& u) {
+  return t < u ? t : u;
+}
+
+template <class T,
+          __MUTE_REQUIRES(is_arithmetic<T>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+abs(T const& t) {
+  if constexpr (is_signed<T>::value) {
+    return t < T(0) ? -t : t;
+  } else {
+    return t;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Returns 1 if x > 0, -1 if x < 0, and 0 if x is zero.
+template <class T,
+          __MUTE_REQUIRES(is_arithmetic<T>::value)>
+MUTE_HOST_DEVICE constexpr
+int
+signum(T const& x) {
+  if constexpr (is_signed<T>::value) {
+    return (T(0) < x) - (x < T(0));
+  } else {
+    return T(0) < x;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// C++17 <numeric> operations
+//
+
+// Greatest common divisor of two positive integers
+template <class T, class U,
+          __MUTE_REQUIRES(is_std_integral<T>::value &&
+                          is_std_integral<U>::value)>
+MUTE_HOST_DEVICE constexpr
+mute::common_type_t<T, U>
+gcd(T t, U u) {
+  while (true) {
+    if (t == 0) { return u; }
+    u %= t;
+    if (u == 0) { return t; }
+    t %= u;
+  }
+}
+
+// Least common multiple of two positive integers
+template <class T, class U,
+          __MUTE_REQUIRES(is_std_integral<T>::value &&
+                          is_std_integral<U>::value)>
+MUTE_HOST_DEVICE constexpr
+mute::common_type_t<T, U>
+lcm(T const& t, U const& u) {
+  return (t / gcd(t,u)) * u;
+}
+
+//
+// C++20 <bit> operations
+//
+
+// Checks if a number is an integral power of two
+template <class T>
+MUTE_HOST_DEVICE constexpr
+bool
+has_single_bit(T x) {
+  return x != 0 && (x & (x - 1)) == 0;
+}
+
+// Smallest number of bits needed to represent the given value
+//   For x == 0, this is 0
+//   For x != 0, this is 1 + floor(log2(x))
+// bit_width( 0b0000 ) = 0
+// bit_width( 0b0001 ) = 1
+// bit_width( 0b0010 ) = 2
+// bit_width( 0b0011 ) = 2
+// bit_width( 0b0100 ) = 3
+// bit_width( 0b0101 ) = 3
+// bit_width( 0b0110 ) = 3
+// bit_width( 0b0111 ) = 3
+template <class T>
+MUTE_HOST_DEVICE constexpr
+int
+bit_width(T x) {
+  static_assert(is_unsigned<T>::value, "Only to be used for unsigned types.");
+  constexpr int N = (numeric_limits<T>::digits == 64 ? 6 :
+                    (numeric_limits<T>::digits == 32 ? 5 :
+                    (numeric_limits<T>::digits == 16 ? 4 :
+                    (numeric_limits<T>::digits ==  8 ? 3 : (assert(false),0)))));
+  T r = 0;
+  for (int i = N - 1; i >= 0; --i) {
+    T shift = (x > ((T(1) << (T(1) << i))-1)) << i;
+    x >>= shift;
+    r  |= shift;
+  }
+  return r + (x != 0);
+}
+
+// Smallest integral power of two not less than the given value
+// bit_ceil( 0b00000000 ) = 0b00000001
+// bit_ceil( 0b00000001 ) = 0b00000001
+// bit_ceil( 0b00000010 ) = 0b00000010
+// bit_ceil( 0b00000011 ) = 0b00000100
+// bit_ceil( 0b00000100 ) = 0b00000100
+// bit_ceil( 0b00000101 ) = 0b00001000
+// bit_ceil( 0b00000110 ) = 0b00001000
+// bit_ceil( 0b00000111 ) = 0b00001000
+// bit_ceil( 0b00001000 ) = 0b00001000
+// bit_ceil( 0b00001001 ) = 0b00010000
+template <class T>
+MUTE_HOST_DEVICE constexpr
+T
+bit_ceil(T x) {
+  return x == 0 ? T(1) : (T(1) << bit_width(x - 1));
+}
+
+// Largest integral power of two not greater than the given value
+// bit_floor( 0b00000000 ) = 0b00000000
+// bit_floor( 0b00000001 ) = 0b00000001
+// bit_floor( 0b00000010 ) = 0b00000010
+// bit_floor( 0b00000011 ) = 0b00000010
+// bit_floor( 0b00000100 ) = 0b00000100
+// bit_floor( 0b00000101 ) = 0b00000100
+// bit_floor( 0b00000110 ) = 0b00000100
+// bit_floor( 0b00000111 ) = 0b00000100
+// bit_floor( 0b00001000 ) = 0b00001000
+// bit_floor( 0b00001001 ) = 0b00001000
+template <class T>
+MUTE_HOST_DEVICE constexpr
+T
+bit_floor(T x) {
+  return x == 0 ? 0 : (T(1) << (bit_width(x) - 1));
+}
+
+template <class T>
+MUTE_HOST_DEVICE constexpr T rotl(T x, int s);
+template <class T>
+MUTE_HOST_DEVICE constexpr T rotr(T x, int s);
+
+// Computes the result of circular bitwise left-rotation
+template <class T>
+MUTE_HOST_DEVICE constexpr
+T
+rotl(T x, int s) {
+  constexpr int N = numeric_limits<T>::digits;
+  return static_cast<T>(s == 0 ? x : s > 0 ? (x << s) | (x >> (N - s)) : rotr(x, -s));
+}
+
+// Computes the result of circular bitwise right-rotation
+template <class T>
+MUTE_HOST_DEVICE constexpr
+T
+rotr(T x, int s) {
+  constexpr int N = numeric_limits<T>::digits;
+  return static_cast<T>(s == 0 ? x : s > 0 ? (x >> s) | (x << (N - s)) : rotl(x, -s));
+}
+
+// Counts the number of consecutive 0 bits, starting from the most significant bit
+// countl_zero( 0b00000000 ) = 8
+// countl_zero( 0b11111111 ) = 0
+// countl_zero( 0b00011100 ) = 3
+template <class T>
+MUTE_HOST_DEVICE constexpr
+int
+countl_zero(T x) {
+  return numeric_limits<T>::digits - bit_width(x);
+}
+
+// Counts the number of consecutive 1 bits, starting from the most significant bit
+// countl_one( 0b00000000 ) = 0
+// countl_one( 0b11111111 ) = 8
+// countl_one( 0b11100011 ) = 3
+template <class T>
+MUTE_HOST_DEVICE constexpr
+int
+countl_one(T x) {
+  return countl_zero(~x);
+}
+
+// Counts the number of consecutive 0 bits, starting from the least significant bit
+// countr_zero( 0b00000000 ) = 8
+// countr_zero( 0b11111111 ) = 0
+// countr_zero( 0b00011100 ) = 2
+template <class T>
+MUTE_HOST_DEVICE constexpr
+int
+countr_zero(T x) {
+  return x == 0 ? numeric_limits<T>::digits : bit_width(T(x & T(-x))) - 1;  // bit_width of the LSB
+}
+
+// Counts the number of consecutive 1 bits, starting from the least significant bit
+// countr_one( 0b00000000 ) = 0
+// countr_one( 0b11111111 ) = 8
+// countr_one( 0b11100011 ) = 2
+template <class T>
+MUTE_HOST_DEVICE constexpr
+int
+countr_one(T x) {
+  return countr_zero(~x);
+}
+
+// Counts the number of 1 bits in an unsigned integer
+// popcount( 0b00000000 ) = 0
+// popcount( 0b11111111 ) = 8
+// popcount( 0b00011101 ) = 4
+template <class T>
+MUTE_HOST_DEVICE constexpr
+int
+popcount(T x) {
+  int c = 0;
+  while (x) {
+    ++c;
+    x &= x - 1; // clear the least significant bit set
+  }
+  return c;
+}
+
+//
+// Custom operations
+//
+
+// Computes the result of bitwise left-shift
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+shiftl(T x, int s) {
+  return s >= 0 ? (x << s) : (x >> -s);
+}
+
+// Computes the result of bitwise right-shift
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+shiftr(T x, int s) {
+  return s >= 0 ? (x >> s) : (x << -s);
+}
+
+// Safe divide
+// @pre t % u == 0
+// @result t / u
+template <class T, class U,
+          __MUTE_REQUIRES(is_std_integral<T>::value &&
+                          is_std_integral<U>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+safe_div(T const& t, U const& u) {
+  //assert(t % u == 0);
+  return t / u;
+}
+
+/**
+ * log2 computation
+ */
+
+template <class T>
+MUTE_HOST_DEVICE constexpr
+int32_t
+log_2(T x) {
+  assert(x > 0);
+  static_assert(is_unsigned<T>::value, "Only to be used for unsigned integral types.");
+  return static_cast<int32_t>(bit_width(x)) - 1;
+}
+
+template <class IntDiv, class IntMod>
+struct DivModReturnType {
+  IntDiv div_;
+  IntMod mod_;
+  MUTE_HOST_DEVICE constexpr
+  DivModReturnType(IntDiv const& div, IntMod const& mod) : div_(div), mod_(mod) {}
+};
+
+// General divmod
+template <class CInt0, class CInt1>
+MUTE_HOST_DEVICE constexpr
+auto
+divmod(CInt0 const& a, CInt1 const& b) {
+  return DivModReturnType{a / b, a % b};
+}
+
+// Specialized function with fastDivmod input
+template <class CInt>
+MUTE_HOST_DEVICE constexpr
+auto
+divmod(CInt const& a, mutlass::FastDivmod const& b) {
+  using val_div_type = typename mutlass::FastDivmod::value_div_type;
+  using val_mod_type = typename mutlass::FastDivmod::value_mod_type;
+  val_div_type div = 0;
+  val_mod_type mod = 0;
+  b(div, mod, a);
+  return DivModReturnType{div, mod};
+}
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/numeric_types.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/numeric_types.hpp
new file mode 100644
index 000000000..ef74b3bed
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/numeric_types.hpp
@@ -0,0 +1,136 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>          // MUTE_HOST_DEVICE
+#include <mute/numeric/int.hpp>     // mute::int2_t, mute::int4_t, etc
+
+#include <mutlass/numeric_size.h>   // mutlass::sizeof_bits
+#include <mutlass/numeric_types.h>  // mutlass::float_e4m3_t, mutlass::float_e5m2_t, etc
+
+namespace mute {
+
+template <typename T>
+struct sizeof_bits : public mutlass::sizeof_bits<T> {};
+
+// DO NOT change auto to int, sizeof_bits<sparse_elem> use integral_ratio instead of int
+template <class T>
+static constexpr auto sizeof_bits_v = sizeof_bits<T>::value;
+
+using mutlass::bytes_to_bits;
+
+using mutlass::is_subbyte;
+
+template <class T>
+static constexpr auto is_subbyte_v = is_subbyte<T>::value;
+
+using mutlass::half_t;
+using mutlass::bfloat16_t;
+
+using mutlass::tfloat32_t;
+
+// Umbrella floating-point 8-bit data type : type_erased_dynamic_float8_t
+// This umbrella datatype can be enabled when a user provides a specific
+// datatype in runtime argument list.
+using mutlass::type_erased_dynamic_float8_t;
+using mutlass::float_e4m3_t;
+using mutlass::float_e5m2_t;
+
+using mutlass::uint1b_t;
+using mutlass::int2b_t;
+using mutlass::uint2b_t;
+using mutlass::int4b_t;
+using mutlass::uint4b_t;
+using mutlass::bin1_t;
+
+
+//
+// Print utility
+//
+
+MUTE_HOST_DEVICE
+void
+print(half_t a) {
+  printf("%f", static_cast<float>(a));
+}
+
+MUTE_HOST_DEVICE
+void
+print(bfloat16_t a) {
+  printf("%f", static_cast<float>(a));
+}
+
+
+MUTE_HOST_DEVICE
+void
+print(tfloat32_t a) {
+  printf("%f", static_cast<float>(a));
+}
+
+MUTE_HOST_DEVICE
+void
+print(float_e4m3_t a) {
+  printf("%f", static_cast<float>(a));
+}
+
+MUTE_HOST_DEVICE
+void
+print(float_e5m2_t a) {
+  printf("%f", static_cast<float>(a));
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(bfloat16_t v) {
+  printf("%*.2f", 8, float(v));
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(half_t v) {
+  printf("%*.2f", 8, float(v));
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(tfloat32_t v) {
+  printf("%*.2e", 10, static_cast<float>(v));
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(float_e4m3_t t) {
+  printf("%*.2f", 8, static_cast<float>(t));
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(float_e5m2_t t) {
+  printf("%*.2f", 8, static_cast<float>(t));
+}
+
+} // namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/real.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/real.hpp
new file mode 100644
index 000000000..371a54c32
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/real.hpp
@@ -0,0 +1,74 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+namespace mute
+{
+
+/// Generic add
+template <class C, class A, class B>
+MUTE_HOST_DEVICE constexpr
+void
+add(C& c, A const& a, B const& b)
+{
+  c = a + b;
+}
+
+/// Generic multiply
+template <class C, class A, class B>
+MUTE_HOST_DEVICE constexpr
+void
+mul(C& c, A const& a, B const& b)
+{
+  c = a * b;
+}
+/// Generic fused multiply-add
+template <class D, class A, class B, class C>
+MUTE_HOST_DEVICE constexpr
+void
+fma(D& d, A const& a, B const& b, C const& c)
+{
+  d = a * b + c;
+}
+
+/// Fused multiply-add for triplets
+template <class A, class B, class C>
+MUTE_HOST_DEVICE constexpr
+void
+fma(A const& a, B const& b, C& c)
+{
+  return fma(c, a, b, c);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/tfloat.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/tfloat.hpp
new file mode 100644
index 000000000..0e853c633
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/tfloat.hpp
@@ -0,0 +1,54 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <vector_types.h>
+#include <mutlass/numeric_types.h>
+
+namespace mute {
+
+using mutlass::tfloat32_t;
+
+//
+// Display utilities
+//
+
+#if !defined(__MUSACC_RTC__)
+MUTE_HOST std::ostream& operator<<(std::ostream& os, tfloat32_t const& v)
+{
+  return os << float(v);
+}
+#endif
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/numeric/uint128.hpp b/vllm_musa/csrc/flash_attn/include/mute/numeric/uint128.hpp
new file mode 100644
index 000000000..7479b3de6
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/numeric/uint128.hpp
@@ -0,0 +1,260 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <cstdint>
+#include <cstdlib>
+#include <cmath>
+#include <type_traits>
+#include <stdexcept>
+#endif
+
+#include <mute/config.hpp>
+
+/// Optionally enable GCC's built-in type
+#if defined(__x86_64) && !defined(__MUSA_ARCH__)
+#  if defined(__GNUC__) && 0
+#    define MUTE_UINT128_NATIVE
+#  elif defined(_MSC_VER)
+#    define MUTE_INT128_ARITHMETIC
+#    include <intrin.h>
+#  endif
+#endif
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mute {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+///! Unsigned 128b integer type
+struct alignas(16) uint128_t
+{
+  /// Size of one part of the uint's storage in bits
+  static constexpr int storage_bits_ = 64;
+
+  struct hilo
+  {
+    uint64_t lo;
+    uint64_t hi;
+  };
+
+  // Use a union to store either low and high parts or, if present, a built-in 128b integer type.
+  union
+  {
+    struct hilo hilo_;
+
+#if defined(MUTE_UINT128_NATIVE)
+    unsigned __int128 native;
+#endif // defined(MUTE_UINT128_NATIVE)
+  };
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  MUTE_HOST_DEVICE constexpr
+  uint128_t() : hilo_{0, 0} {}
+
+  /// Constructor from uint64
+  MUTE_HOST_DEVICE constexpr
+  uint128_t(uint64_t lo_) : hilo_{lo_, 0} {}
+
+  /// Constructor from two 64b unsigned integers
+  MUTE_HOST_DEVICE constexpr
+  uint128_t(uint64_t lo_, uint64_t hi_) : hilo_{lo_, hi_} {}
+
+  /// Optional constructor from native value
+#if defined(MUTE_UINT128_NATIVE)
+  uint128_t(unsigned __int128 value) : native(value) { }
+#endif
+
+  /// Lossily cast to uint64
+  MUTE_HOST_DEVICE constexpr
+  explicit operator uint64_t() const
+  {
+    return hilo_.lo;
+  }
+
+  template <class Dummy = bool>
+  MUTE_HOST_DEVICE constexpr
+  static void exception()
+  {
+    //static_assert(sizeof(Dummy) == 0, "Not implemented exception!");
+    //abort();
+    //printf("uint128 not implemented!\n");
+  }
+
+  /// Add
+  MUTE_HOST_DEVICE constexpr
+  uint128_t operator+(uint128_t const& rhs) const
+  {
+    uint128_t y;
+#if defined(MUTE_UINT128_NATIVE)
+    y.native = native + rhs.native;
+#else
+    y.hilo_.lo = hilo_.lo + rhs.hilo_.lo;
+    y.hilo_.hi = hilo_.hi + rhs.hilo_.hi + (!y.hilo_.lo && (rhs.hilo_.lo));
+#endif
+    return y;
+  }
+
+  /// Subtract
+  MUTE_HOST_DEVICE constexpr
+  uint128_t operator-(uint128_t const& rhs) const
+  {
+    uint128_t y;
+#if defined(MUTE_UINT128_NATIVE)
+    y.native = native - rhs.native;
+#else
+    y.hilo_.lo = hilo_.lo - rhs.hilo_.lo;
+    y.hilo_.hi = hilo_.hi - rhs.hilo_.hi - (rhs.hilo_.lo && y.hilo_.lo > hilo_.lo);
+#endif
+    return y;
+  }
+
+  /// Multiply by unsigned 64b integer yielding 128b integer
+  MUTE_HOST_DEVICE constexpr
+  uint128_t operator*(uint64_t const& rhs) const
+  {
+    uint128_t y;
+#if defined(MUTE_UINT128_NATIVE)
+    y.native = native * rhs;
+#elif defined(MUTE_INT128_ARITHMETIC)
+    // Multiply by the low part
+    y.hilo_.lo = _umul128(hilo_.lo, rhs, &y.hilo_.hi);
+
+    // Add the high part and ignore the overflow
+    uint64_t overflow;
+    y.hilo_.hi += _umul128(hilo_.hi, rhs, &overflow);
+#else
+    exception();
+#endif
+    return y;
+  }
+
+  /// Divide 128b operation by 64b operation yielding a 64b quotient
+  MUTE_HOST_DEVICE constexpr
+  uint64_t operator/(uint64_t const& divisor) const
+  {
+    uint64_t quotient = 0;
+#if defined(MUTE_UINT128_NATIVE)
+    quotient = uint64_t(native / divisor);
+#elif defined(MUTE_INT128_ARITHMETIC)
+    // implemented using MSVC's arithmetic intrinsics
+    uint64_t remainder = 0;
+    quotient = _udiv128(hilo_.hi, hilo_.lo, divisor, &remainder);
+#else
+    exception();
+#endif
+    return quotient;
+  }
+
+  /// Divide 128b operation by 64b operation yielding a 64b quotient
+  MUTE_HOST_DEVICE constexpr
+  uint64_t operator%(uint64_t const& divisor) const
+  {
+    uint64_t remainder = 0;
+#if defined(MUTE_UINT128_NATIVE)
+    remainder = uint64_t(native % divisor);
+#elif defined(MUTE_INT128_ARITHMETIC)
+    // implemented using MSVC's arithmetic intrinsics
+    (void)_udiv128(hilo_.hi, hilo_.lo, divisor, &remainder);
+#else
+    exception();
+#endif
+    return remainder;
+  }
+
+  /// Computes the quotient and remainder in a single method.
+  MUTE_HOST_DEVICE constexpr
+  uint64_t divmod(uint64_t &remainder, uint64_t divisor) const
+  {
+    uint64_t quotient = 0;
+#if defined(MUTE_UINT128_NATIVE)
+    quotient = uint64_t(native / divisor);
+    remainder = uint64_t(native % divisor);
+#elif defined(MUTE_INT128_ARITHMETIC)
+    // implemented using MSVC's arithmetic intrinsics
+    quotient = _udiv128(hilo_.hi, hilo_.lo, divisor, &remainder);
+#else
+    exception();
+#endif
+    return quotient;
+  }
+
+  /// Left-shifts a 128b unsigned integer
+  MUTE_HOST_DEVICE constexpr
+  uint128_t operator<<(int sh) const
+  {
+    if (sh == 0) {
+      return *this;
+    }
+    else if (sh >= storage_bits_) {
+      return uint128_t(0, hilo_.lo << (sh - storage_bits_));
+    }
+    else {
+      return uint128_t(
+        (hilo_.lo << sh),
+        (hilo_.hi << sh) | uint64_t(hilo_.lo >> (storage_bits_ - sh))
+      );
+    }
+  }
+
+  /// Right-shifts a 128b unsigned integer
+  MUTE_HOST_DEVICE constexpr
+  uint128_t operator>>(int sh) const
+  {
+    if (sh == 0) {
+      return *this;
+    }
+    else if (sh >= storage_bits_) {
+      return uint128_t((hilo_.hi >> (sh - storage_bits_)), 0);
+    }
+    else {
+      return uint128_t(
+        (hilo_.lo >> sh) | (hilo_.hi << (storage_bits_ - sh)),
+        (hilo_.hi >> sh)
+      );
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mute
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mute/pointer.hpp b/vllm_musa/csrc/flash_attn/include/mute/pointer.hpp
new file mode 100644
index 000000000..1fa44e769
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/pointer.hpp
@@ -0,0 +1,332 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE
+#include <mute/pointer_base.hpp>               // mute::iter_adaptor
+#include <mute/pointer_sparse.hpp>
+#include <mute/container/array_subbyte.hpp>    // mute::subbyte_iterator
+#include <mute/numeric/integral_constant.hpp>  // mute::true_type, mute::false_type
+#include <mute/numeric/numeric_types.hpp>      // sizeof_bits
+namespace mute
+{
+
+//
+// recast_ptr<T> -- Create an iterator over values of type T.
+// For most types this will simply be T*, but certain types require more care.
+// Subbyte Types: uint2_t, uint4_t, etc
+//   Requires construction of a subbyte_iterator<T> in order to properly
+//   resolve each element in byte-addressed memory.
+// Sparse Types: sparse_elem<int S, class T>
+//   A type that holds one physical element meant to represent S number of logical elements.
+//   Requires construction of a sparse_ptr that emulates access to the S logical elements.
+//
+
+template <class NewT>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(void* ptr)
+{
+  if constexpr (is_sparse<NewT>::value) {
+    constexpr int sparsity = NewT::sparsity;
+    NewT* p = reinterpret_cast<NewT*>(ptr);
+    return make_sparse_ptr<sparsity>(p);
+  } else
+  if constexpr (mute::is_subbyte_v<NewT>) {
+    return subbyte_iterator<NewT>(ptr);
+  } else {
+    return reinterpret_cast<NewT*>(ptr);
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class NewT>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(void const* ptr)
+{
+  if constexpr (is_sparse<NewT>::value) {
+    constexpr int sparsity = NewT::sparsity;
+    NewT const* p = reinterpret_cast<NewT const*>(ptr);
+    return make_sparse_ptr<sparsity>(p);
+  } else
+  if constexpr (mute::is_subbyte_v<NewT>) {
+    return subbyte_iterator<NewT const>(ptr);
+  } else {
+    return reinterpret_cast<NewT const*>(ptr);
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Disambiguate nullptr
+template <class NewT>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(decltype(nullptr)) {   // nullptr_t
+  return recast_ptr<NewT>(static_cast<NewT*>(nullptr));
+}
+
+//
+// gmem_ptr
+//
+
+template <class P>
+struct gmem_ptr : iter_adaptor<P, gmem_ptr<P>> {
+  using iter_adaptor<P, gmem_ptr<P>>::iter_adaptor;
+};
+
+template <class T, class = void>
+struct is_gmem : false_type {};
+template <class P>                     // Found the gmem
+struct is_gmem<gmem_ptr<P>> : true_type {};
+template <class P>                     // Recurse on ::iterator, if possible
+struct is_gmem<P, void_t<typename P::iterator>> : is_gmem<typename P::iterator> {};
+template <class P>
+constexpr bool is_gmem_v = is_gmem<P>::value;
+
+// Idempotent gmem tag on an iterator
+template <class Iterator>
+MUTE_HOST_DEVICE constexpr
+auto
+make_gmem_ptr(Iterator iter) {
+  if constexpr (is_gmem<Iterator>::value) {
+    return iter;
+  } else {
+    return gmem_ptr<Iterator>{iter};
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Explicitly typed construction from a raw pointer
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_gmem_ptr(void* ptr) {
+  return make_gmem_ptr(recast_ptr<T>(ptr));
+}
+
+// Explicitly typed construction from a raw pointer
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_gmem_ptr(void const* ptr) {
+  return make_gmem_ptr(recast_ptr<T const>(ptr));
+}
+
+// nullptr_t overload for make_gmem_ptr<float>(nullptr) disambiguation
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_gmem_ptr(decltype(nullptr)) { // nullptr_t
+  return make_gmem_ptr(recast_ptr<T>(nullptr));
+}
+
+// The gmem tag is invariant over type-recast
+template <class NewT, class P>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(gmem_ptr<P> const& ptr) {
+  return make_gmem_ptr(recast_ptr<NewT>(ptr.get()));
+}
+
+//
+// smem_ptr
+//
+
+template <class P>
+struct smem_ptr : iter_adaptor<P, smem_ptr<P>> {
+  using iter_adaptor<P, smem_ptr<P>>::iter_adaptor;
+};
+
+template <class T, class = void>
+struct is_smem : false_type {};
+template <class P>                     // Found the smem
+struct is_smem<smem_ptr<P>> : true_type {};
+template <class P>                     // Recurse on ::iterator, if possible
+struct is_smem<P, void_t<typename P::iterator>> : is_smem<typename P::iterator> {};
+template <class P>
+constexpr bool is_smem_v = is_smem<P>::value;
+
+// Idempotent smem tag on an iterator
+template <class Iterator>
+MUTE_HOST_DEVICE constexpr
+auto
+make_smem_ptr(Iterator iter) {
+  if constexpr (is_smem<Iterator>::value) {
+    return iter;
+  } else {
+    return smem_ptr<Iterator>{iter};
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Make a smem swizzle pointer, common operation
+template <class Iterator, class Swizzle>
+MUTE_HOST_DEVICE constexpr
+auto
+make_smem_ptr(Iterator ptr, Swizzle sw)
+{
+  return make_swizzle_ptr(make_smem_ptr(ptr), sw);
+}
+
+// Explicitly typed construction from a raw pointer
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_smem_ptr(void* ptr) {
+  return make_smem_ptr(recast_ptr<T>(ptr));
+}
+
+// Explicitly typed construction from a raw pointer
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_smem_ptr(void const* ptr) {
+  return make_smem_ptr(recast_ptr<T const>(ptr));
+}
+
+
+// nullptr_t overload for make_smem_ptr<float>(nullptr) disambiguation
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_smem_ptr(decltype(nullptr)) { // nullptr_t
+  return make_smem_ptr(recast_ptr<T>(nullptr));
+}
+
+
+// The smem tag is invariant over type-recast
+template <class NewT, class P>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(smem_ptr<P> const& ptr) {
+  return make_smem_ptr(recast_ptr<NewT>(ptr.get()));
+}
+
+//
+// rmem_ptr
+//
+
+template <class P>
+struct rmem_ptr : iter_adaptor<P, rmem_ptr<P>> {
+  using iter_adaptor<P, rmem_ptr<P>>::iter_adaptor;
+};
+
+// Anything that is not gmem or smem is rmem
+template <class T, class = void>
+struct is_rmem : bool_constant<not (is_gmem<T>::value || is_smem<T>::value)> {};
+template <class P>
+struct is_rmem<rmem_ptr<P>> : true_type {};
+template <class P>
+constexpr bool is_rmem_v = is_rmem<P>::value;
+
+// Idempotent rmem tag on an iterator
+template <class Iterator>
+MUTE_HOST_DEVICE constexpr
+auto
+make_rmem_ptr(Iterator iter) {
+  if constexpr (is_rmem<Iterator>::value) {
+    return iter;
+  } else {
+    return rmem_ptr<Iterator>{iter};
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Explicitly typed construction from a raw pointer
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_rmem_ptr(void* ptr) {
+  return make_rmem_ptr(recast_ptr<T>(ptr));
+}
+
+// Explicitly typed construction from a raw pointer
+template <class T>
+MUTE_HOST_DEVICE constexpr
+auto
+make_rmem_ptr(void const* ptr) {
+  return make_rmem_ptr(recast_ptr<T const>(ptr));
+}
+
+// The rmem tag is invariant over type-recast
+template <class NewT, class P>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(rmem_ptr<P> const& ptr) {
+  return make_rmem_ptr(recast_ptr<NewT>(ptr.get()));
+}
+
+//
+// Display utilities
+//
+
+template <class T>
+MUTE_HOST_DEVICE void print(gmem_ptr<T> ptr)
+{
+  printf("gmem_"); print(ptr.get());
+}
+
+template <class T>
+MUTE_HOST_DEVICE void print(smem_ptr<T> ptr)
+{
+  printf("smem_"); print(ptr.get());
+}
+
+template <class T>
+MUTE_HOST_DEVICE void print(rmem_ptr<T> ptr)
+{
+  printf("rmem_"); print(ptr.get());
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class T>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, gmem_ptr<T> ptr)
+{
+  return os << "gmem_[" << int(sizeof_bits<iter_value_t<T>>::value) << "b]";
+}
+
+template <class T>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, smem_ptr<T> ptr)
+{
+  return os << "smem_[" << int(sizeof_bits<iter_value_t<T>>::value) << "b]";
+}
+
+template <class T>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, rmem_ptr<T> ptr)
+{
+  return os << "rmem_[" << int(sizeof_bits<iter_value_t<T>>::value) << "b]";
+}
+
+#endif // !defined(__MUSACC_RTC__)
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/pointer_base.hpp b/vllm_musa/csrc/flash_attn/include/mute/pointer_base.hpp
new file mode 100644
index 000000000..f4fcca8c8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/pointer_base.hpp
@@ -0,0 +1,264 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE
+#include <mute/numeric/numeric_types.hpp>      // mute::sizeof_bits
+#include <mute/numeric/integral_constant.hpp>  // Int<0>
+#include <mute/util/type_traits.hpp>           // mute::declval, mute::void_t, etc
+
+namespace mute
+{
+
+//
+// C++20 <iterator> iterator_traits
+//
+
+namespace detail {
+// Default reference type of an iterator
+template <class T, class = void>
+struct iter_ref { using type = decltype(*declval<T&>()); };
+// Prefer to propagate ::reference
+template <class T>
+struct iter_ref<T,void_t<typename T::reference>> { using type = typename T::reference; };
+} // end namespace detail
+
+template <class T>
+using iter_reference = detail::iter_ref<T>;
+template <class T>
+using iter_reference_t = typename iter_reference<T>::type;
+
+namespace detail {
+// Default element_type of an iterator
+template <class T, class = void>
+struct iter_e { using type = remove_reference_t<typename iter_ref<T>::type>; };
+// Prefer to propagate ::element_type
+template <class T>
+struct iter_e<T,void_t<typename T::element_type>> { using type = typename T::element_type; };
+} // end namespace detail
+
+template <class T>
+using iter_element = detail::iter_e<T>;
+template <class T>
+using iter_element_t = typename iter_element<T>::type;
+
+namespace detail {
+// Default value_type of an iterator
+template <class T, class = void>
+struct iter_v { using type = remove_cv_t<typename iter_e<T>::type>; };
+// Prefer to propagate ::value_type
+template <class T>
+struct iter_v<T,void_t<typename T::value_type>> { using type = typename T::value_type; };
+} // end namespace detail
+
+template <class T>
+using iter_value = detail::iter_v<T>;
+template <class T>
+using iter_value_t = typename iter_value<T>::type;
+
+template <class Iterator>
+struct iterator_traits {
+  using reference    = iter_reference_t<Iterator>;
+  using element_type = iter_element_t<Iterator>;
+  using value_type   = iter_value_t<Iterator>;
+};
+
+//
+// has_dereference to determine if a type is an iterator concept
+//
+
+namespace detail {
+template <class T, class = void>
+struct has_dereference : MUTE_STL_NAMESPACE::false_type {};
+template <class T>
+struct has_dereference<T, void_t<decltype(*declval<T&>())>> : MUTE_STL_NAMESPACE::true_type {};
+} // end namespace detail
+
+template <class T>
+using has_dereference = detail::has_dereference<T>;
+
+//
+// raw_pointer_cast
+//
+
+template <class T>
+MUTE_HOST_DEVICE constexpr
+T*
+raw_pointer_cast(T* ptr) {
+  return ptr;
+}
+
+// The statically-known alignment of a dynamic pointer is unknown
+template <class T>
+MUTE_HOST_DEVICE constexpr
+Int<0>
+max_alignment(T*) {
+  return {};
+}
+
+//
+// A very simplified iterator adaptor.
+// Derived classed may override methods, but be careful to reproduce interfaces exactly.
+// Clients should never have an instance of this class. Do not write methods that take this as a param.
+//
+
+template <class Iterator, class DerivedType>
+struct iter_adaptor
+{
+  using iterator     = Iterator;
+  using reference    = typename iterator_traits<iterator>::reference;
+  using element_type = typename iterator_traits<iterator>::element_type;
+  using value_type   = typename iterator_traits<iterator>::value_type;
+
+  iterator ptr_;
+
+  MUTE_HOST_DEVICE constexpr
+  iter_adaptor(iterator ptr = {}) : ptr_(ptr) {}
+
+  MUTE_HOST_DEVICE constexpr
+  reference operator*() const { return *ptr_; }
+
+  template <class Index>
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](Index const& i) const { return ptr_[i]; }
+
+  template <class Index>
+  MUTE_HOST_DEVICE constexpr
+  DerivedType operator+(Index const& i) const { return {ptr_ + i}; }
+
+  MUTE_HOST_DEVICE constexpr
+  iterator get() const { return ptr_; }
+
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator==(DerivedType const& x, DerivedType const& y) { return x.ptr_ == y.ptr_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator!=(DerivedType const& x, DerivedType const& y) { return x.ptr_ != y.ptr_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator< (DerivedType const& x, DerivedType const& y) { return x.ptr_ <  y.ptr_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator<=(DerivedType const& x, DerivedType const& y) { return x.ptr_ <= y.ptr_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator> (DerivedType const& x, DerivedType const& y) { return x.ptr_ >  y.ptr_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator>=(DerivedType const& x, DerivedType const& y) { return x.ptr_ >= y.ptr_; }
+};
+
+template <class I, class D>
+MUTE_HOST_DEVICE constexpr
+auto
+raw_pointer_cast(iter_adaptor<I,D> const& x) {
+  return raw_pointer_cast(x.ptr_);
+}
+
+template <class I, class D>
+MUTE_HOST_DEVICE constexpr
+auto
+max_alignment(iter_adaptor<I,D> const& x) {
+  return max_alignment(x.ptr_);
+}
+
+
+//
+// counting iterator -- quick and dirty
+//
+
+template <class T = int>
+struct counting_iterator
+{
+  using index_type = T;
+  using value_type = T;
+  using reference  = T;
+
+  index_type n_;
+
+  MUTE_HOST_DEVICE constexpr
+  counting_iterator(index_type n = 0) : n_(n) {}
+
+  MUTE_HOST_DEVICE constexpr
+  index_type operator*() const { return n_; }
+
+  MUTE_HOST_DEVICE constexpr
+  index_type operator[](index_type i) const { return n_ + i; }
+
+  MUTE_HOST_DEVICE constexpr
+  counting_iterator operator+(index_type i) const { return {n_ + i}; }
+  MUTE_HOST_DEVICE constexpr
+  counting_iterator& operator++() { ++n_; return *this; }
+  MUTE_HOST_DEVICE constexpr
+  counting_iterator operator++(int) { counting_iterator ret = *this; ++n_; return ret; }
+
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator==(counting_iterator const& x, counting_iterator const& y) { return x.n_ == y.n_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator!=(counting_iterator const& x, counting_iterator const& y) { return x.n_ != y.n_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator< (counting_iterator const& x, counting_iterator const& y) { return x.n_ <  y.n_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator<=(counting_iterator const& x, counting_iterator const& y) { return x.n_ <= y.n_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator> (counting_iterator const& x, counting_iterator const& y) { return x.n_ >  y.n_; }
+  MUTE_HOST_DEVICE constexpr
+  friend bool operator>=(counting_iterator const& x, counting_iterator const& y) { return x.n_ >= y.n_; }
+};
+
+template <class T>
+MUTE_HOST_DEVICE constexpr
+T
+raw_pointer_cast(counting_iterator<T> const& x) {
+  return x.n_;
+}
+
+//
+// Display utilities
+//
+
+template <class T>
+MUTE_HOST_DEVICE void print(T const* const ptr)
+{
+  printf("ptr["); print(sizeof_bits<T>::value); printf("b](%p)", ptr);
+}
+
+template <class T>
+MUTE_HOST_DEVICE void print(counting_iterator<T> ptr)
+{
+  printf("counting_iter("); print(ptr.n_); printf(")");
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class T>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, counting_iterator<T> ptr)
+{
+  return os << "counting_iter(" << ptr.n_ << ")";
+}
+#endif // !defined(__MUSACC_RTC__)
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/pointer_flagged.hpp b/vllm_musa/csrc/flash_attn/include/mute/pointer_flagged.hpp
new file mode 100644
index 000000000..3cf4c45ea
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/pointer_flagged.hpp
@@ -0,0 +1,200 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE
+#include <mute/layout_composed.hpp>            // mute::ComposedLayout
+#include <mute/pointer.hpp>                    // mute::make_smem_ptr
+#include <mute/pointer_sparse.hpp>             // mute::is_sparse
+#include <mute/pointer_swizzle.hpp>            // mute::make_swizzle_ptr
+#include <mute/arch/util.hpp>                  // mute::cast_smem_ptr_to_uint
+#include <mute/numeric/integral_constant.hpp>  // mute::Int
+
+namespace mute
+{
+
+//
+// Stand-in Swizzle Layout
+//   A model of a nullptr smem_ptr<T> with B == sizeof_bits<T>::value
+//   That represents an unset pointer. This is a placeholder type that is waiting for an smem_ptr
+//
+
+template <int Bits>
+struct smem_ptr_flag_bits : Int<0> {};
+
+using smem_ptr_flag = smem_ptr_flag_bits<1>;
+
+// A flagged construction method to transform ComposedLayout
+// Make a swizzle pointer tensor and check that the intended type size matches
+template <class Iterator, class SwizzleFn, int B, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tensor(Iterator const& ptr,
+            ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
+{
+  static_assert(is_smem<Iterator>::value, "Expected smem.");
+  static_assert(B == sizeof_bits<iter_value_t<Iterator>>::value, "Expected a B-bit pointer type.");
+  return make_tensor(make_smem_ptr(ptr.get(), layout.layout_a()),
+                     layout.layout_b());
+}
+
+// NOTE: To preserve smem_ptr_flag_bits under recast ops
+template <int N, class SwizzleFn, int B, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+upcast(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
+{
+  return composition(layout.layout_a(), smem_ptr_flag_bits<B*N>{}, upcast<N>(layout.layout_b()));
+}
+
+template <int N, class SwizzleFn, int B, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+downcast(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
+{
+  return composition(layout.layout_a(), smem_ptr_flag_bits<B/N>{}, downcast<N>(layout.layout_b()));
+}
+
+//
+// Conversion with swizzle_layout
+//
+
+template <class SwizzleFn, int B, class Layout>
+MUTE_HOST_DEVICE
+auto
+as_position_independent_swizzle_layout(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
+{
+  return composition(recast_layout<uint8_t,uint_bit_t<B>>(layout.layout_a()), Int<0>{}, layout.layout_b());
+}
+
+template <class Tensor>
+MUTE_HOST_DEVICE
+auto
+as_position_independent_swizzle_tensor(Tensor&& tensor)
+{
+  static_assert(is_smem<remove_cvref_t<Tensor>>::value, "Expected smem tensor.");
+  using SwizzleFn = get_swizzle_t<remove_cvref_t<Tensor>>;
+  if constexpr (SwizzleFn::num_bits == 0) {
+    return tensor;
+  } else {
+#if !defined(NDEBUG)
+    {
+    uint32_t address = cast_smem_ptr_to_uint(raw_pointer_cast(static_cast<Tensor&&>(tensor).data()));
+    uint32_t mask    = ((uint32_t(1) << SwizzleFn::num_base) - 1) | SwizzleFn::swizzle_code;
+    assert((address & mask) == 0);  // Alignment to the Base, Z, and Y of Swizzle
+    }
+#endif
+    using T = typename remove_cvref_t<Tensor>::value_type;
+    // Recast swizzle from acting on byte-addressed pointers to elements of type-T
+    auto new_swizzle = recast_layout<uint8_t, T>(SwizzleFn{});
+    // Strip off everything and create a new smem_ptr for type-T
+    auto new_ptr = make_smem_ptr<T>(raw_pointer_cast(static_cast<Tensor&&>(tensor).data()));
+    return make_tensor(new_ptr, composition(new_swizzle, Int<0>{}, tensor.layout()));
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+// A model of a nullptr sparse_ptr<S, smem_ptr<T>> with B == sizeof_bits<T>::value
+// That represents an unset pointer. This is a placeholder type that is waiting for an smem_ptr
+template <int Sparsity, int Bits>
+struct smem_sparse_ptr_flag_bits : Int<0> {};
+
+template <int Sparsity>
+using smem_sparse_ptr_flag = smem_sparse_ptr_flag_bits<Sparsity, 1>;
+
+// A flagged construction method to transform ComposedLayout
+// Make a swizzle pointer tensor and check that the intended type size matches
+template <class Iterator, class SwizzleFn, int S, int B, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tensor(Iterator const& ptr,
+            ComposedLayout<SwizzleFn,smem_sparse_ptr_flag_bits<S,B>,Layout> const& layout)
+{
+  static_assert(is_smem<Iterator>::value, "Expected smem.");
+  static_assert(is_sparse_ptr<Iterator>::value, "Expected sparse iter");
+  static_assert(is_sparse<iter_value_t<Iterator>>::value, "Expected sparse elem");
+  static_assert(S == iter_value_t<Iterator>::sparsity, "Expected sparsity S");
+  static_assert(B == sizeof_bits<typename iter_value_t<Iterator>::raw_type>::value, "Expected B-bit pointer type");
+  return make_tensor(make_swizzle_ptr(ptr, layout.layout_a()), layout.layout_b());
+}
+
+// NOTE: To preserve smem_ptr_flag_bits under recast ops
+template <int N, class SwizzleFn, int S, int B, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+upcast(ComposedLayout<SwizzleFn,smem_sparse_ptr_flag_bits<S,B>,Layout> const& layout)
+{
+  static_assert(dependent_false<SwizzleFn>, "Not implemented for safety");
+}
+
+template <int N, class SwizzleFn, int S, int B, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+downcast(ComposedLayout<SwizzleFn,smem_sparse_ptr_flag_bits<S,B>,Layout> const& layout)
+{
+  static_assert(dependent_false<SwizzleFn>, "Not implemented for safety");
+}
+
+//
+// Display utilities
+//
+
+// Capture and cast smem_ptr_flag Layouts to offset-0 layouts
+template <class SwizzleFn, int B, class Layout>
+MUTE_HOST_DEVICE
+void
+print_layout(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
+{
+  print_layout(as_position_independent_swizzle_layout(layout));
+}
+
+template <class SwizzleFn, int B, class Layout>
+MUTE_HOST_DEVICE
+void
+print_latex(ComposedLayout<SwizzleFn,smem_ptr_flag_bits<B>,Layout> const& layout)
+{
+  print_latex(as_position_independent_swizzle_layout(layout));
+}
+
+template <int B>
+MUTE_HOST_DEVICE void print(smem_ptr_flag_bits<B> ptr)
+{
+  printf("smem_ptr[%db](unset)", B);
+}
+
+template <int S, int B>
+MUTE_HOST_DEVICE void print(smem_sparse_ptr_flag_bits<S,B>)
+{
+  printf("smem_sparse<%d>_ptr[%db](unset)", S, B);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/pointer_sparse.hpp b/vllm_musa/csrc/flash_attn/include/mute/pointer_sparse.hpp
new file mode 100644
index 000000000..c4ae27688
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/pointer_sparse.hpp
@@ -0,0 +1,172 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE
+#include <mute/pointer_base.hpp>               // mute::iter_adaptor
+#include <mute/numeric/integral_constant.hpp>  // mute::false_type, mute::true_type
+#include <mute/numeric/integral_ratio.hpp>     // mute::ratio
+
+namespace mute
+{
+
+// A data type that holds one physical element meant to represent Sparsity number of logical elements
+// This class is purposely not compatible with anything -- know what you're doing if you attempt to use it
+template <int Sparsity, class T>
+struct sparse_elem
+{
+  static constexpr int sparsity = Sparsity;
+  using raw_type = T;
+  T elem_;
+
+  MUTE_HOST_DEVICE constexpr
+  explicit sparse_elem(T const& elem = {}) : elem_(elem) {}
+
+  MUTE_HOST_DEVICE constexpr friend bool operator==(sparse_elem const& a, sparse_elem const& b) { return a.elem_ == b.elem_; }
+  MUTE_HOST_DEVICE constexpr friend bool operator!=(sparse_elem const& a, sparse_elem const& b) { return a.elem_ != b.elem_; }
+  MUTE_HOST_DEVICE constexpr friend bool operator< (sparse_elem const& a, sparse_elem const& b) { return a.elem_ <  b.elem_; }
+  MUTE_HOST_DEVICE constexpr friend bool operator<=(sparse_elem const& a, sparse_elem const& b) { return a.elem_ <= b.elem_; }
+  MUTE_HOST_DEVICE constexpr friend bool operator> (sparse_elem const& a, sparse_elem const& b) { return a.elem_ >  b.elem_; }
+  MUTE_HOST_DEVICE constexpr friend bool operator>=(sparse_elem const& a, sparse_elem const& b) { return a.elem_ >= b.elem_; }
+};
+
+template <class T>
+struct is_sparse : false_type {};
+template <class T>
+struct is_sparse<T const> : is_sparse<T> {};
+template <int S, class T>
+struct is_sparse<sparse_elem<S,T>> : true_type {};
+template<class T>
+static constexpr auto is_sparse_v = is_sparse<T>::value;
+
+// Overload sizeof_bits for sparse_elem.
+//   Much like subbyte element types, this is the effective number of bits in a sparse_elem
+//   rather than actual physical bits that may be used in storing one. Also like subbyte element
+//   types, modified iterators are required to properly index and access sparse_elems.
+//
+//   Defining sizeof_bits like this makes reasonable expressions like N * sizeof_bits_v<E> meaningful
+//   even when E is subbyte or sparse. However, this also means that sparse_elem can rather easily be
+//   confused with subbyte elements and special care should be taken with each.
+template <int S, class T>
+struct sizeof_bits<sparse_elem<S,T>> {
+  // Simple implementation that conforms to sizeof_bits
+  //static constexpr auto value = sizeof_bits<T>::value / S;
+  //static_assert(value != 0, "sizeof_bits=0 detected. Sparsity is larger than width.");
+  //static_assert((sizeof_bits<T>::value % S) == 0, "Width needs to be a multiple of sparsity.")
+
+  // Interesting experiment that allows any sparsity level to be used by potentially presenting
+  // an integral_ratio rather than size_t. This is valid in most integer expressions as well.
+  static constexpr auto value = mute::ratio(mute::Int<mute::sizeof_bits_v<T>>{}, mute::Int<S>{});
+};
+
+//
+// sparse_ptr
+//
+
+template <class T, class = void>
+struct is_sparse_ptr : false_type {};
+template <class T>
+struct is_sparse_ptr<T, void_t<typename T::iterator>> : is_sparse_ptr<typename T::iterator> {};
+
+template <int Sparsity, class Iterator>
+struct sparse_ptr : iter_adaptor<Iterator, sparse_ptr<Sparsity, Iterator>>
+{
+  using reference    = typename iterator_traits<Iterator>::reference;
+  using element_type = typename iterator_traits<Iterator>::element_type;
+  using value_type   = typename iterator_traits<Iterator>::value_type;
+
+  // Sanity, for now
+  static_assert(is_sparse<value_type>::value, "Enforce sparse value-type");
+  static_assert(Sparsity == iter_value_t<Iterator>::sparsity, "Enforce sparsity S");
+  static_assert(not is_sparse_ptr<Iterator>::value, "Enforce sparse singleton");
+
+  template <class Index>
+  MUTE_HOST_DEVICE constexpr
+  sparse_ptr operator+(Index const& i) const {
+    // Only allow offset by multiples of the sparsity factor,
+    // else the misalignments become a bug. E.g. (sparse_ptr<8,I>{} + 7) + 7
+    // Motivation for subsparse_iterator or generalization of subbyte_iterator?
+    assert(i % Sparsity == 0);
+    return {this->get() + i / Sparsity};
+  }
+
+  template <class Index>
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](Index const& i) const {
+    // Allow offset by any value and dereference.
+    // Not implemented in terms of sparse_ptr::op+()
+    return *(this->get() + i / Sparsity);
+  }
+};
+
+template <int S, class I>
+struct is_sparse_ptr<sparse_ptr<S,I>> : true_type {};
+
+template <int Sparsity, class Iter>
+MUTE_HOST_DEVICE constexpr
+auto
+make_sparse_ptr(Iter const& iter) {
+  if constexpr (Sparsity == 1) {
+    return iter;
+  } else {
+    return sparse_ptr<Sparsity, Iter>{iter};
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class NewT, int S, class Iter>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(sparse_ptr<S,Iter> const& ptr) {
+  static_assert(not is_sparse<NewT>::value);
+  return recast_ptr<NewT>(ptr.get());
+}
+
+//
+// Display utilities
+//
+
+template <int S, class Iter>
+MUTE_HOST_DEVICE void print(sparse_ptr<S,Iter> ptr)
+{
+  printf("sparse<%d>_", S); print(ptr.get());
+}
+
+#if !defined(__MUSACC_RTC__)
+template <int S, class Iter>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, sparse_ptr<S,Iter> ptr)
+{
+  return os << "sparse<" << S << ">_" << ptr.get();
+}
+#endif
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/pointer_swizzle.hpp b/vllm_musa/csrc/flash_attn/include/mute/pointer_swizzle.hpp
new file mode 100644
index 000000000..6b3514bb0
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/pointer_swizzle.hpp
@@ -0,0 +1,190 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                   // MUTE_HOST_DEVICE
+#include <mute/pointer_base.hpp>             // mute::iter_adaptor
+#include <mute/swizzle.hpp>                  // mute::Swizzle, mute::get_swizzle primary template
+#include <mute/util/type_traits.hpp>         // mute::iterator_traits
+#include <mute/container/array_subbyte.hpp>  // mute::subbyte_iterator
+
+/* This implements a swizzle pointer of the form
+ *   InvolutionFn o PtrAdd
+ * where the InvolutionFn need not be linear.
+ *
+ * This differs subtly from swizzle_layout because the smem pointer is used
+ * as the offset. That means that swizzle_layout will implement position-independent
+ * swizzle layouts, while swizzle_ptr implements position-dependent swizzle tensors.
+ * Arch chose to design hardware with position-dependent swizzles.
+ *
+ * For clarity:
+ *   NormalLayout  : DeRef <- PtrAdd <- [Layout]
+ *   ComposedLayout: DeRef <- PtrAdd <- [Swizzle <- OffsetAdd <- Layout]
+ *   SwizzlePtr    : [DeRef <- Swizzle <- PtrAdd] <- Layout
+ *
+ * Furthermore, for known swizzles, this pointer attempts to decay itself
+ *    to a normal-pointer with a new layout containing dynamic or static strides.
+ * This is possible by determining the subdomain of the InvolutionFn
+ *    that is identity and testing if the Layout's codomain is contained
+ *    within it.
+ */
+
+namespace mute
+{
+
+// concept SwizzleFn {
+//   MUTE_HOST_DEVICE constexpr static uint apply(uint);
+// }
+// See Swizzle<B,M,S> in swizzle.hpp for common swizzle-functions.
+
+template <class SwizzleFn, class Iterator>
+struct swizzle_ptr : iter_adaptor<Iterator,swizzle_ptr<SwizzleFn,Iterator>>
+{
+  using iterator     = Iterator;
+  using reference    = typename iterator_traits<iterator>::reference;
+  using element_type = typename iterator_traits<iterator>::element_type;
+  using value_type   = typename iterator_traits<iterator>::value_type;
+
+  using iter_adaptor<Iterator,swizzle_ptr<SwizzleFn,Iterator>>::iter_adaptor;
+
+  template <class Iter>
+  MUTE_HOST_DEVICE constexpr static
+  Iter apply_swizzle(Iter ptr) {
+    if constexpr (mute::is_gmem_v<Iter>) {
+      return {apply_swizzle<typename iterator_traits<decltype(ptr.get())>::value_type,
+                            AddressSpace::Global>(ptr.get())};
+    } else if constexpr (mute::is_smem_v<Iter>) {
+      return {apply_swizzle<typename iterator_traits<decltype(ptr.get())>::value_type,
+                            AddressSpace::Shared>(ptr.get())};
+    } else {
+      return {apply_swizzle<typename iterator_traits<decltype(ptr.get())>::value_type,
+                            AddressSpace::Generic>(ptr.get())};
+    }
+  }
+
+  template <class T, AddressSpace AS = AddressSpace::Generic>
+  MUTE_HOST_DEVICE constexpr static
+  T* apply_swizzle(T* ptr) {
+    // reinterpret_cast can't be used to cast to a different address space
+    // compiler only allows c-style pointer cast
+    return (T*)make_ptr_with_address_space<AS>(SwizzleFn::apply(reinterpret_cast<uintptr_t>(ptr)));
+  }
+
+  template <class T, AddressSpace AS = AddressSpace::Generic>
+  MUTE_HOST_DEVICE constexpr static
+  subbyte_iterator<T> apply_swizzle(subbyte_iterator<T> ptr) {
+    return {apply_swizzle<T, AS>(ptr.ptr_), ptr.idx_};
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  reference operator*() const {
+    return *apply_swizzle(this->get());
+  }
+
+  template <class Int>
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](Int const& i) const {
+    return *apply_swizzle(this->get() + i);
+  }
+};
+
+template <class T, class = void>                      // Default No-Swizzle
+struct get_swizzle { using type = Swizzle<0,4,3>; };
+template <class SwizzleFn, class P>                   // Found the SwizzleFn
+struct get_swizzle<swizzle_ptr<SwizzleFn,P>> { using type = SwizzleFn; };
+template <class T>                                    // Recurse into anything with a ::iterator
+struct get_swizzle<T, void_t<typename T::iterator>> : get_swizzle<typename T::iterator> {};
+
+template <class Iter>
+using get_swizzle_t = typename get_swizzle<Iter>::type;
+
+template <class Iterator, class SwizzleFn>
+MUTE_HOST_DEVICE constexpr
+swizzle_ptr<SwizzleFn,Iterator>
+make_swizzle_ptr(Iterator ptr, SwizzleFn) {
+  return {ptr};
+}
+
+// Swizzle-0 specialization for immediate decay
+template <class Iterator, int M, int S>
+MUTE_HOST_DEVICE constexpr
+Iterator
+make_swizzle_ptr(Iterator ptr, Swizzle<0,M,S>) {
+  return ptr;
+}
+
+//
+// Recast
+//
+
+template <class SwizzleFn, class P>
+MUTE_HOST_DEVICE constexpr
+auto
+raw_pointer_cast(swizzle_ptr<SwizzleFn,P> const& ptr) {
+  return raw_pointer_cast(ptr.get());
+}
+
+// SwizzleFn operates on the pointer address, so it doesn't care about the type
+template <class NewT, class SwizzleFn, class P>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_ptr(swizzle_ptr<SwizzleFn,P> const& ptr) {
+  return make_swizzle_ptr(recast_ptr<NewT>(ptr.get()), SwizzleFn{});
+}
+
+// The statically-known alignment of a swizzle pointer is the alignment of the swizzle function converted to bits
+template <class SwizzleFn, class P>
+MUTE_HOST_DEVICE constexpr
+auto
+max_alignment(swizzle_ptr<SwizzleFn,P> const&) {
+  return Int<8>{} * max_alignment(SwizzleFn{});
+}
+
+//
+// Display utilities
+//
+
+template <class SwizzleFn, class P>
+MUTE_HOST_DEVICE void print(swizzle_ptr<SwizzleFn,P> ptr)
+{
+  print(SwizzleFn{}); printf("_"); print(ptr.get());
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class SwizzleFn, class P>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, swizzle_ptr<SwizzleFn,P> ptr)
+{
+  return os << SwizzleFn{} << "_" << ptr.get();
+}
+#endif
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/stride.hpp b/vllm_musa/csrc/flash_attn/include/mute/stride.hpp
new file mode 100644
index 000000000..5ee51d55e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/stride.hpp
@@ -0,0 +1,598 @@
+/***************************************************************************************************
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_HOST_DEVICE
+#include <mute/util/type_traits.hpp>           // mute::__MUTE_REQUIRES
+#include <mute/container/tuple.hpp>            // mute::is_tuple
+#include <mute/numeric/integral_constant.hpp>  // mute::is_integral
+#include <mute/numeric/integer_sequence.hpp>   // mute::seq
+#include <mute/numeric/math.hpp>               // mute::divmod
+#include <mute/numeric/arithmetic_tuple.hpp>   // mute::basis_get
+#include <mute/algorithm/functional.hpp>       // mute::identity
+#include <mute/algorithm/tuple_algorithms.hpp> // mute::fold
+#include <mute/int_tuple.hpp>                  // mute::is_congruent
+
+namespace mute
+{
+
+/** crd2idx(c,s,d) maps a coordinate within <Shape,Stride> to an index
+ *
+ * This is computed as follows:
+ *  [coord, shape, and stride are all integers => step forward by stride]
+ * op(c, s, d)             => c * d
+ *  [coord is integer, shape and stride are tuple => divmod coord for each mode]
+ * op(c, (s,S), (d,D))     => op(c % prod(s), s, d) + op(c / prod(s), (S), (D))
+ *  [coord, shape, and stride are all tuples => consider each mode independently]
+ * op((c,C), (s,S), (d,D)) => op(c, s, d) + op((C), (S), (D))
+ */
+template <class Coord, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+crd2idx(Coord  const& coord,
+        Shape  const& shape,
+        Stride const& stride);
+
+namespace detail {
+
+template <class Coord, class Shape, class Stride, int... Is>
+MUTE_HOST_DEVICE constexpr
+auto
+crd2idx_ttt(Coord  const& coord,
+            Shape  const& shape,
+            Stride const& stride, seq<Is...>)
+{
+  return (... + crd2idx(get<Is>(coord), get<Is>(shape), get<Is>(stride)));
+}
+
+template <class CInt, class STuple, class DTuple, int I0, int... Is>
+MUTE_HOST_DEVICE constexpr
+auto
+crd2idx_itt(CInt   const& coord,
+            STuple const& shape,
+            DTuple const& stride, seq<I0,Is...>)
+{
+  if constexpr (sizeof...(Is) == 0) {  // Avoid recursion and mod on single/last iter
+    return crd2idx(coord, get<I0>(shape), get<I0>(stride));
+  } else if constexpr (is_constant<0, CInt>::value) {
+    return crd2idx(_0{}, get<I0>(shape), get<I0>(stride))
+         + (_0{} + ... + crd2idx(_0{}, get<Is>(shape), get<Is>(stride)));
+  } else {                             // General case
+    auto [div, mod] = divmod(coord, product(get<I0>(shape)));
+    return crd2idx(mod, get<I0>(shape), get<I0>(stride))
+         + crd2idx_itt(div, shape, stride, seq<Is...>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+template <class Coord, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+crd2idx(Coord  const& coord,
+        Shape  const& shape,
+        Stride const& stride)
+{
+  if constexpr (is_tuple<Coord>::value) {
+    if constexpr (is_tuple<Shape>::value) {      // tuple tuple tuple
+      static_assert(tuple_size<Coord>::value == tuple_size< Shape>::value, "Mismatched Ranks");
+      static_assert(tuple_size<Coord>::value == tuple_size<Stride>::value, "Mismatched Ranks");
+      return detail::crd2idx_ttt(coord, shape, stride, tuple_seq<Coord>{});
+    } else {                                     // tuple "int" "int"
+      static_assert(sizeof(Coord) == 0, "Invalid parameters");
+    }
+  } else {
+    if constexpr (is_tuple<Shape>::value) {      // "int" tuple tuple
+      static_assert(tuple_size<Shape>::value == tuple_size<Stride>::value, "Mismatched Ranks");
+      return detail::crd2idx_itt(coord, shape, stride, tuple_seq<Shape>{});
+    } else {                                     // "int" "int" "int"
+      return coord * stride;
+    }
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+namespace detail {
+
+template <class CTuple, class STuple, int I0, int... Is>
+MUTE_HOST_DEVICE constexpr
+auto
+crd2idx_horner(CTuple const& coord,
+               STuple const& shape, seq<I0,Is...>)
+{
+  if constexpr (sizeof...(Is) == 0) {  // No recursion on single/last iter
+    return get<I0>(coord);
+  } else {                             // General case
+    return get<I0>(coord) + get<I0>(shape) * crd2idx_horner(coord, shape, seq<Is...>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+/** crd2idx(c,s) maps a coordinate within Shape to an index
+ * via a colexicographical enumeration of coordinates in Shape.
+ * i = c0 + s0 * (c1 + s1 * (c2 + s2 * ...))
+ */
+template <class Coord, class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+crd2idx(Coord const& coord,
+        Shape const& shape)
+{
+  if constexpr (is_integral<Coord>::value) {  // Coord is already an index
+    return coord;
+  } else if constexpr (is_integral<Shape>::value) {
+    static_assert(dependent_false<Shape>, "Invalid parameters");
+  } else {                                    // Make congruent, flatten, and apply Horner's method
+    static_assert(tuple_size<Coord>::value == tuple_size<Shape>::value, "Mismatched Ranks");
+    auto flat_coord = flatten(coord);
+    auto flat_shape = flatten(product_like(shape, coord));
+    return detail::crd2idx_horner(flat_coord, flat_shape, tuple_seq<decltype(flat_shape)>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/** idx2crd(i,s,d) splits an index into a coordinate within <Shape,Stride>.
+ *
+ * This is computed as follows:
+ *  [index, shape, and stride are all integers => determine 1D coord]
+ * op(i, s, d)             => (i / d) % s
+ *  [index is integer, shape and stride are tuple => determine component for each mode]
+ * op(i, (s,S), (d,D))     => (op(i, s, d), op(i, S, D)...)
+ *  [index, shape, and stride are all tuples => consider each mode independently]
+ * op((i,I), (s,S), (d,D)) => (op(i, s, d), op((I), (S), (D)))
+ *
+ * NOTE: This only works for compact shape+stride layouts. A more general version would
+ *       apply to all surjective layouts
+ */
+template <class Index, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+idx2crd(Index  const& idx,
+        Shape  const& shape,
+        Stride const& stride)
+{
+  if constexpr (is_tuple<Index>::value) {
+    if constexpr (is_tuple<Shape>::value) {      // tuple tuple tuple
+      static_assert(tuple_size<Index>::value == tuple_size< Shape>::value, "Mismatched Ranks");
+      static_assert(tuple_size<Index>::value == tuple_size<Stride>::value, "Mismatched Ranks");
+      return transform(idx, shape, stride, [](auto const& i, auto const& s, auto const& d){ return idx2crd(i,s,d); });
+    } else {                                     // tuple "int" "int"
+      static_assert(sizeof(Index) == 0, "Invalid parameters");
+    }
+  } else {
+    if constexpr (is_tuple<Shape>::value) {
+      if constexpr (is_tuple<Stride>::value) {   // "int" tuple tuple
+        static_assert(tuple_size<Shape>::value == tuple_size<Stride>::value, "Mismatched Ranks");
+        return transform(shape, stride, [&](auto const& s, auto const& d){ return idx2crd(idx,s,d); });
+      } else {                                   // "int" tuple "int"
+        return transform(shape, compact_col_major(shape, stride), [&](auto const& s, auto const& d){ return idx2crd(idx,s,d); });
+      }
+    } else {                                     // "int" "int" "int"
+      if constexpr (is_constant<1, Shape>::value) {
+        // Skip potential stride-0 division
+        return Int<0>{};
+      } else {
+        return (idx / stride) % shape;
+      }
+    }
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/** idx2crd(i,s) splits an index into a coordinate within Shape
+ * via a colexicographical enumeration of coordinates in Shape.
+ * c0 = (idx / 1) % s0
+ * c1 = (idx / s0) % s1
+ * c2 = (idx / (s0 * s1)) % s2
+ * ...
+ */
+template <class Index, class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+idx2crd(Index const& idx,
+        Shape const& shape)
+{
+  if constexpr (is_tuple<Index>::value) {
+    if constexpr (is_tuple<Shape>::value) {      // tuple tuple
+      static_assert(tuple_size<Index>::value == tuple_size<Shape>::value, "Mismatched Ranks");
+      return transform(idx, shape, [](auto const& i, auto const& s) { return idx2crd(i,s); });
+    } else {                                     // tuple "int"
+      static_assert(sizeof(Index) == 0, "Invalid parameters");
+    }
+  } else {
+    if constexpr (is_tuple<Shape>::value) {      // "int" tuple
+      return transform_leaf(as_arithmetic_tuple(crd2idx(idx, shape, make_basis_like(shape))), identity{});
+    } else {                                     // "int" "int"
+      return idx;
+    }
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// crd2crd
+//
+
+template <class Coord, class SShape, class DShape>
+MUTE_HOST_DEVICE constexpr
+auto
+crd2crd(Coord  const& coord,
+        SShape const& src_shape,
+        DShape const& dst_shape)
+{
+  if constexpr (is_tuple<Coord>::value && is_tuple<SShape>::value && is_tuple<DShape>::value) {
+    static_assert(tuple_size<Coord>::value == tuple_size<SShape>::value, "Mismatched Ranks");
+    static_assert(tuple_size<Coord>::value == tuple_size<DShape>::value, "Mismatched Ranks");
+    return transform(coord, src_shape, dst_shape, [](auto const& c, auto const& s, auto const& d) { return crd2crd(c,s,d); });
+  } else {
+    // assert(size(src_shape) == size(dst_shape))
+    return idx2crd(crd2idx(coord, src_shape), dst_shape);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Compact Major
+//
+
+// Tags for common layouts and dispatching
+struct LayoutLeft;               // Col-major layout mapping; leftmost extent has stride 1
+using GenColMajor = LayoutLeft;  // Alias
+
+struct LayoutRight;              // Row-major layout mapping; rightmost extent has stride 1
+using GenRowMajor = LayoutRight; // Alias
+
+namespace detail {
+
+// For GCC8.5 -- Use of lambdas in unevaluated contexts. Instead use function objects.
+template <class Major>
+struct CompactLambda;
+
+// @pre is_integral<Current>
+// Return (result, current * product(shape)) to enable recurrence
+template <class Major, class Shape, class Current>
+MUTE_HOST_DEVICE constexpr
+auto
+compact(Shape   const& shape,
+        Current const& current)
+{
+  if constexpr (is_tuple<Shape>::value) { // Shape::tuple Current::int
+    using Lambda = CompactLambda<Major>;                  // Append or Prepend
+    using Seq    = typename Lambda::template seq<Shape>;  // Seq or RSeq
+    return mute::detail::fold(shape, mute::make_tuple(mute::make_tuple(), current), Lambda{}, Seq{});
+  } else {                                // Shape::int Current::int
+    if constexpr (is_constant<1, Shape>::value) {
+      return mute::make_tuple(Int<0>{}, current); // If current is dynamic, this could save a reg
+    } else {
+      return mute::make_tuple(current, current * shape);
+    }
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// For GCC8.5 -- Specialization LayoutLeft
+template <>
+struct CompactLambda<LayoutLeft>
+{
+  template <class Init, class Shape>
+  MUTE_HOST_DEVICE constexpr auto
+  operator()(Init const& init, Shape const& si) {
+    auto result = detail::compact<LayoutLeft>(si, get<1>(init));
+    return mute::make_tuple(append(get<0>(init), get<0>(result)), get<1>(result));  // Append
+  }
+
+  template <class Shape>
+  using seq = tuple_seq<Shape>;                                                     // Seq
+};
+
+// For GCC8.5 -- Specialization LayoutRight
+template <>
+struct CompactLambda<LayoutRight>
+{
+  template <class Init, class Shape>
+  MUTE_HOST_DEVICE constexpr auto
+  operator()(Init const& init, Shape const& si) {
+    auto result = detail::compact<LayoutRight>(si, get<1>(init));
+    return mute::make_tuple(prepend(get<0>(init), get<0>(result)), get<1>(result));  // Prepend
+  }
+
+  template <class Shape>
+  using seq = tuple_rseq<Shape>;                                                     // RSeq
+};
+
+} // end namespace detail
+
+template <class Major, class Shape, class Current = Int<1>,
+          __MUTE_REQUIRES(is_tuple<Shape>::value || is_integral<Shape>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+compact_major(Shape   const& shape,
+              Current const& current = {})
+{
+  if constexpr (is_tuple<Current>::value) {    // Shape::tuple Current::tuple
+    static_assert(is_tuple<Shape>::value, "Invalid parameters");
+    static_assert(tuple_size<Shape>::value == tuple_size<Current>::value, "Mismatched Ranks");
+    // Recurse to apply to the terminals of current
+    return transform(shape, current, [&](auto const& s, auto const& c){ return compact_major<Major>(s,c); });
+  } else {
+    return get<0>(detail::compact<Major>(shape, current));
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Compact Col Major
+//
+
+struct LayoutLeft {
+  template <class Shape>
+  using Apply = decltype(compact_major<LayoutLeft>(declval<Shape>()));
+};
+
+template <class Shape, class Current = Int<1>>
+MUTE_HOST_DEVICE constexpr
+auto
+compact_col_major(Shape   const& shape,
+                  Current const& current = {})
+{
+  return compact_major<LayoutLeft>(shape, current);
+}
+
+//
+// Compact Row Major
+//
+
+struct LayoutRight {
+  template <class Shape>
+  using Apply = decltype(compact_major<LayoutRight>(declval<Shape>()));
+};
+
+template <class Shape, class Current = Int<1>>
+MUTE_HOST_DEVICE constexpr
+auto
+compact_row_major(Shape   const& shape,
+                  Current const& current = {})
+{
+  return compact_major<LayoutRight>(shape, current);
+}
+
+//
+// Compact Order -- compute a compact stride based on an ordering of the modes
+//
+
+namespace detail {
+
+// @pre weakly_congruent(order, shape)
+// @pre is_congruent<RefShape, RefOrder>
+// @pre is_static<Order>
+// @pre is_static<RefOrder>
+template <class Shape, class Order, class RefShape, class RefOrder>
+MUTE_HOST_DEVICE constexpr
+auto
+compact_order(Shape const& shape, Order const& order,
+              RefShape const& ref_shape, RefOrder const& ref_order)
+{
+  if constexpr (is_tuple<Order>::value) {
+    static_assert(tuple_size<Shape>::value == tuple_size<Order>::value, "Need equal rank of shape and order");
+    return transform(shape, order, [&](auto const& s, auto const& o) { return compact_order(s, o, ref_shape, ref_order); });
+  } else {
+    // Compute the starting stride for this shape by accumulating all shapes corresponding to lesser orders
+    auto stride_start = product(transform(ref_shape, ref_order,
+                                          [&](auto const& s, auto const& o) {
+                                            return conditional_return(o < order, s, Int<1>{});
+                                          }));
+    return compact_col_major(shape, stride_start);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+template <class Shape, class Order>
+MUTE_HOST_DEVICE constexpr
+auto
+compact_order(Shape const& shape, Order const& order)
+{
+  auto ref_shape = flatten_to_tuple(product_like(shape, order));
+
+  auto flat_order = flatten_to_tuple(order);
+  // Find the largest static element of order
+  auto max_order = mute::fold(flat_order, Int<0>{}, [](auto v, auto order) {
+    if constexpr (is_constant<true, decltype(v < order)>::value) {
+      return order;
+    } else {
+      return v;
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  });
+  // Replace any dynamic elements within order with large-static elements
+  auto max_seq = make_range<max_order+1, max_order+1+rank(flat_order)>{};
+  auto ref_order = mute::transform(max_seq, flat_order, [](auto seq_v, auto order) {
+    if constexpr (is_static<decltype(order)>::value) {
+      return order;
+    } else {
+      return seq_v;
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  });
+
+  auto new_order = unflatten(ref_order, order);
+
+  return detail::compact_order(shape, new_order, ref_shape, ref_order);
+}
+
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+compact_order(Shape const& shape, GenColMajor const& major)
+{
+  return compact_major<LayoutLeft>(shape);
+}
+
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+compact_order(Shape const& shape, GenRowMajor const& major)
+{
+  return compact_major<LayoutRight>(shape);
+}
+
+//
+// Coordinate iterator
+//
+
+namespace detail {
+
+template <class Coord, class Shape, class Order>
+MUTE_HOST_DEVICE constexpr
+void
+increment(Coord& coord, Shape const& shape, Order const& order)
+{
+  ++basis_get(get<0>(order), coord);
+  mute::for_each(make_range<1, tuple_size<Order>::value>{}, [&](auto i){
+    if (basis_get(get<i-1>(order), coord) == basis_get(get<i-1>(order), shape)) {
+      basis_get(get<i-1>(order), coord) = 0;
+      ++basis_get(get<i>(order), coord);
+    }
+  });
+}
+
+/** Increment a (dynamic) coord colexicographically within a shape
+ * @pre is_congruent<Coord,Shape>::value
+ * \code
+ *   auto shape = make_shape(1,2,make_shape(2,3),3);
+ *   auto coord = repeat_like(shape, 0);
+ *
+ *   for (int i = 0; i < size(shape); ++i) {
+ *     std::cout << i << ": " << coord << std::endl;
+ *     increment(coord, shape);
+ *   }
+ * \endcode
+ */
+template <class Coord, class Shape>
+MUTE_HOST_DEVICE constexpr
+void
+increment(Coord& coord, Shape const& shape)
+{
+  increment(coord, shape, flatten_to_tuple(make_basis_like(shape)));
+}
+
+} // end namespace detail
+
+struct ForwardCoordIteratorSentinel
+{};
+
+// A forward iterator for a starting coordinate in a shape's domain, and a shape.
+// The starting coordinate may be zero but need not necessarily be.
+template <class Coord, class Shape, class Order>
+struct ForwardCoordIterator
+{
+  static_assert(is_congruent<Coord, Shape>::value);
+
+  MUTE_HOST_DEVICE constexpr
+  Coord const& operator*() const { return coord; }
+  MUTE_HOST_DEVICE constexpr
+  ForwardCoordIterator& operator++() { detail::increment(coord, shape, Order{}); return *this; }
+  // Sentinel for the end of the implied range
+  MUTE_HOST_DEVICE constexpr
+  bool operator==(ForwardCoordIteratorSentinel const&) const { return basis_get(back(Order{}), coord) == basis_get(back(Order{}), shape); }
+  MUTE_HOST_DEVICE constexpr
+  bool operator!=(ForwardCoordIteratorSentinel const&) const { return basis_get(back(Order{}), coord) != basis_get(back(Order{}), shape); }
+  // NOTE: These are expensive, avoid use
+  MUTE_HOST_DEVICE constexpr
+  bool operator==(ForwardCoordIterator const& other) const { return coord == other.coord; }
+  MUTE_HOST_DEVICE constexpr
+  bool operator!=(ForwardCoordIterator const& other) const { return coord != other.coord; }
+
+  Coord coord;
+  Shape const& shape;
+};
+
+// A forward iterator for a coordinate that starts from a provided coordinate and increments in a prescribed order
+template <class Order, class Shape, class Coord>
+MUTE_HOST_DEVICE constexpr
+auto
+make_coord_iterator(Coord const& coord, Shape const& shape)
+{
+  static_assert(is_congruent<Coord, Shape>::value);
+  static_assert(is_congruent<Order, Coord>::value);
+  static_assert(is_congruent<Order, Shape>::value);
+  auto flat_order  = flatten_to_tuple(Order{});
+  auto inv_order   = transform(make_seq<rank(flat_order)>{}, [&](auto i){ return find(flat_order, i); });
+  auto basis_order = transform_leaf(inv_order, [&](auto i) { return get<i>(flatten_to_tuple(make_basis_like(shape))); });
+  return ForwardCoordIterator<Coord,Shape,decltype(basis_order)>{coord,shape};
+}
+
+// A forward iterator for a coordinate that starts from a provided coordinate and increments colex
+template <class Shape, class Coord>
+MUTE_HOST_DEVICE constexpr
+auto
+make_coord_iterator(Coord const& coord, Shape const& shape)
+{
+  static_assert(is_congruent<Coord, Shape>::value);
+  auto basis_order = flatten_to_tuple(make_basis_like(shape));
+  return ForwardCoordIterator<Coord,Shape,decltype(basis_order)>{coord,shape};
+}
+
+// A forward iterator for a coordinate that starts from zero and increments in a prescribed order
+template <class Order, class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+make_coord_iterator(Shape const& shape)
+{
+  return make_coord_iterator<Order>(repeat_like(shape, int(0)), shape);
+}
+
+// A forward iterator for a coordinate that starts from zero and increments colex
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+make_coord_iterator(Shape const& shape)
+{
+  return make_coord_iterator(repeat_like(shape, int(0)), shape);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/swizzle.hpp b/vllm_musa/csrc/flash_attn/include/mute/swizzle.hpp
new file mode 100644
index 000000000..d3b233911
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/swizzle.hpp
@@ -0,0 +1,474 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                      // MUTE_HOST_DEVICE
+#include <mute/container/tuple.hpp>             // mute::is_tuple
+#include <mute/numeric/integral_constant.hpp>   // mute::constant
+#include <mute/numeric/math.hpp>                // mute::max, mute::min
+#include <mute/algorithm/tuple_algorithms.hpp>  // mute::transform_apply
+
+namespace mute
+{
+
+// A generic Swizzle functor
+/* 0bxxxxxxxxxxxxxxxYYYxxxxxxxZZZxxxx
+ *                               ^--^ MBase is the number of least-sig bits to keep constant
+ *                  ^-^       ^-^     BBits is the number of bits in the mask
+ *                    ^---------^     SShift is the distance to shift the YYY mask
+ *                                       (pos shifts YYY to the right, neg shifts YYY to the left)
+ *
+ * e.g. Given
+ * 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxZZxxx
+ * the result is
+ * 0bxxxxxxxxxxxxxxxxYYxxxxxxxxxAAxxx where AA = ZZ xor YY
+ */
+template <int BBits, int MBase, int SShift = BBits>
+struct Swizzle
+{
+  static constexpr int num_bits = BBits;
+  static constexpr int num_base = MBase;
+  static constexpr int num_shft = SShift;
+
+  static_assert(num_base >= 0,             "MBase must be positive.");
+  static_assert(num_bits >= 0,             "BBits must be positive.");
+  static_assert(abs(num_shft) >= num_bits, "abs(SShift) must be more than BBits.");
+
+  // using 'int' type here to avoid unintentially casting to unsigned... unsure.
+  using bit_msk = mute::constant<int, (1 << num_bits) - 1>;
+  using yyy_msk = mute::constant<int, bit_msk{} << (num_base + max(0,num_shft))>;
+  using zzz_msk = mute::constant<int, bit_msk{} << (num_base - min(0,num_shft))>;
+  using msk_sft = mute::constant<int, num_shft>;
+
+  static constexpr uint32_t swizzle_code = uint32_t(yyy_msk{} | zzz_msk{});
+
+  template <class Offset>
+  MUTE_HOST_DEVICE constexpr static
+  auto
+  apply(Offset const& offset)
+  {
+    return offset ^ shiftr(offset & yyy_msk{}, msk_sft{});   // ZZZ ^= YYY
+  }
+
+  template <class Offset>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator()(Offset const& offset) const
+  {
+    return apply(offset);
+  }
+
+  template <int B, int M, int S>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator==(Swizzle<B,M,S> const&) const
+  {
+    return B == BBits && M == MBase && S == SShift;
+  }
+};
+
+//
+// make_swizzle<0b1000, 0b0100>()         ->  Swizzle<1,2,1>
+// make_swizzle<0b11000000, 0b00000110>() ->  Swizzle<2,1,5>
+//
+
+template <uint32_t Y, uint32_t Z>
+MUTE_HOST_DEVICE constexpr
+auto
+make_swizzle()
+{
+  constexpr uint32_t BZ = popcount(Y);                    // Number of swizzle bits
+  constexpr uint32_t BY = popcount(Z);                    // Number of swizzle bits
+  static_assert(BZ == BY, "Number of bits in Y and Z don't match");
+  constexpr uint32_t TZ_Y = countr_zero(Y);               // Number of trailing zeros in Y
+  constexpr uint32_t TZ_Z = countr_zero(Z);               // Number of trailing zeros in Z
+  constexpr uint32_t M = mute::min(TZ_Y, TZ_Z) % 32;
+  constexpr  int32_t S = int32_t(TZ_Y) - int32_t(TZ_Z);   // Difference in trailing zeros
+  static_assert((Y | Z) == Swizzle<BZ,M,S>::swizzle_code, "Something went wrong.");
+  return Swizzle<BZ,M,S>{};
+}
+
+template <int B0, int M0, int S0,
+          int B1, int M1, int S1>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(Swizzle<B0,M0,S0>, Swizzle<B1,M1,S1>)
+{
+  static_assert(S0 == S1, "Can only merge swizzles of the same shift.");
+  constexpr uint32_t Y = Swizzle<B0,M0,S0>::yyy_msk::value ^ Swizzle<B1,M1,S1>::yyy_msk::value;
+  constexpr uint32_t Z = Swizzle<B0,M0,S0>::zzz_msk::value ^ Swizzle<B1,M1,S1>::zzz_msk::value;
+  return make_swizzle<Y,Z>();
+
+  //return ComposedFn<Swizzle<B0,M0,S0>, Swizzle<B1,M1,S1>>{};
+}
+
+//
+// Utility for slicing and swizzle "offsets"
+//
+
+// For swizzle functions, it is often needed to keep track of which bits are
+//   consumed and which bits are free. Furthermore, it is useful to know whether
+// each of these bits is known statically or dynamically.
+
+// MixedBits is an 32-bit unsigned integer class where some bits are known statically
+//   and some bits are known dynamically. These sets of bits are disjoint and it is
+//   known statically which bits are known dynamically.
+
+// MixedBits can only be manipulated through bitwise operations
+
+// Abstract value:  StaticInt | (dynamic_int_ & StaticFlags)
+template <uint32_t StaticInt,
+          uint32_t StaticFlags>    // 0: static, 1: dynamic
+struct MixedBits
+{
+  // Representation invariants
+  static_assert(StaticFlags != 0, "Should be at least one dynamic bit in MixedBits.");
+  static_assert((StaticInt & StaticFlags) == 0, "No static/dynamic overlap allowed in MixedBits.");
+
+  uint32_t dynamic_int_;
+  // assert((dynamic_int_ & ~StaticFlags) == 0);
+
+  MUTE_HOST_DEVICE constexpr operator uint32_t() const noexcept { return StaticInt | dynamic_int_; }
+};
+
+// Return a value representing (C<s>{} | (d & C<f>)) potentially using MixedBits to track s and f.
+// This maker does allow ((s & f) != 0) and enforces the MixedBits invariant before creation.
+template <auto s, class DynamicType, auto f>
+MUTE_HOST_DEVICE constexpr
+auto
+make_mixed_bits(C<s>, DynamicType const& d, C<f>)
+{
+  static_assert(is_integral<DynamicType>::value);
+  constexpr uint32_t new_f = uint32_t(f) & ~uint32_t(s);        // StaticBits take precedence, M<0,f>{d} | C<s>{}
+  if constexpr (new_f == 0 || is_static<DynamicType>::value) {
+    return C<s>{} | (d & C<new_f>{});                           // Just return a static int
+  } else {
+    return MixedBits<s, new_f>{uint32_t(d) & new_f};            // MixedBits
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Operators
+//
+
+// Equality
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator==(MixedBits<S0,F0> const& m, C<S1>)
+{
+  return (S0 == (uint32_t(S1) & ~F0)) && (m.dynamic_int_ == (uint32_t(S1) & F0));
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator==(C<S1> s, MixedBits<S0,F0> const& m)
+{
+  return m == s;
+}
+
+// Bitwise AND
+template <uint32_t S0, uint32_t F0,
+          uint32_t S1, uint32_t F1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator&(MixedBits<S0,F0> const& m0, MixedBits<S1,F1> const& m1)
+{
+  // Truth table for (S0,D0,F0) & (S1,D1,F1) -> (S,D,F)
+  //   S0D0F0  | 0X0 | 001 | 011 | 1X0 |
+  // S1D1F1
+  //  0X0      | 0X0 | 0X0 | 0X0 | 0X0 |
+  //  001      | 0X0 | 001 | 001 | 001 |
+  //  011      | 0X0 | 001 | 011 | 011 |
+  //  1X0      | 0X0 | 001 | 011 | 1X0 |
+
+  return make_mixed_bits(C<S0 & S1>{},
+                         //(S0 | m0.dynamic_int_) & (S1 | m1.dynamic_int_),
+                         ((S1 & F0) & m0.dynamic_int_) | ((S0 & F1) & m1.dynamic_int_) | (m0.dynamic_int_ & m1.dynamic_int_),
+                         C<(S1 & F0) | (S0 & F1) | (F0 & F1)>{});
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator&(MixedBits<S0,F0> const& m, C<S1>)
+{
+  return make_mixed_bits(C<S0 & uint32_t(S1)>{},
+                         m.dynamic_int_,
+                         C<F0 & uint32_t(S1)>{});
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator&(C<S1> s, MixedBits<S0,F0> const& m)
+{
+  return m & s;
+}
+
+// Bitwise OR
+template <uint32_t S0, uint32_t F0,
+          uint32_t S1, uint32_t F1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator|(MixedBits<S0,F0> const& m0, MixedBits<S1,F1> const& m1)
+{
+  // Truth table for (S0,D0,F0) | (S1,D1,F1) -> (S,D,F)
+  //   S0D0F0 | 0X0 | 001 | 011 | 1X0 |
+  // S1D1F1
+  //  0X0     | 0X0 | 001 | 011 | 1X0 |
+  //  001     | 001 | 001 | 011 | 1X0 |
+  //  011     | 011 | 011 | 011 | 1X0 |
+  //  1X0     | 1X0 | 1X0 | 1X0 | 1X0 |
+
+  return make_mixed_bits(C<S0 | S1>{},
+                         ((~S1 & F0) & m0.dynamic_int_) | ((~S0 & F1) & m1.dynamic_int_),
+                         C<(~S0 & F1) | (~S1 & F0)>{});
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator|(MixedBits<S0,F0> const& m, C<S1>)
+{
+  return make_mixed_bits(C<S0 |  uint32_t(S1)>{},
+                         m.dynamic_int_,
+                         C<F0 & ~uint32_t(S1)>{});
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator|(C<S1> s, MixedBits<S0,F0> const& m)
+{
+  return m | s;
+}
+
+// Bitwise XOR
+template <uint32_t S0, uint32_t F0,
+          uint32_t S1, uint32_t F1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator^(MixedBits<S0,F0> const& m0, MixedBits<S1,F1> const& m1)
+{
+  // Truth table for (S0,D0,F0) ^ (S1,D1,F1) -> (S,D,F)
+  //   S0D0F0 | 0X0 | 001 | 011 | 1X0 |
+  // S1D1F1
+  //  0X0     | 0X0 | 001 | 011 | 1X0 |
+  //  001     | 001 | 001 | 011 | 011 |
+  //  011     | 011 | 011 | 001 | 001 |
+  //  1X0     | 1X0 | 011 | 001 | 0X0 |
+
+  return make_mixed_bits(C<(~S0 & S1 & ~F0) | (S0 & ~S1 & ~F1)>{},
+                         (S0 | m0.dynamic_int_) ^ (S1 | m1.dynamic_int_),
+                         C<F0 | F1>{});
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator^(MixedBits<S0,F0> const& m, C<S1>)
+{
+  return make_mixed_bits(C<(~S0 & uint32_t(S1) & ~F0) | (S0 & ~uint32_t(S1))>{},
+                         (S0 | m.dynamic_int_) ^ uint32_t(S1),
+                         C<F0>{});
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator^(C<S1> s, MixedBits<S0,F0> const& m)
+{
+  return m ^ s;
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator<<(MixedBits<S0,F0> const& m, C<S1>)
+{
+  return make_mixed_bits(C<(S0 << S1)>{},
+                         m.dynamic_int_ << S1,
+                         C<(F0 << S1)>{});
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+operator>>(MixedBits<S0,F0> const& m, C<S1>)
+{
+  return make_mixed_bits(C<(S0 >> S1)>{},
+                         m.dynamic_int_ >> S1,
+                         C<(F0 >> S1)>{});
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+shiftl(MixedBits<S0,F0> const& m, C<S1> s)
+{
+  if constexpr (S1 >= 0) {
+    return m << s;
+  } else {
+    return m >> -s;
+  }
+}
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+shiftr(MixedBits<S0,F0> const& m, C<S1> s)
+{
+  if constexpr (S1 >= 0) {
+    return m >> s;
+  } else {
+    return m << -s;
+  }
+}
+
+//
+// upcast and downcast
+//
+
+template <uint32_t S0, uint32_t F0, auto S1>
+MUTE_HOST_DEVICE constexpr
+auto
+safe_div(MixedBits<S0,F0> const& m, C<S1> s)
+{
+  static_assert(has_single_bit(uint32_t(S1)), "Only divide MixedBits by powers of two.");
+  return make_mixed_bits(safe_div(C<S0>{}, s),
+                         safe_div(m.dynamic_int_, s),
+                         safe_div(C<F0>{}, s));
+}
+
+template <uint32_t N, uint32_t S0, uint32_t F0>
+MUTE_HOST_DEVICE constexpr
+auto
+upcast(MixedBits<S0,F0> const& m)
+{
+  static_assert(has_single_bit(N), "Only divide MixedBits by powers of two.");
+  return safe_div(m, C<N>{});
+}
+
+template <uint32_t N, class T, __MUTE_REQUIRES(mute::is_integral<T>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+upcast(T const& m)
+{
+  return safe_div(m, C<N>{});
+}
+
+template <uint32_t N, uint32_t S0, uint32_t F0>
+MUTE_HOST_DEVICE constexpr
+auto
+downcast(MixedBits<S0,F0> const& m)
+{
+  static_assert(has_single_bit(N), "Only scale MixedBits by powers of two.");
+  return make_mixed_bits(C<S0 * N>{},
+                         m.dynamic_int_ * N,
+                         C<F0 * N>{});
+}
+
+template <uint32_t N, class T, __MUTE_REQUIRES(mute::is_integral<T>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+downcast(T const& m)
+{
+  return m * C<N>{};
+}
+
+//
+// Convert a Pow2Layout+Coord to a MixedBits
+//
+
+template <class Shape, class Stride, class Coord>
+MUTE_HOST_DEVICE constexpr
+auto
+to_mixed_bits(Shape const& shape, Stride const& stride, Coord const& coord)
+{
+  if constexpr (is_tuple<Shape>::value && is_tuple<Stride>::value && is_tuple<Coord>::value) {
+    static_assert(tuple_size<Shape>::value == tuple_size<Stride>::value, "Mismatched ranks");
+    static_assert(tuple_size<Shape>::value == tuple_size<Coord >::value, "Mismatched ranks");
+    return transform_apply(shape, stride, coord, [](auto const& s, auto const& d, auto const& c) { return to_mixed_bits(s,d,c); },
+                                                 [](auto const&... a) { return (a ^ ...); });
+  } else if constexpr (is_integral<Shape>::value && is_integral<Stride>::value && is_integral<Coord>::value) {
+    static_assert(decltype(shape*stride)::value == 0 || has_single_bit(decltype(shape*stride)::value), "Requires pow2 shape*stride.");
+    return make_mixed_bits(Int<0>{}, coord * stride, (shape - Int<1>{}) * stride);
+  } else {
+    static_assert(is_integral<Shape>::value && is_integral<Stride>::value && is_integral<Coord>::value, "Either Shape, Stride, and Coord must be all tuples, or they must be all integral (in the sense of mute::is_integral).");
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <class Layout, class Coord>
+MUTE_HOST_DEVICE constexpr
+auto
+to_mixed_bits(Layout const& layout, Coord const& coord)
+{
+  return to_mixed_bits(layout.shape(), layout.stride(), idx2crd(coord, layout.shape()));
+}
+
+//
+// Display utilities
+//
+
+template <int B, int M, int S>
+MUTE_HOST_DEVICE void print(Swizzle<B,M,S> const&)
+{
+  printf("Sw<%d,%d,%d>", B, M, S);
+}
+
+template <uint32_t S, uint32_t F>
+MUTE_HOST_DEVICE void print(MixedBits<S,F> const& m)
+{
+  printf("M_%u|(%u&%u)=%u", S, m.dynamic_int_, F, uint32_t(m));
+}
+
+#if !defined(__MUSACC_RTC__)
+template <int B, int M, int S>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, Swizzle<B,M,S> const&)
+{
+  return os << "Sw<" << B << "," << M << "," << S << ">";
+}
+
+template <uint32_t S, class D, uint32_t F>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, MixedBits<S,F> const& m)
+{
+  return os << "M_" << S << "|(" << m.dynamic_int_ << "&" << F << ")=" << uint32_t(m);
+}
+#endif // !defined(__MUSACC_RTC__)
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/swizzle_layout.hpp b/vllm_musa/csrc/flash_attn/include/mute/swizzle_layout.hpp
new file mode 100644
index 000000000..aa1c79430
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/swizzle_layout.hpp
@@ -0,0 +1,579 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>           // MUTE_HOST_DEVICE
+#include <mute/layout.hpp>           // mute::Layout
+#include <mute/layout_composed.hpp>  // mute::ComposedLayout
+#include <mute/swizzle.hpp>          // mute::Swizzle, mute::get_swizzle primary template
+
+/* Specialized functionality for a ComposedLayout of the form
+ *   InvolutionFn o Offset o LayoutB
+ * where the InvolutionFn is a Swizzle<B,M,S> and is not linear (hence the need for the Offset).
+ *
+ * Because these are specializations for core functions of ComposedLayout, these Swizzle Layouts
+ * provide similar functionality to Layout including tiling, partitioning,
+ * coordinate-to-index mapping and layout manipulations, but are not considered "normal" layouts.
+ * For example, these provide shape() and size() functions, but do not provide stride() functions.
+ *
+ * Furthermore, each of these specializations uses Swizzle<>-specific knowledge in its implementation and
+ * attempts to decay itself to a normal-layout with dynamic or static strides when certain slicing conditions
+ * are met. This is possible by determining the subdomain of the Swizzle<> function that is identity and
+ * testing if LayoutB's codomain is contained within it. In general, MizedBits is used as the Offset to track
+ * statically-vs-dynamically known bits in the Offset to improve the decay to static or dynamic normal layouts.
+ */
+
+namespace mute
+{
+
+//
+// Constructors
+//
+
+template <int B, int M, int S>
+MUTE_HOST_DEVICE constexpr
+auto
+make_layout(Swizzle<B,M,S> const& sxor)
+{
+  return composition(sxor, Layout<Int<M+B+abs(S)>,Int<1>>{});
+}
+
+namespace detail {
+
+template <int B, int M, int S, class OldShape, class OldStride, class NewShape, class NewStride>
+MUTE_HOST_DEVICE constexpr
+auto
+transfer_swizzle(Layout<OldShape,OldStride> const& old_layout,
+                 Layout<NewShape,NewStride> const& new_layout)
+{
+  // Our goal is to determine a new swizzle for the strides in new_layout for consistent vectorizations
+
+  // This is accomplished by identifying
+  //  S o L  :=:  S? o L*
+  // We identify the "active" portion of S by computing (P o L)(c*) where P is a projection generated by S
+  // Then that active identifier is transformed through the layouts:
+  //  L*(L[(P o L)(c*)])
+  // which is a new swizzle identifier for S?, the new swizzle
+
+  // Projections of the swizzle layout for composition, P
+  auto swizzle_only_zy = make_layout(make_shape (Int<(1 << M)>{}, Int<(1 << B)>{}, Int<(1 << (abs(S)-B))>{}, Int<(1 <<  B        )>{}, Int<1>{}),
+                                     make_stride(       Int<0>{}, Int<(1 << M)>{},                 Int<0>{}, Int<(1 << (M+abs(S)))>{}, Int<0>{}));
+
+  // Compose with the tile to get the swizzle projection, P o L  [The Z and Y contributing portions of L]
+  auto layout_only_zy       = composition(swizzle_only_zy, old_layout);
+  // Transform the end coordinate to get the active bits of the swizzle, (P o L)(c*)
+  auto swizzle_active_bits  = layout_only_zy(size(layout_only_zy)-Int<1>{});
+
+  // Get the Z bit and the Y bits -- keep only those that are active in Z *and* Y
+  auto zzz_msk = typename Swizzle<B,M,S>::zzz_msk{};
+  auto yyy_msk = typename Swizzle<B,M,S>::yyy_msk{};
+  auto msk_sft = typename Swizzle<B,M,S>::msk_sft{};
+  auto active_Z = swizzle_active_bits & shiftr(swizzle_active_bits,  msk_sft) & zzz_msk;
+  auto active_Y = swizzle_active_bits & shiftr(swizzle_active_bits, -msk_sft) & yyy_msk;
+
+  // Pass the identifiers through the old layout and new layout to make a new swizzle identifier, L*(L[(P o L)(c*)])
+  auto new_active_Z = new_layout(old_layout.get_1d_coord(active_Z));
+  auto new_active_Y = new_layout(old_layout.get_1d_coord(active_Y));
+
+  // Use this new swizzle identifier to construct the new swizzle for new_layout
+  //   (this also makes sure it's a "valid" swizzle that Swizzle can represent)
+  return composition(make_swizzle<new_active_Y,new_active_Z>(), new_layout);
+}
+
+} // end namespace detail
+
+template <int B, int M, int S, class Offset, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_fragment_like(ComposedLayout<Swizzle<B,M,S>,Offset,Layout> const& layout)
+{
+  return make_fragment_like(layout.layout_b());
+}
+
+//
+// Utilities
+//
+
+namespace detail {
+
+// Get just the Swizzle part of a composed layout.
+template <int B, int M, int S, class Offset, class LayoutB>
+MUTE_HOST_DEVICE constexpr
+auto
+get_swizzle_portion(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB>)
+{
+  return Swizzle<B,M,S>{};
+}
+
+// A non-swizzled layout's "Swizzle part" is the identity swizzle.
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+get_swizzle_portion(Layout<Shape,Stride>)
+{
+  return Swizzle<0,4,3>{};
+}
+
+// Get the "non-swizzle" part of a composed layout,
+// which is the underlying (non-composed) Layout.
+template <int B, int M, int S, class Offset, class LayoutB>
+MUTE_HOST_DEVICE constexpr
+auto
+get_nonswizzle_portion(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& slayout)
+{
+  return slayout.layout_b();
+}
+
+// The non-swizzle part of a non-swizzled layout is just the Layout.
+template <class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+get_nonswizzle_portion(Layout<Shape,Stride> const& slayout)
+{
+  return slayout;
+}
+
+} // namespace detail
+
+//
+// Slice a Swizzled ComposedLayout
+//
+
+namespace detail {
+
+template <class IntZ, class IntY, class Offset, int... I>
+MUTE_HOST_DEVICE constexpr
+auto
+make_swizzle_strides(true_type,
+                     IntZ   const& Z,
+                     IntY   const& Y,
+                     Offset const& offset,
+                     int_sequence<I...>)
+{
+  // Below is an optimized/compressed version of:
+  //return mute::make_tuple((swizzle(offset + Z*Int<(1 << I)>{}) - swizzle(offset))...);
+  // with knowledge of Swizzle, I... ranges for each B bits,
+  //    and the layout won't slice along z-bits that are already set
+
+  // y\z  0   1
+  //   0  Z  DC
+  //   1 -Z  DC
+
+  return mute::make_tuple(conditional_return((offset & (Y << Int<I>{})) == Int<0>{}, Z * Int<(1 << I)>{}, -Z * Int<(1 << I)>{})...);
+}
+
+template <class IntZ, class IntY, class Offset, int... I>
+MUTE_HOST_DEVICE constexpr
+auto
+make_swizzle_strides(false_type,
+                     IntZ   const& Z,
+                     IntY   const& Y,
+                     Offset const& offset,
+                     int_sequence<I...>)
+{
+  // Below is an optimized/compressed version of:
+  //return mute::make_tuple((swizzle(offset + Y*Int<(1 << I)>{}) - swizzle(offset))...);
+  // with knowledge of Swizzle, I... ranges for each B bits,
+  //    and the layout won't slice along y-bits that are already set
+
+  // y\z  0   1
+  //   0 Y+Z Y-Z
+  //   1 DC  DC
+
+  return mute::make_tuple(conditional_return((offset & (Z << Int<I>{})) == Int<0>{}, (Y+Z) * Int<(1 << I)>{}, (Y-Z) * Int<(1 << I)>{})...);
+}
+
+} // end namespace detail
+
+template <class Coord, int B, int M, int S, class Offset, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+slice_and_offset(Coord const& coord, ComposedLayout<Swizzle<B,M,S>,Offset,Layout> const& layout)
+{
+  if constexpr (all_underscore<Coord>::value) {
+    // Skip the expensive/complicated attempt to decay to a normal layout and just reshape
+    return mute::make_tuple(composition(layout.layout_a(), layout.offset(), slice(coord, layout.layout_b())), Int<0>{});
+  } else {
+
+    // Projections of the swizzle layout for composition
+    auto sw = make_layout(make_shape(Int<(1 << M)>{}, Int<(1 << B)>{}, Int<(1 << (abs(S)-B))>{}, Int<(1 << B)>{}, Int<1>{}));
+
+    auto swizzle_anti_zy = make_layout(shape(sw),
+                                       make_stride(stride<0>(sw),      Int<0>{}, stride<2>(sw),      Int<0>{}, size(sw)));
+    auto swizzle_only_zy = make_layout(shape(sw),
+                                       make_stride(     Int<0>{}, stride<1>(sw),      Int<0>{}, stride<3>(sw), Int<0>{}));
+
+    // The portion of the layout that is not yet consumed
+    auto sliced_layout = slice(coord, layout.layout_b());
+
+    // The portion of the layout that we are consuming now
+    auto diced_layout = dice(coord, layout.layout_b());
+    auto diced_coord  = dice(coord, coord);
+
+    auto diced_layout_anti_zy = composition(swizzle_anti_zy, diced_layout);
+    auto diced_layout_only_zy = composition(swizzle_only_zy, diced_layout);
+
+    // New swizzle and offset
+    auto swizzle = layout.layout_a();
+    // offset_only_zy interacts with swizzle and gets accumulated with layout.offset()
+    //   being careful about the static/dynamic contributions from diced_layout and diced_coord
+    auto offset_only_zy = layout.offset() ^ to_mixed_bits(diced_layout_only_zy, diced_coord);
+    // offset_anti_zy always gets passed through, no interaction with swizzle
+    auto offset_anti_zy = diced_layout_anti_zy(diced_coord);
+
+    // If Layout's codomain hits on         Y AND Z, then it's not reducible
+    // If Layout's codomain hits on         Y XOR Z, then it's dynamic-normal
+    // If Layout's codomain hits on neither Y NOR Z, then it's static-normal
+
+    // If the sliced_layout hits two bits that are swizzled together, then don't attempt to decay
+
+    // Compose with the layout to get the swizzle projection, P o L  [The Z and Y contributing portions of L]
+    //   (this also tests that shape/stride of layout compose with swizzle)
+    auto sliced_layout_only_zy = composition(swizzle_only_zy, sliced_layout);
+    // Transform the end coordinate to get the active bits of the swizzle, (P o L)(c*)
+    [[maybe_unused]] auto swizzle_active_bits = sliced_layout_only_zy(size(sliced_layout_only_zy)-Int<1>{});
+
+    // Determine if any active bits collide under the swizzle for potential decay
+    if constexpr (is_constant<0, decltype(not (swizzle_active_bits & ~swizzle(swizzle_active_bits)))>::value)
+    { // Hits on Y AND Z, so it's not reducible
+      return mute::make_tuple(composition(swizzle, offset_only_zy, sliced_layout), offset_anti_zy);
+    } else
+    { // Misses on Y or Z, so it's static-normal or dynamic-normal
+
+      // Lowest bit of the Z and Y masks
+      auto Z = typename Swizzle<B,M,S>::zzz_msk{} & -typename Swizzle<B,M,S>::zzz_msk{};
+      auto Y = typename Swizzle<B,M,S>::yyy_msk{} & -typename Swizzle<B,M,S>::yyy_msk{};
+      auto stride_lo = detail::make_swizzle_strides(Z < Y, Z, Y, offset_only_zy, make_int_sequence<B>{});
+      auto stride_hi = detail::make_swizzle_strides(Z > Y, Z, Y, offset_only_zy, make_int_sequence<B>{});
+
+      // Construct a (dynamic) layout that we can perform the composition with
+      auto swizzle_layout = make_layout(make_shape (Int<(1 << M)>{}, repeat<B>(Int<2>{}), Int<(1 << (abs(S)-B))>{}, repeat<B>(Int<2>{}), Int<                  1>{}),
+                                        make_stride(Int<       1>{},           stride_lo, Int<(1 <<      (M+B))>{},          stride_hi , Int<(1 << (M+B+abs(S)))>{}));
+
+      // Decay to a normal layout with offset
+      return mute::make_tuple(composition(swizzle_layout, sliced_layout),
+                              swizzle(offset_only_zy) + offset_anti_zy);
+    }
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// composition
+//
+
+// Ignore identity case
+template <int M, int S,
+          class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(Swizzle<0,M,S> const&,
+            Int<0> const&,
+            Layout<Shape,Stride> const& layout)
+{
+  return layout;
+}
+
+template <int B, int M, int S,
+          class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(Swizzle<B,M,S> const& sxor,
+            Layout<Shape,Stride> const& layout)
+{
+  return composition(sxor, Int<0>{}, layout);
+}
+
+template <class ShapeA, class StrideA,
+          int B, int M, int S>
+MUTE_HOST_DEVICE constexpr
+auto
+composition(Layout<ShapeA,StrideA> const& a,
+            Swizzle<B,M,S>         const& b)
+{
+  // Get the Z bits and the Y bits
+  auto active_Y = a(typename Swizzle<B,M,S>::yyy_msk{});
+  auto active_Z = a(typename Swizzle<B,M,S>::zzz_msk{});
+
+  // Works in simple cases... but could be greatly generalized
+
+  return composition(make_swizzle<active_Y,active_Z>(), a);
+}
+
+//
+// inverse
+//
+
+// Specialization to attempt to pass-through the Swizzle back to the left -- Needed?
+template <int B, int M, int S, class Offset, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+right_inverse(ComposedLayout<Swizzle<B,M,S>,Offset,Layout> const& layout)
+{
+  if constexpr (is_constant<0, Offset>::value) {
+    return composition(right_inverse(layout.layout_b()), layout.layout_a());
+  } else {
+    return composition(right_inverse(layout.layout_b()), right_inverse(layout.offset()), right_inverse(layout.layout_a()));
+  }
+}
+
+// Specialization to attempt to pass-through the Swizzle back to the left -- Needed?
+template <int B, int M, int S, class Offset, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+left_inverse(ComposedLayout<Swizzle<B,M,S>,Offset,Layout> const& layout)
+{
+  if constexpr (is_constant<0, Offset>::value) {
+    return composition(left_inverse(layout.layout_b()), layout.layout_a());
+  } else {
+    return composition(left_inverse(layout.layout_b()), left_inverse(layout.offset()), left_inverse(layout.layout_a()));
+  }
+}
+
+template <int B, int M, int S>
+MUTE_HOST_DEVICE constexpr
+Swizzle<B,M,S>
+right_inverse(Swizzle<B,M,S> const& sw)
+{
+  return sw;
+}
+
+template <int B, int M, int S>
+MUTE_HOST_DEVICE constexpr
+Swizzle<B,M,S>
+left_inverse(Swizzle<B,M,S> const& sw)
+{
+  return sw;
+}
+
+// Kludge -- Probably want an OffsetFn<T> here instead
+template <class T, __MUTE_REQUIRES(is_integral<T>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+right_inverse(T const& t)
+{
+  return -t;
+}
+
+// Kludge -- Probably want an OffsetFn<T> here instead
+template <class T, __MUTE_REQUIRES(is_integral<T>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+left_inverse(T const& t)
+{
+  return -t;
+}
+
+//
+// Upcast and Downcast
+//
+
+template <int N, int B, int M, int S>
+MUTE_HOST_DEVICE constexpr
+auto
+upcast(Swizzle<B,M,S> const& swizzle)
+{
+  static_assert(has_single_bit(N), "N must be a power of two");
+  constexpr int log2_n = bit_width(uint32_t(N)) - 1;
+  constexpr int NewM   = M - log2_n;
+  if constexpr (NewM >= 0) {
+    return Swizzle<B,NewM,S>{};
+  } else {
+    return Swizzle<mute::max(B+NewM,0), 0, S>{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <int N, int B, int M, int S>
+MUTE_HOST_DEVICE constexpr
+auto
+downcast(Swizzle<B,M,S> const& swizzle)
+{
+  static_assert(has_single_bit(N), "N must be a power of two");
+  constexpr int log2_n = bit_width(uint32_t(N)) - 1;
+  return Swizzle<B,(M + log2_n),S>{};
+}
+
+template <class OldType, class NewType,
+          int B, int M, int S>
+MUTE_HOST_DEVICE constexpr
+auto
+recast_layout(Swizzle<B,M,S> const& swizzle)
+{
+  using scale = decltype(trait_ratio(sizeof_bits<NewType>{}, sizeof_bits<OldType>{}));
+  if constexpr (scale::num == 1 && scale::den == 1) {
+    return swizzle;
+  }
+  else if constexpr (scale::num == 1) {
+    return downcast<scale::den>(swizzle);
+  }
+  else if constexpr (scale::den == 1) {
+    return upcast<scale::num>(swizzle);
+  }
+  else {
+    return downcast<scale::den>(upcast<scale::num>(layout));
+  }
+  MUTE_GCC_UNREACHABLE;
+}
+
+template <int B, int M, int S>
+MUTE_HOST_DEVICE constexpr
+auto
+max_alignment(Swizzle<B,M,S> const&)
+{
+  return Int<(1 << M)>{};
+}
+
+template <int B, int M, int S, class Offset, class LayoutB>
+MUTE_HOST_DEVICE constexpr
+auto
+max_alignment(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& layout)
+{
+  return gcd(max_alignment(layout.layout_a()),
+             max_alignment(layout.offset()),
+             max_alignment(layout.layout_b()));
+}
+
+//
+// Other operations
+//
+
+template <int B, int M, int S, class Offset, class LayoutB, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_layout(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& a,
+                  Layout<Shape,Stride>                          const& b)
+{
+  auto common = max_common_layout(a.layout_b(), b);
+  auto base = Int<(1 << M)>{};
+  if constexpr (base < size(common)) {
+    return common.compose(base);       // Truncate common to size base
+  } else {
+    return common;
+  }
+}
+
+template <class Shape, class Stride, int B, int M, int S, class Offset, class LayoutB>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_layout(Layout<Shape,Stride>                          const& a,
+                  ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& b)
+{
+  return max_common_layout(b, a);
+}
+
+template <int B, int M, int S, class Offset, class LayoutB, class Shape, class Stride>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_vector(ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& a,
+                  Layout<Shape,Stride>                          const& b)
+{
+  // This assumes that Offset is in the YZ domain of the Swizzle...
+  return mute::min(max_common_vector(a.layout_b(), b), Int<(1 << M)>{});
+}
+
+template <class Shape, class Stride, int B, int M, int S, class Offset, class LayoutB>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_vector(Layout<Shape,Stride>                          const& a,
+                  ComposedLayout<Swizzle<B,M,S>,Offset,LayoutB> const& b)
+{
+  return max_common_vector(b, a);
+}
+
+template <int B0, int M0, int S0, class Offset0, class LayoutB0,
+          int B1, int M1, int S1, class Offset1, class LayoutB1>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_vector(ComposedLayout<Swizzle<B0,M0,S0>,Offset0,LayoutB0> const& a,
+                  ComposedLayout<Swizzle<B1,M1,S1>,Offset1,LayoutB1> const& b)
+{
+  // Typical impl is composition(a, right_inverse(b))
+  // so this is  Sw0 o B0 o rinv(Sw1 o B1) = Sw0 o B0 o rinv(B1) o Sw1
+  auto vec = max_common_vector(a.layout_b(), b.layout_b());
+
+  // This assumes that Offset is in the YZ domain of the Swizzle...
+  if constexpr (Swizzle<B0,M0,S0>{} == Swizzle<B1,M1,S1>{}) {
+    return vec;
+  } else {
+    return mute::min(vec, Int<(1 << M0)>{}, Int<(1 << M1)>{});
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+// ComposedLayout as second argument is often more difficult...
+
+template <class Shape, class Stride,
+          int B, int M, int S, class Offset, class LayoutT>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_product(Layout<Shape,Stride>                          const& layout,
+                ComposedLayout<Swizzle<B,M,S>,Offset,LayoutT> const& tiler)
+{
+  MUTE_STATIC_ASSERT_V(tiler.offset() == Int<0>{}, "Require Swizzle offset == 0.");
+  // The new layout -- if swizzle wasn't an issue, this is the result
+  //   our goal is to determine a new swizzle for these strides
+  auto new_layout = logical_product(layout, tiler.layout_b());
+
+  // This is accomplished by identifying
+  //  S o L  :=:  S? o L*
+  // We identify the "active" portion of S by computing (P o L)(c*) where P is a projection generated by S
+  // Then that active identifier is transformed through the layouts:
+  //  L*(L[(P o L)(c*)])
+  // which is a new swizzle identifier for S?, the new swizzle
+
+  // Projections of the swizzle layout for composition, P
+  auto swizzle_only_zy = make_layout(make_shape (Int<(1 << M)>{}, Int<(1 << B)>{}, Int<(1 << (abs(S)-B))>{}, Int<(1 <<  B        )>{}, Int<1>{}),
+                                     make_stride(       Int<0>{}, Int<(1 << M)>{},                 Int<0>{}, Int<(1 << (M+abs(S)))>{}, Int<0>{}));
+
+  // Compose with the tiler to get the swizzle projection, P o L  [The Z and Y contributing portions of L]
+  auto layout_only_zy       = composition(swizzle_only_zy, tiler.layout_b());
+  // Transform the end coordinate to get the active bits of the swizzle, (P o L)(c*)
+  auto swizzle_active_bits  = layout_only_zy(size(layout_only_zy)-Int<1>{});
+  // Get the Z bit and the Y bits
+  auto active_Z = swizzle_active_bits & typename Swizzle<B,M,S>::zzz_msk{};
+  auto active_Y = swizzle_active_bits & typename Swizzle<B,M,S>::yyy_msk{};
+
+  // Pass the identifiers through the old layout and new layout to make a new swizzle identifier, L*(L[(P o L)(c*)])
+  auto new_active_Z = new_layout(Int<0>{}, tiler.layout_b()[active_Z]);
+  auto new_active_Y = new_layout(Int<0>{}, tiler.layout_b()[active_Y]);
+
+  // Use this new swizzle identifier to construxt the new swizzle for new_layout
+  //   (this also makes sure it's a "valid" swizzle that Swizzle can represent)
+  return composition(make_swizzle<new_active_Y,new_active_Z>(), new_layout);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/tensor.hpp b/vllm_musa/csrc/flash_attn/include/mute/tensor.hpp
new file mode 100644
index 000000000..358cadf5c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/tensor.hpp
@@ -0,0 +1,1102 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/util/type_traits.hpp>
+#include <mute/numeric/integral_constant.hpp>
+#include <mute/numeric/integer_sequence.hpp>
+
+#include <mute/container/tuple.hpp>
+#include <mute/container/array_aligned.hpp>
+#include <mute/container/array_subbyte.hpp>
+
+#include <mute/pointer.hpp>
+#include <mute/layout.hpp>
+
+namespace mute
+{
+
+//
+// Engine -- owning or non-owning data store
+//
+
+// concept Engine {
+//   using iterator     = ;
+//   using value_type   = ;
+//   using element_type = ;
+//   using reference    = ;
+//   iterator begin();
+// };
+
+template <class T, int N>
+struct ArrayEngine
+{
+  using Storage = typename conditional<(sizeof_bits<T>::value % 8 == 0),
+                                       array_aligned<T,N>,
+                                       array_subbyte<T,N>>::type;
+  using iterator     = typename Storage::iterator;
+  using reference    = typename iterator_traits<iterator>::reference;
+  using element_type = typename iterator_traits<iterator>::element_type;
+  using value_type   = typename iterator_traits<iterator>::value_type;
+  Storage storage_;
+
+  MUTE_HOST_DEVICE constexpr auto begin() const { return storage_.begin(); }
+  MUTE_HOST_DEVICE constexpr auto begin()       { return storage_.begin(); }
+};
+
+template <class Iterator>
+struct ViewEngine
+{
+  using iterator     = Iterator;
+  using reference    = typename iterator_traits<iterator>::reference;
+  using element_type = typename iterator_traits<iterator>::element_type;
+  using value_type   = typename iterator_traits<iterator>::value_type;
+  iterator storage_;
+
+  MUTE_HOST_DEVICE constexpr iterator const& begin() const { return storage_; }
+  MUTE_HOST_DEVICE constexpr iterator      & begin()       { return storage_; }
+};
+
+template <class Iterator>
+struct ConstViewEngine
+{
+  using iterator     = Iterator;
+  using reference    = typename iterator_traits<iterator>::reference;
+  using element_type = typename iterator_traits<iterator>::element_type;
+  using value_type   = typename iterator_traits<iterator>::value_type;
+  iterator storage_;
+
+  MUTE_HOST_DEVICE constexpr iterator const& begin() const { return storage_; }
+};
+
+//
+// Tensor
+//
+
+template <class Engine, class Layout>
+struct Tensor
+{
+  using iterator     = typename Engine::iterator;
+  using value_type   = typename Engine::value_type;
+  using element_type = typename Engine::element_type;
+  using reference    = typename Engine::reference;
+
+  using engine_type  = Engine;
+  using layout_type  = Layout;
+
+  MUTE_HOST_DEVICE constexpr
+  Tensor() {}
+
+  template <class Ptr>
+  MUTE_HOST_DEVICE constexpr
+  Tensor(Ptr const& ptr, Layout const& layout)
+      : rep_(layout, ptr) {
+  }
+
+  //
+  // Accessors
+  //
+
+  static constexpr int rank  = Layout::rank;
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  tensor() const {
+    return *this;
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  layout() const {
+    return get<0>(rep_);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  engine() const {
+    return get<1>(rep_);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  engine() {
+    return get<1>(rep_);
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  data() const {
+    return engine().begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  data() {
+    return engine().begin();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  shape() const {
+    return layout().shape();
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  auto
+  size() const {
+    return mute::size(shape());
+  }
+
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  stride() const {
+    return layout().stride();
+  }
+
+  //
+  // Indexing op() and op[]
+  //
+
+  // Index into this tensor like an array by computing the offset via layout()
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  operator[](Coord const& coord) {
+    return data()[layout()(coord)];
+  }
+
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  operator[](Coord const& coord) const {
+    return data()[layout()(coord)];
+  }
+
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  operator()(Coord const& coord) {
+    if constexpr (has_underscore<Coord>::value) {
+      auto const& [sliced_layout,offset] = slice_and_offset(coord, layout());
+      return make_tensor(data() + offset, sliced_layout);
+    } else {
+      return data()[layout()(coord)];
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  operator()(Coord const& coord) const {
+    if constexpr (has_underscore<Coord>::value) {
+      auto const& [sliced_layout,offset] = slice_and_offset(coord, layout());
+      return make_tensor(data() + offset, sliced_layout);
+    } else {
+      return data()[layout()(coord)];
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+
+  // op() convenience function for multi-dimensional coordinates
+  template <class Coord0, class Coord1, class... Coords>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) {
+    return operator()(make_coord(c0,c1,cs...));
+  }
+
+  template <class Coord0, class Coord1, class... Coords>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const {
+    return operator()(make_coord(c0,c1,cs...));
+  }
+
+  //
+  // Compose
+  //
+
+  template <class... Layouts>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  compose(Layouts const&... layouts) {
+    return make_tensor(data(), layout().compose(layouts...));
+  }
+
+  template <class... Layouts>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  compose(Layouts const&... layouts) const {
+    return make_tensor(data(), layout().compose(layouts...));
+  }
+
+  //
+  // Tile
+  //
+
+  template <class... Layouts>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  tile(Layouts const&... layouts) {
+    return make_tensor(data(), layout().tile(layouts...));
+  }
+
+  template <class... Layouts>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  tile(Layouts const&... layouts) const {
+    return make_tensor(data(), layout().tile(layouts...));
+  }
+
+  //
+  // Utility
+  //
+
+  template <class Int,
+            __MUTE_REQUIRES(is_integral<Int>::value)>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_1d_coord(Int const& linear_idx) const {
+    return layout().get_1d_coord(linear_idx);
+  }
+
+  template <class Int,
+            __MUTE_REQUIRES(is_integral<Int>::value)>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_hier_coord(Int const& linear_idx) const {
+    return layout().get_hier_coord(linear_idx);
+  }
+
+  template <class Int,
+            __MUTE_REQUIRES(is_integral<Int>::value)>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  get_flat_coord(Int const& linear_idx) const {
+    return layout().get_flat_coord(linear_idx);
+  }
+
+  mute::tuple<layout_type, engine_type> rep_;
+};
+
+template <class T>
+struct is_tensor : false_type {};
+template <class Engine, class Layout>
+struct is_tensor<Tensor<Engine,Layout>> : true_type {};
+template <class T>
+constexpr bool is_tensor_v = is_tensor<T>::value;
+
+// Customization point for creation of owning and non-owning Tensors
+template <class T>
+struct MakeTensor
+{
+  template <class Layout,
+            __MUTE_REQUIRES(not has_dereference<T>::value &&
+                            is_layout<Layout>::value)>
+  MUTE_HOST_DEVICE constexpr auto
+  operator()(Layout const& layout) const
+  {
+    static_assert(is_static<Layout>::value, "Dynamic owning tensors not supported");
+    using Engine = ArrayEngine<T, cosize_v<Layout>>;
+    return Tensor<Engine,Layout>();
+  }
+
+  template <class Layout,
+            __MUTE_REQUIRES(has_dereference<T>::value &&
+                            is_layout<Layout>::value)>
+  MUTE_HOST_DEVICE constexpr auto
+  operator()(T const& iter, Layout const& layout)
+  {
+    using Engine = ViewEngine<T>;
+    return Tensor<Engine,Layout>(iter, layout);
+  }
+
+  template <class LayoutArg, class... LayoutArgs,
+            __MUTE_REQUIRES(not is_layout<LayoutArg>::value)>
+  MUTE_HOST_DEVICE constexpr auto
+  operator()(LayoutArg const& arg, LayoutArgs const&... args) const
+  {
+    return operator()(make_layout(arg, args...));
+  }
+
+  template <class LayoutArg, class... LayoutArgs,
+            __MUTE_REQUIRES(not is_layout<LayoutArg>::value)>
+  MUTE_HOST_DEVICE constexpr auto
+  operator()(T const& iter, LayoutArg const& arg, LayoutArgs const&... args)
+  {
+    return operator()(iter, make_layout(arg, args...));
+  }
+};
+
+//
+// make_tensor
+//
+
+// Make an owning Tensor that will allocate a static array
+// e.g. make_tensor<float>(Int<12>{})
+template <class T, class... Args>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tensor(Args const&... args)
+{
+  return MakeTensor<T>{}(args...);
+}
+
+// Make a non-owning Tensor that will use a pointer (view)
+// e.g. make_tensor(vec.data(), 12)
+template <class Iterator, class... Args>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tensor(Iterator const& iter, Args const&... args)
+{
+  return MakeTensor<Iterator>{}(iter, args...);
+}
+
+//
+// make_tensor_like
+//   Make a register tensor the same type and shape and (if possible) order as another tensor
+//
+
+template <class NewT, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tensor_like(Layout const& layout)
+{
+  return make_tensor<NewT>(make_layout_like(layout));
+}
+
+template <class NewT, class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tensor_like(Tensor<Engine,Layout> const& tensor)
+{
+  return make_tensor_like<NewT>(tensor.layout());
+}
+
+template <class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tensor_like(Tensor<Engine,Layout> const& tensor)
+{
+  return make_tensor_like<typename Engine::value_type>(tensor.layout());
+}
+
+//
+// make_fragment_like --
+//   Make a tensor the same shape and (if possible) order as another tensor, with special
+//   consideration of the 0th mode. The 0th mode is commonly used for MMA_Atoms or Copy_Atoms
+//   so this allocates the 0th mode with LayoutLeft regardless of the reference layout.
+//
+
+template <class NewT, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_fragment_like(Layout const& layout)
+{
+  return make_tensor<NewT>(make_fragment_like(layout));
+}
+
+template <class NewT, class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_fragment_like(Tensor<Engine,Layout> const& tensor)
+{
+  return make_fragment_like<NewT>(tensor.layout());
+}
+
+template <class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_fragment_like(Tensor<Engine,Layout> const& tensor)
+{
+  return make_fragment_like<typename Engine::value_type>(tensor.layout());
+}
+
+//
+// make_counting_tensor
+//   Make a tensor from a layout by binding it to a counting iter with 0-offset of the same profile as the codomain.
+//
+
+template <class Layout, __MUTE_REQUIRES(is_layout<Layout>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+make_counting_tensor(Layout const& layout)
+{
+  return make_tensor(make_inttuple_iter(repeat_like(coshape(layout), Int<0>{})), layout);
+}
+
+//
+// make_identity_tensor
+//   Make a tensor that maps coordinates within a shape to themselves.
+//
+
+template <class Shape>
+MUTE_HOST_DEVICE constexpr
+auto
+make_identity_tensor(Shape const& shape)
+{
+  return make_counting_tensor(make_identity_layout(shape));
+}
+
+//
+// Utilities
+//
+
+// Return the subtensor of a mode
+template <class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+tensor(Tensor&& tensor)
+{
+  return static_cast<Tensor&&>(tensor);
+}
+
+template <int I, int... Is, class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+tensor(Tensor&& tensor)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(), get<I,Is...>(tensor.layout()));
+}
+
+// Return the layout of a mode
+template <int... Is, class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+layout(Tensor<Engine,Layout> const& tensor)
+{
+  return layout<Is...>(tensor.layout());
+}
+
+// Return the shape of a mode
+template <int... Is, class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+shape(Tensor<Engine,Layout> const& tensor)
+{
+  return shape<Is...>(tensor.layout());
+}
+
+// Return the stride of a mode
+template <int... Is, class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+stride(Tensor<Engine,Layout> const& tensor)
+{
+  return stride<Is...>(tensor.layout());
+}
+
+// Return the number of elements in a mode
+template <int... Is, class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+size(Tensor<Engine,Layout> const& tensor)
+{
+  return size<Is...>(tensor.layout());
+}
+
+// Return the rank of a mode
+template <int... Is, class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+rank(Tensor<Engine,Layout> const& tensor)
+{
+  return rank<Is...>(tensor.layout());
+}
+
+// Return the depth of a mode
+template <int... Is, class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+depth(Tensor<Engine, Layout> const& tensor)
+{
+  return depth<Is...>(tensor.layout());
+}
+
+//
+// Operations to manipulate Tensors like a Layout
+//
+
+template <class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+flatten(Tensor&& tensor)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(), flatten(tensor.layout()));
+}
+
+template <class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce(Tensor&& tensor)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(), coalesce(tensor.layout()));
+}
+
+template <class Tensor, class Profile,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+coalesce(Tensor&& tensor, Profile const& profile)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(), coalesce(tensor.layout(), profile));
+}
+
+template <class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+filter_zeros(Tensor&& tensor)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(), filter_zeros(tensor.layout()));
+}
+
+template <class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+filter(Tensor&& tensor)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(), filter(tensor.layout()));
+}
+
+template <class Tensor, class Profile,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+filter(Tensor&& tensor, Profile const& profile)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(), filter(tensor.layout(), profile));
+}
+
+// Return a tensor with the same shape as input but offset by a given coordinate
+template <class Coord, class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+domain_offset(Coord const& coord, Tensor&& tensor)
+{
+  auto [layout, ptr_offset] = domain_offset(coord, tensor.layout());
+  return make_tensor(static_cast<Tensor&&>(tensor).data() + ptr_offset, layout);
+}
+
+// Group the modes [B,E) into a single mode
+// e.g. group<2,4>(make_tensor<int>(Layout<Shape<_1,_2,_3,_4,_5,_6>>{}))
+//      => make_tensor<int>(Layout<Shape<_1,_2,Shape<_3,_4>,_5,_6>>{})
+template <int B, int E, class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+group_modes(Tensor&& tensor)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(),
+                     group<B,E>(tensor.layout()));
+}
+
+// Return the subtensor of a range of modes
+template <int B, int E, class Tensor,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+decltype(auto)
+take(Tensor&& tensor)
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(), take<B,E>(tensor.layout()));
+}
+
+//
+// Recast
+//
+
+// NOTE: This is very dangerous to do
+//   -- doesn't check dynamic integer divisibility
+//   -- doesn't check alignment
+
+template <class NewType, class Tensor>
+MUTE_HOST_DEVICE constexpr
+auto
+recast(Tensor&& tensor)
+{
+  using OldType = typename remove_cvref_t<Tensor>::value_type;
+  auto old_layout = tensor.layout();
+  auto new_layout = recast_layout<OldType,NewType>(old_layout);
+
+  // If this is an upcast of a normal Layout with static negative strides, then offset as well
+  if constexpr (sizeof(OldType) < sizeof(NewType) && not is_composed_layout<decltype(old_layout)>::value) {
+    auto shape_diff = transform(flatten(old_layout.shape()), flatten(new_layout.shape()), minus{});
+    auto extent_diff = transform(shape_diff, flatten(old_layout.stride()), multiplies{});
+    auto offset = fold(extent_diff, Int<0>{}, [](auto const& i, auto const& a) { return i + mute::min(a,Int<0>{}); });
+
+    return make_tensor(recast_ptr<NewType>(static_cast<Tensor&&>(tensor).data() + offset), new_layout);
+  } else {
+    return make_tensor(recast_ptr<NewType>(static_cast<Tensor&&>(tensor).data()         ), new_layout);
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// max_common_vector
+//
+
+/* Return Int<N> such that N is the maximum number of contiguous elements
+ * that logically correspond in the tensors of @a a and @a b. This is,
+ * the number of elements that could reasonably be vectorized into a single load/store.
+ *
+ * @returns Int<N> with N >= 0
+ *
+ * A return value of Int<0> indicates that no such conclusion can be made and no
+ * vectorization should be attempted.
+ *
+ * Note that the return value does NOT include alignment concerns such as the pointer value and
+ * the divisbility of dynamic strides.
+ */
+template <class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_vector(Tensor<SrcEngine,SrcLayout> const& a,
+                  Tensor<DstEngine,DstLayout> const& b)
+{
+  using SrcType = typename Tensor<SrcEngine,SrcLayout>::value_type;
+  using DstType = typename Tensor<DstEngine,DstLayout>::value_type;
+  using SrcRef  = typename Tensor<SrcEngine,SrcLayout>::reference;
+  using DstRef  = typename Tensor<SrcEngine,SrcLayout>::reference;
+
+  // Determine if vectorization candidates at all
+  if constexpr (// Should be the same value_types, else the copy is also performing a cast
+                sizeof_bits_v<SrcType> == sizeof_bits_v<DstType> &&
+                // The types should be trivially copyable so that vectorization is valid
+                is_trivially_copyable<SrcType>::value &&
+                is_trivially_copyable<DstType>::value &&
+                // Should be load/storing real data, rather than implicit iterators or such
+                is_reference<SrcRef>::value &&
+                is_reference<DstRef>::value)
+  {
+    return max_common_vector(a.layout(), b.layout());
+  } else {
+    return Int<0>{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+/* Return a layout that points to the maximum number of contiguous elements
+ * that logically correspond in the tensors of @a a and @a b. This is,
+ * the elements that could reasonably be "vectorized" into a single load/store.
+ *
+ * @returns Layout R such that composition(a.layout(), R) and composition(b.layout(), R)
+ *          are both identity Layouts.
+ *
+ * Note that the returned layout does NOT include alignment concerns such as the pointer value and
+ * the divisbility of dynamic strides.
+ */
+template <class SrcEngine, class SrcLayout,
+          class DstEngine, class DstLayout>
+MUTE_HOST_DEVICE constexpr
+auto
+max_common_layout(Tensor<SrcEngine,SrcLayout> const& a,
+                  Tensor<DstEngine,DstLayout> const& b)
+{
+  using SrcType = typename Tensor<SrcEngine,SrcLayout>::value_type;
+  using DstType = typename Tensor<DstEngine,DstLayout>::value_type;
+  using SrcRef  = typename Tensor<SrcEngine,SrcLayout>::reference;
+  using DstRef  = typename Tensor<SrcEngine,SrcLayout>::reference;
+
+  // Determine if vectorization candidates at all
+  if constexpr (// Should be the same value_types, else the copy is also performing a cast
+                sizeof_bits_v<SrcType> == sizeof_bits_v<DstType> &&
+                // The types should be trivially copyable so that vectorization is valid
+                is_trivially_copyable<SrcType>::value &&
+                is_trivially_copyable<DstType>::value &&
+                // Should be load/storing real data, rather than implicit iterators or such
+                is_reference<SrcRef>::value &&
+                is_reference<DstRef>::value)
+  {
+    return max_common_layout(a.layout(), b.layout());
+  } else {
+    return Layout<_1,_0>{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Key algebraic operations -- Divide and Product
+//
+
+// Apply a Tiler to the Tensor.
+//
+// Consider a Tensor with shape (A,B,x,y)
+// And a Tiler that is:
+//
+// * A Layout with shape (BLK_A,BLK_B)
+// ** Result Tensor shape ((BLK_A,BLK_B),Rest).
+// ** That is, the Tensor and Tile are treated as 1D for the tiling.
+// ** See logical_divide(Layout,Layout)
+//
+// * A Tile<Layout...> with shape <BLK_A,BLK_B>
+// ** Result Tensor shape ((BLK_A,a),(BLK_B,b),x,y).
+// ** Each mode of the Tile<Layout...> is applied to the corresponding mode of the Tensor.
+// ** See logical_divide(Layout,Tuple)
+//
+// * A Shape (BLK_A,BLK_B)
+// ** Result Tensor shape ((BLK_A,a),(BLK_B,b),x,y).
+// ** Equivalent to applying Tile<BLK_A:_1,BLK_B:_1>.
+// ** See logical_divide(Layout,Tuple) and logical_divide(Layout,Int)
+//
+// Note that the Tile<Layout...>/Shape Tilers must be weakly_congruent to the Tensor
+template <class Tensor, class Tiler,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_divide(Tensor    && tensor,
+               Tiler const& tiler)   // Layout or Tile<Layout...> or Shape
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(),
+                     logical_divide(tensor.layout(), tiler));
+}
+
+// zipped_divide is logical_divide with Tiler modes and Rest modes gathered together: (Tiler,Rest)
+// When Tiler is Layout, this has no effect as logical_divide results in the same.
+// When Tiler is Tile<Layout...> or Shape, this zips modes into standard form ((BLK_A,BLK_B),(a,b,x,y))
+template <class Tensor, class Tiler,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+zipped_divide(Tensor    && tensor,
+              Tiler const& tiler)    // Layout or Tile<Layout...> or Shape
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(),
+                     zipped_divide(tensor.layout(), tiler));
+}
+
+// tiled_divide is zipped_divide with the second output mode flattened ((BLK_A,BLK_B),a,b,x,y)
+template <class Tensor, class Tiler,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+tiled_divide(Tensor    && tensor,
+             Tiler const& tiler)     // Layout or Tile<Layout...> or Shape
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(),
+                     tiled_divide(tensor.layout(), tiler));
+}
+
+// flat_divide is zipped_divide with the both modes flattened (BLK_A,BLK_B,a,b,x,y)
+template <class Tensor, class Tiler,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+flat_divide(Tensor    && tensor,
+            Tiler const& tiler)      // Layout or Tile<Layout...> or Shape
+{
+  return make_tensor(static_cast<Tensor&&>(tensor).data(),
+                     flat_divide(tensor.layout(), tiler));
+}
+
+// logical_product on a Tensor doesn't make sense since it often increases cosize
+//   though this might make sense for creating Tensors with broadcasted (stride-0) modes
+
+//
+// Tensor partitioning utilities
+//
+
+// Apply a Tiler to the Tensor, then slice out one of those tiles by slicing into the "Rest" modes.
+// With an inner_partition, you get everything that's inside the Tiler. Everything that the Tiler is pointing to.
+// Split the modes of tensor according to the Tiler
+//   zipped_divide returns something like ((BLK_A,BLK_B,...),(a,b,...,x,y))
+// Then slice into the second mode (the "Rest" mode) with Coord
+template <class Tensor, class Tiler, class Coord,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+inner_partition(Tensor    && tensor,
+                Tiler const& tiler,
+                Coord const& coord)
+{
+  auto tensor_tiled = zipped_divide(static_cast<Tensor&&>(tensor), tiler);
+  constexpr int R0 = decltype(rank<0>(tensor_tiled))::value;
+
+  // The coord slices into the second mode (the "rest" mode), flatten the first
+  if constexpr (is_tuple<Coord>::value) {
+    // Append trailing modes if coord is tuple
+    constexpr int R1 = decltype(rank<1>(tensor_tiled))::value;;
+    return tensor_tiled(repeat<R0>(_), append<R1>(coord,_));
+  } else {
+    // Flat indexing if coord is not tuple
+    return tensor_tiled(repeat<R0>(_), coord);
+  }
+}
+
+// Apply a Tiler to the Tensor, then slice out the remainder by slicing into the "Tile" modes.
+// With an outer_partition, you get everything that's outside the Tiler. The layout of the Tile in the Tensor.
+// Split the modes of tensor according to the Tiler
+//   zipped_divide returns something like ((BLK_A,BLK_B,...),(a,b,...,x,y))
+// Then slice into the first mode (the "Tile" mode) with Coord
+template <class Tensor, class Tiler, class Coord,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+outer_partition(Tensor    && tensor,
+                Tiler const& tiler,
+                Coord const& coord)
+{
+  auto tensor_tiled = zipped_divide(static_cast<Tensor&&>(tensor), tiler);
+  constexpr int R1 = decltype(rank<1>(tensor_tiled))::value;
+
+  // The coord slices into the first mode (the "tile" mode), flatten the second
+  if constexpr (is_tuple<Coord>::value) {
+    // Append trailing modes if coord is tuple
+    constexpr int R0 = decltype(rank<0>(tensor_tiled))::value;
+    return tensor_tiled(append<R0>(coord,_), repeat<R1>(_));
+  } else {
+    // Flat indexing if coord is not tuple
+    return tensor_tiled(coord, repeat<R1>(_));
+  }
+}
+
+// Tile a tensor according to @a tiler and use @a coord to index into the remainder, keeping the tile.
+// This is typical at the CTA level where tiles of data are extracted:
+//   Tensor data = ...                                                                         // (  M,  N)
+//   Tensor cta_data = local_tile(data, Shape<_32,_64>{}, make_coord(blockIdx.x,blockIdx.y));  // (_32,_64)
+template <class Tensor, class Tiler, class Coord,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE constexpr
+auto
+local_tile(Tensor    && tensor,
+           Tiler const& tiler,   // tiler to apply
+           Coord const& coord)   // coord to slice into "remainder"
+{
+  return inner_partition(static_cast<Tensor&&>(tensor),
+                         tiler,
+                         coord);
+}
+
+// Same as above, but with a projection parameter to strip out unwanted tiling modes for convenience
+//   when using projections of the same tiler.
+// This is typical at the CTA level where tiles of data are extracted as projections:
+//   Tensor dataA = ...                                                        // (M,K)
+//   Tensor dataB = ...                                                        // (N,K)
+//   Tensor dataC = ...                                                        // (M,N)
+//   auto cta_tiler = Shape<_32, _64, _4>{};
+//   auto cta_coord = make_coord(blockIdx.x, blockIdx.y, _);
+//   Tensor ctaA = local_tile(dataA, cta_tiler, cta_coord, Step<_1, X,_1>{});  // (_32,_4,k)
+//   Tensor ctaB = local_tile(dataA, cta_tiler, cta_coord, Step< X,_1,_1>{});  // (_64,_4,k)
+//   Tensor ctaC = local_tile(dataA, cta_tiler, cta_coord, Step<_1,_1, X>{});  // (_32,_64)
+template <class Tensor, class Tiler, class Coord, class Proj,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE
+auto
+local_tile(Tensor    && tensor,
+           Tiler const& tiler,   // tiler to apply
+           Coord const& coord,   // coord to slice into "remainder"
+           Proj  const& proj)    // projection to apply to tiler and coord
+{
+  return local_tile(static_cast<Tensor&&>(tensor),
+                    dice(proj, tiler),
+                    dice(proj, coord));
+}
+
+// Tile a tensor according to the flat shape of a layout that provides the coordinate of the target index.
+// This is typical at the Thread level where data is partitioned across repeated patterns of threads:
+//   Tensor data = ...                                                            // (_16,_64)
+//   Tensor thr_data = local_partition(data, Layout<Shape<_2,_16>>{}, thr_idx);   // ( _8, _4)
+template <class Tensor, class LShape, class LStride, class Index,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE
+auto
+local_partition(Tensor                     && tensor,
+                Layout<LShape,LStride> const& tile,    // coord -> index
+                Index                  const& index)   // index to slice for
+{
+  static_assert(is_integral<Index>::value);
+  return outer_partition(static_cast<Tensor&&>(tensor),
+                         product_each(shape(tile)),
+                         tile.get_flat_coord(index));
+}
+
+// Same as above, but with a projection parameter to strip out unwanted tiling modes for convenience
+//   when using projections of the same tiler.
+// This is typical at the Thread level where data is partitioned across projected layouts of threads:
+//   Tensor dataA = ...                                                            // (M,K)
+//   Tensor dataB = ...                                                            // (N,K)
+//   Tensor dataC = ...                                                            // (M,N)
+//   auto thr_layout = Layout<Shape<_2,_16,_1>, Stride<_16,_1,_0>>{};
+//   Tensor thrA = local_partition(dataA, thr_layout, thr_idx, Step<_1, X,_1>{});  // (M/2,K/1)
+//   Tensor thrB = local_partition(dataB, thr_layout, thr_idx, Step< X,_1,_1>{});  // (N/16,K/1)
+//   Tensor thrC = local_partition(dataC, thr_layout, thr_idx, Step<_1,_1, X>{});  // (M/2,N/16)
+template <class Tensor, class LShape, class LStride, class Index, class Projection,
+          __MUTE_REQUIRES(is_tensor<remove_cvref_t<Tensor>>::value)>
+MUTE_HOST_DEVICE
+auto
+local_partition(Tensor                     && tensor,
+                Layout<LShape,LStride> const& tile,   // coord -> index
+                Index                  const& index,  // index to slice for
+                Projection             const& proj)
+{
+  return local_partition(static_cast<Tensor&&>(tensor),
+                         dice(proj, tile),
+                         index);
+}
+
+//
+// Display utilities
+//
+
+template <class Engine, class Layout>
+MUTE_HOST_DEVICE void print(Tensor<Engine,Layout> const& tensor)
+{
+  print(tensor.data()); print(" o "); print(tensor.layout());
+}
+
+template <class Engine, class Layout>
+MUTE_HOST_DEVICE void print_tensor(Tensor<Engine,Layout> const& tensor, bool print_type = true)
+{
+  if (print_type) {
+    print(tensor); print(":\n");
+  }
+
+  if constexpr (Layout::rank == 1)
+  {
+    for (int m = 0; m < size(tensor); ++m) {
+      pretty_print(tensor(m));
+      printf("\n");
+    }
+  } else
+  if constexpr (Layout::rank == 2)
+  {
+    for (int m = 0; m < size<0>(tensor); ++m) {
+      for (int n = 0; n < size<1>(tensor); ++n) {
+        pretty_print(tensor(m,n));
+      }
+      printf("\n");
+    }
+  } else
+  if constexpr (Layout::rank == 3)
+  {
+    print_tensor(tensor(_,_,0), false);
+    for (int k = 1; k < size<2>(tensor); ++k) {
+      for (int i = 0; i < 5*size<1>(tensor); ++i) { print("-"); } print("\n");
+      print_tensor(tensor(_,_,k), false);
+    }
+  } else
+  if constexpr (Layout::rank == 4)
+  {
+    print_tensor(tensor(_,_,_,0), false);
+    for (int p = 1; p < size<3>(tensor); ++p) {
+      for (int i = 0; i < 5*size<1>(tensor); ++i) { print("="); } print("\n");
+      print_tensor(tensor(_,_,_,p), false);
+    }
+  }
+}
+
+#if !defined(__MUSACC_RTC__)
+template <class Engine, class Layout>
+MUTE_HOST std::ostream& print_tensor_os(std::ostream& os, Tensor<Engine,Layout> const& tensor)
+{
+  int digits = 9;
+
+  if constexpr (Layout::rank == 1)
+  {
+    for (int m = 0; m < size(tensor); ++m) {
+      os << std::setw(digits) << tensor(m) << std::endl;
+    }
+  } else
+  if constexpr (Layout::rank == 2)
+  {
+    for (int m = 0; m < size<0>(tensor); ++m) {
+      for (int n = 0; n < size<1>(tensor); ++n) {
+        os << std::setw(digits) << tensor(m,n);
+      }
+      os << std::endl;
+    }
+  } else
+  if constexpr (Layout::rank == 3)
+  {
+    print_tensor_os(os, tensor(_,_,0));
+    for (int k = 1; k < size<2>(tensor); ++k) {
+      for (int i = 0; i < digits*size<1>(tensor); ++i) { os << "-"; } os << std::endl;
+      print_tensor_os(os, tensor(_,_,k));
+    }
+  } else
+  if constexpr (Layout::rank == 4)
+  {
+    print_tensor_os(os, tensor(_,_,_,0));
+    for (int p = 1; p < size<3>(tensor); ++p) {
+      for (int i = 0; i < digits*size<1>(tensor); ++i) { os << "="; } os << std::endl;
+      print_tensor_os(os, tensor(_,_,_,p));
+    }
+  }
+
+  return os;
+}
+
+template <class Engine, class Layout>
+MUTE_HOST std::ostream& operator<<(std::ostream& os, Tensor<Engine,Layout> const& tensor)
+{
+  os << tensor.layout() << std::endl;
+  return print_tensor_os(os, tensor);
+}
+#endif // !defined(__MUSACC_RTC__)
+
+} // end namespace mute
+
+//
+// Extended Engines
+//
+
+#include <mute/pointer_swizzle.hpp>
+#include <mute/pointer_flagged.hpp>
+#include <mute/tensor_zip.hpp>
+//
+// Tensor Algorithms
+//
+
+#include <mute/algorithm/tensor_algorithms.hpp>
+#include <mute/algorithm/fill.hpp>
+#include <mute/algorithm/clear.hpp>
+#include <mute/algorithm/copy.hpp>
+#include <mute/algorithm/prefetch.hpp>
+#include <mute/algorithm/axpby.hpp>
+#include <mute/algorithm/gemm.hpp>
+
+#include <mute/algorithm/cooperative_copy.hpp>
+#include <mute/algorithm/cooperative_gemm.hpp>
diff --git a/vllm_musa/csrc/flash_attn/include/mute/tensor_predicate.hpp b/vllm_musa/csrc/flash_attn/include/mute/tensor_predicate.hpp
new file mode 100644
index 000000000..ab5e64738
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/tensor_predicate.hpp
@@ -0,0 +1,80 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/numeric/integral_constant.hpp>
+
+namespace mute
+{
+
+template <class T>
+struct ConstantTensor
+{
+  template <class... Coords>
+  MUTE_HOST_DEVICE constexpr
+  T const&
+  operator()(Coords const&...) const {
+    return val_;
+  }
+
+  T val_;
+};
+
+struct TrivialPredTensor
+{
+  template <class... Coords>
+  MUTE_HOST_DEVICE constexpr
+  true_type
+  operator()(Coords const&...) const {
+    return {};
+  }
+};
+
+template <class Fn>
+struct FunctionPredTensor
+{
+  MUTE_HOST_DEVICE constexpr
+  FunctionPredTensor(Fn const& fn) : fn_(fn) {}
+
+  template <class... Coords>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator()(Coords const&... coords) const {
+    return fn_(coords...);
+  }
+
+  Fn const& fn_;
+};
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/tensor_zip.hpp b/vllm_musa/csrc/flash_attn/include/mute/tensor_zip.hpp
new file mode 100644
index 000000000..5f1e8697e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/tensor_zip.hpp
@@ -0,0 +1,243 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2024 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>           // MUTE_HOST_DEVICE
+#include <mute/tensor.hpp>           // mute::Tensor
+#include <mute/container/tuple.hpp>  // mute::tuple
+
+namespace mute
+{
+
+// A tuple of Iterators that can be offset asymmetrically
+// Note that this only accepts op+(tuple<Index...>) and op[tuple<Index...>]
+//   where each iterator will be offset by its respective index only.
+// READ-ONLY for now until mute::tuple can be constructed with references.
+template <class... Iters>
+struct ZipIterator
+{
+  using value_type   = mute::tuple<iter_value_t<Iters>...>;
+  using element_type = mute::tuple<iter_element_t<Iters>...>;
+  // NOTE: mute::tuple does not support constructions with references at the moment.
+  //       Consider fixes and/or an implementation of std::forward_as_tuple.
+  //       For now, use a mute::tuple of value_types instead, which makes this Iterator READ-ONLY.
+  //using reference    = mute::tuple<iter_reference_t<Iters>...>;
+  using reference  = value_type;
+
+  ZipIterator() = delete;
+
+  MUTE_HOST_DEVICE constexpr
+  ZipIterator(Iters... iters)
+    : iters_(iters...)
+  {}
+
+  MUTE_HOST_DEVICE constexpr
+  ZipIterator(mute::tuple<Iters...> const& iters)
+    : iters_(iters)
+  {}
+
+  MUTE_HOST_DEVICE constexpr
+  reference operator*() const {
+    return mute::apply(iters_, [](auto&&... args) { return reference(*args...); });
+  }
+
+  template <class... Index>
+  MUTE_HOST_DEVICE constexpr
+  ZipIterator operator+(mute::tuple<Index...> const& idxs) const {
+    static_assert(sizeof...(Index) == sizeof...(Iters), "Expect same number of offsets as iterators.");
+    return mute::transform(iters_, idxs, [](auto&& iter, auto&& idx) { return iter + idx; });
+  }
+
+  template <class... Index>
+  MUTE_HOST_DEVICE constexpr
+  reference operator[](mute::tuple<Index...> const& idxs) const {
+    return *(*this + idxs);
+  }
+
+  mute::tuple<Iters...> iters_;
+};
+
+//------------------------------------------------------------------------------
+// type traits
+
+template <class... Iters>
+struct is_rmem<ZipIterator<Iters...>> : conjunction<is_rmem<Iters>...> {};
+template <class... Iters>
+struct is_smem<ZipIterator<Iters...>> : conjunction<is_smem<Iters>...> {};
+template <class... Iters>
+struct is_gmem<ZipIterator<Iters...>> : conjunction<is_gmem<Iters>...> {};
+// A tuple of Layouts that operates on each Layout symmetrically
+// The Layouts need to have compatible shapes and ranks.
+// The ZipLayout presents the intersection of the domain of its component Layouts.
+//   E.g. all Layouts accept 1D coords and ZipLayout does as well.
+// The ZipLayout returns the union of the codomain of its component Layouts.
+//   E.g. all Layouts return an integer so ZipLayout returns a tuple of integers.
+template <class... Layouts>
+struct ZipLayout
+{
+  static constexpr int rank = (int(0) | ... | Layouts::rank);
+
+  static_assert((is_layout<Layouts>::value && ...), "All template parameters must be layouts");
+  static_assert(((Layouts::rank == rank) && ...),   "All layouts must have the same rank");
+
+  MUTE_HOST_DEVICE constexpr
+  ZipLayout(Layouts const&... layouts)
+    : layouts_(layouts...)
+  {}
+
+  MUTE_HOST_DEVICE constexpr
+  ZipLayout(mute::tuple<Layouts...> const& layouts)
+    : layouts_(layouts)
+  {}
+
+  template <class Coord>
+  MUTE_HOST_DEVICE constexpr
+  auto
+  operator()(Coord const& coord) const {
+    if constexpr (has_underscore<Coord>::value) {
+      return ZipLayout(mute::transform(layouts_, [&] (auto layout) { return layout(coord); }));
+    } else {
+      return mute::transform(layouts_, [&] (auto layout) { return layout(coord); });
+    }
+
+    MUTE_GCC_UNREACHABLE;
+  }
+
+  // op() convenience function for multi-dimensional coordinates
+  template <class Coord0, class Coord1, class... Coords>
+  MUTE_HOST_DEVICE constexpr
+  decltype(auto)
+  operator()(Coord0 const& c0, Coord1 const& c1, Coords const&... cs) const {
+    return operator()(make_coord(c0,c1,cs...));
+  }
+
+  mute::tuple<Layouts...> layouts_;
+};
+
+template <class... Layouts>
+struct is_layout<ZipLayout<Layouts...>> : true_type {};
+
+//
+// make_zip_tensor and unzip_tensor
+//
+
+template <class... Engines, class... Layouts>
+MUTE_HOST_DEVICE constexpr
+auto
+make_zip_tensor(Tensor<Engines,Layouts> const&... tensors)
+{
+  return make_tensor(ZipIterator(tensors.data()...),
+                     ZipLayout(tensors.layout()...));
+}
+
+template <class Engine, class Layout>
+MUTE_HOST_DEVICE constexpr
+auto
+unzip_tensor(Tensor<Engine,Layout> const& tensor)
+{
+  return mute::transform(tensor.data().iters_, tensor.layout().layouts_,
+                         [](auto iter, auto layout) { return make_tensor(iter, layout); });
+}
+
+//
+// Utilities
+//
+
+template <int... Is, class... Layouts>
+MUTE_HOST_DEVICE constexpr
+auto
+rank(ZipLayout<Layouts...> const& layouts)
+{
+  return rank<Is...>(get<0>(layouts.layouts_));
+}
+
+template <int... Is, class... Layouts>
+MUTE_HOST_DEVICE constexpr
+auto
+size(ZipLayout<Layouts...> const& layouts)
+{
+  return size<Is...>(get<0>(layouts.layouts_));
+}
+
+//
+// Manipulation
+//
+
+// Extend each component layout to rank-N by appending Layout @a x.
+template <int N, class... Layouts, class ShapeX = _1, class StrideX = _0>
+MUTE_HOST_DEVICE constexpr
+auto
+append(ZipLayout<Layouts...>  const& layouts,
+       Layout<ShapeX,StrideX> const& x = {})
+{
+  return ZipLayout(mute::transform(layouts.layouts_, [&](auto t){ return append<N>(t, x); }));
+}
+
+// Extend each component layout to rank-N by prepending Layout @a x.
+template <int N, class... Layouts, class ShapeX = _1, class StrideX = _0>
+MUTE_HOST_DEVICE constexpr
+auto
+prepend(ZipLayout<Layouts...>  const& layouts,
+        Layout<ShapeX,StrideX> const& x = {})
+{
+  return ZipLayout(mute::transform(layouts.layouts_, [&](auto t){ return prepend<N>(t, x); }));
+}
+
+template <class... Layouts, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+logical_divide(ZipLayout<Layouts...> const& layouts,
+               Tiler                 const& tiler)
+{
+  return ZipLayout(mute::transform(layouts.layouts_, [&](auto t){ return logical_divide(t, tiler); }));
+}
+
+template <class... Layouts, class Tiler>
+MUTE_HOST_DEVICE constexpr
+auto
+zipped_divide(ZipLayout<Layouts...> const& layouts,
+              Tiler                 const& tiler)
+{
+  return ZipLayout(mute::transform(layouts.layouts_, [&](auto t){ return zipped_divide(t, tiler); }));
+}
+
+// Return <SlicedZipLayout, ZipOffsets> by calling slice_and_offset and all component layouts.
+template <class Coord, class... Layouts>
+MUTE_HOST_DEVICE constexpr
+auto
+slice_and_offset(Coord const& c, ZipLayout<Layouts...> const& layouts)
+{
+  auto result = mute::zip(mute::transform(layouts.layouts_, [&c](auto const& layout) { return slice_and_offset(c, layout); }));
+  return mute::make_tuple(ZipLayout(get<0>(result)), get<1>(result));
+}
+
+} // end namespace mute
\ No newline at end of file
diff --git a/vllm_musa/csrc/flash_attn/include/mute/tile.hpp b/vllm_musa/csrc/flash_attn/include/mute/tile.hpp
new file mode 100644
index 000000000..aa6c5adde
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/tile.hpp
@@ -0,0 +1,59 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>
+
+#include <mute/layout.hpp>
+
+namespace mute
+{
+
+//
+// A Tile is not a Layout, it's a tuple of Layouts or Tiles or Underscores
+//
+
+template <class... Layouts>
+using Tile = tuple<Layouts...>;
+
+template <class Tile>
+using is_tile = is_tuple<Tile>;
+
+template <class... Layouts>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tile(Layouts const&... layouts)
+{
+  return Tile<Layouts...>(layouts...);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/underscore.hpp b/vllm_musa/csrc/flash_attn/include/mute/underscore.hpp
new file mode 100644
index 000000000..9d71bf2a7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/underscore.hpp
@@ -0,0 +1,195 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>                     // MUTE_INLINE_CONSTANT, MUTE_HOST_DEVICE
+#include <mute/container/tuple.hpp>            // mute::is_tuple
+#include <mute/numeric/integral_constant.hpp>  // mute::false_type, mute::true_type
+
+namespace mute
+{
+
+// For slicing
+struct Underscore : Int<0> {};
+
+MUTE_INLINE_CONSTANT Underscore _;
+
+// Convenient alias
+using X = Underscore;
+
+// Treat Underscore as an integral like integral_constant
+template <>
+struct is_integral<Underscore> : true_type {};
+
+template <class T>
+struct is_underscore : false_type {};
+template <>
+struct is_underscore<Underscore> : true_type {};
+
+// Tuple trait for detecting static member element
+template <class Tuple, class Elem, class Enable = void>
+struct has_elem : false_type {};
+template <class Elem>
+struct has_elem<Elem, Elem> : true_type {};
+template <class Tuple, class Elem>
+struct has_elem<Tuple, Elem, enable_if_t<is_tuple<Tuple>::value> >
+    : has_elem<Tuple, Elem, tuple_seq<Tuple> > {};
+template <class Tuple, class Elem, int... Is>
+struct has_elem<Tuple, Elem, seq<Is...>>
+    : disjunction<has_elem<tuple_element_t<Is, Tuple>, Elem>...> {};
+
+// Tuple trait for detecting static member element
+template <class Tuple, class Elem, class Enable = void>
+struct all_elem : false_type {};
+template <class Elem>
+struct all_elem<Elem, Elem> : true_type {};
+template <class Tuple, class Elem>
+struct all_elem<Tuple, Elem, enable_if_t<is_tuple<Tuple>::value> >
+    : all_elem<Tuple, Elem, tuple_seq<Tuple> > {};
+template <class Tuple, class Elem, int... Is>
+struct all_elem<Tuple, Elem, seq<Is...>>
+    : conjunction<all_elem<tuple_element_t<Is, Tuple>, Elem>...> {};
+
+// Tuple trait for detecting Underscore member
+template <class Tuple>
+using has_underscore = has_elem<Tuple, Underscore>;
+
+template <class Tuple>
+using all_underscore = all_elem<Tuple, Underscore>;
+
+template <class Tuple>
+using has_int1 = has_elem<Tuple, Int<1>>;
+
+template <class Tuple>
+using has_int0 = has_elem<Tuple, Int<0>>;
+
+//
+// Slice keeps only the elements of Tuple B that are paired with an Underscore
+//
+
+namespace detail {
+
+template <class A, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+lift_slice(A const& a, B const& b)
+{
+  if constexpr (is_tuple<A>::value) {
+    static_assert(tuple_size<A>::value == tuple_size<B>::value, "Mismatched Ranks");
+    return filter_tuple(a, b, [](auto const& x, auto const& y) { return lift_slice(x,y); });
+  } else if constexpr (is_underscore<A>::value) {
+    return mute::tuple<B>{b};
+  } else {
+    return mute::tuple<>{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+// Entry point overrides the lifting so that slice(_,b) == b
+template <class A, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+slice(A const& a, B const& b)
+{
+  if constexpr (is_tuple<A>::value) {
+    static_assert(tuple_size<A>::value == tuple_size<B>::value, "Mismatched Ranks");
+    return filter_tuple(a, b, [](auto const& x, auto const& y) { return detail::lift_slice(x,y); });
+  } else if constexpr (is_underscore<A>::value) {
+    return b;
+  } else {
+    return mute::tuple<>{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Dice keeps only the elements of Tuple B that are paired with an Int
+//
+
+namespace detail {
+
+template <class A, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+lift_dice(A const& a, B const& b)
+{
+  if constexpr (is_tuple<A>::value) {
+    static_assert(tuple_size<A>::value == tuple_size<B>::value, "Mismatched Ranks");
+    return filter_tuple(a, b, [](auto const& x, auto const& y) { return lift_dice(x,y); });
+  } else if constexpr (is_underscore<A>::value) {
+    return mute::tuple<>{};
+  } else {
+    return mute::tuple<B>{b};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+} // end namespace detail
+
+// Entry point overrides the lifting so that dice(1,b) == b
+template <class A, class B>
+MUTE_HOST_DEVICE constexpr
+auto
+dice(A const& a, B const& b)
+{
+  if constexpr (is_tuple<A>::value) {
+    static_assert(tuple_size<A>::value == tuple_size<B>::value, "Mismatched Ranks");
+    return filter_tuple(a, b, [](auto const& x, auto const& y) { return detail::lift_dice(x,y); });
+  } else if constexpr (is_underscore<A>::value) {
+    return mute::tuple<>{};
+  } else {
+    return b;
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+//
+// Display utilities
+//
+
+MUTE_HOST_DEVICE void print(Underscore const&) {
+  printf("_");
+}
+
+#if !defined(__MUSACC_RTC__)
+MUTE_HOST std::ostream& operator<<(std::ostream& os, Underscore const&) {
+  return os << "_";
+}
+#endif
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/util/debug.hpp b/vllm_musa/csrc/flash_attn/include/mute/util/debug.hpp
new file mode 100644
index 000000000..f4982d546
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/util/debug.hpp
@@ -0,0 +1,165 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+/**
+ * \file
+ * \brief Debugging and logging functionality
+ */
+
+#include <musa_runtime_api.h>
+
+#include <mute/config.hpp>
+
+namespace mute
+{
+
+/******************************************************************************
+ * Debug and logging macros
+ ******************************************************************************/
+
+/**
+ * Formats and prints the given message to stdout
+ */
+#if !defined(MUTE_LOG)
+#  if !defined(__MUSA_ARCH__)
+#    define MUTE_LOG(format, ...) printf(format, __VA_ARGS__)
+#  else
+#    define MUTE_LOG(format, ...)                                \
+        printf("[block (%d,%d,%d), thread (%d,%d,%d)]: " format, \
+               blockIdx.x,  blockIdx.y,  blockIdx.z,             \
+               threadIdx.x, threadIdx.y, threadIdx.z,            \
+               __VA_ARGS__);
+#  endif
+#endif
+
+/**
+ * Formats and prints the given message to stdout only if DEBUG is defined
+ */
+#if !defined(MUTE_LOG_DEBUG)
+#  ifdef DEBUG
+#    define MUTE_LOG_DEBUG(format, ...) MUTE_LOG(format, __VA_ARGS__)
+#  else
+#    define MUTE_LOG_DEBUG(format, ...)
+#  endif
+#endif
+
+/**
+ * \brief Perror macro with exit
+ */
+#if !defined(MUTE_ERROR_EXIT)
+#  define MUTE_ERROR_EXIT(e)                                         \
+      do {                                                           \
+        musaError_t code = (e);                                      \
+        if (code != musaSuccess) {                                   \
+          fprintf(stderr, "<%s:%d> %s:\n    %s: %s\n",               \
+                  __FILE__, __LINE__, #e,                            \
+                  musaGetErrorName(code), musaGetErrorString(code)); \
+          fflush(stderr);                                            \
+          exit(1);                                                   \
+        }                                                            \
+      } while (0)
+#endif
+
+#if !defined(MUTE_CHECK_LAST)
+#  define MUTE_CHECK_LAST() MUTE_ERROR_EXIT(musaPeekAtLastError()); MUTE_ERROR_EXIT(musaDeviceSynchronize())
+#endif
+
+#if !defined(MUTE_CHECK_ERROR)
+#  define MUTE_CHECK_ERROR(e) MUTE_ERROR_EXIT(e)
+#endif
+
+// A dummy function that uses compilation failure to print a type
+template <class... T>
+MUTE_HOST_DEVICE void
+print_type() {
+  static_assert(sizeof...(T) < 0, "Printing type T.");
+}
+
+template <class... T>
+MUTE_HOST_DEVICE void
+print_type(T&&...) {
+  static_assert(sizeof...(T) < 0, "Printing type T.");
+}
+
+//
+// Device-specific helpers
+//
+// e.g.
+// if (thread0()) print(...);
+// if (block0()) print(...);
+// if (thread(42)) print(...);
+
+MUTE_HOST_DEVICE
+bool
+block(int bid)
+{
+#if defined(__MUSA_ARCH__)
+  return blockIdx.x + blockIdx.y*gridDim.x + blockIdx.z*gridDim.x*gridDim.y == static_cast<unsigned int>(bid);;
+#else
+  return true;
+#endif
+}
+
+MUTE_HOST_DEVICE
+bool
+thread(int tid, int bid)
+{
+#if defined(__MUSA_ARCH__)
+  return (threadIdx.x + threadIdx.y*blockDim.x + threadIdx.z*blockDim.x*blockDim.y == static_cast<unsigned int>(tid)) && block(bid);
+#else
+  return true;
+#endif
+}
+
+MUTE_HOST_DEVICE
+bool
+thread(int tid)
+{
+  return thread(tid,0);
+}
+
+MUTE_HOST_DEVICE
+bool
+thread0()
+{
+  return thread(0,0);
+}
+
+MUTE_HOST_DEVICE
+bool
+block0()
+{
+  return block(0);
+}
+
+}  // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/util/print.hpp b/vllm_musa/csrc/flash_attn/include/mute/util/print.hpp
new file mode 100644
index 000000000..47a0f43eb
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/util/print.hpp
@@ -0,0 +1,200 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/config.hpp>           // MUTE_HOST_DEVICE
+#include <mute/util/type_traits.hpp> // mute::is_valid
+
+//
+// MUSA compatible print and printf
+//
+
+namespace mute
+{
+
+MUTE_HOST_DEVICE
+int
+num_digits(int x)
+{
+  return (x < 10 ? 1 :
+          (x < 100 ? 2 :
+           (x < 1000 ? 3 :
+            (x < 10000 ? 4 :
+             (x < 100000 ? 5 :
+              (x < 1000000 ? 6 :
+               (x < 10000000 ? 7 :
+                (x < 100000000 ? 8 :
+                 (x < 1000000000 ? 9 :
+                  10)))))))));
+}
+
+//
+// print dispatcher
+//
+
+MUTE_HOST_DEVICE
+void
+print(char c) {
+  printf("%c", c);
+}
+
+MUTE_HOST_DEVICE
+void
+print(signed char a) {
+  printf("%d", static_cast<int>(a));
+}
+
+MUTE_HOST_DEVICE
+void
+print(unsigned char a) {
+  printf("%u", static_cast<unsigned int>(a));
+}
+
+MUTE_HOST_DEVICE
+void
+print(short a) {
+  printf("%hd", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(unsigned short a) {
+  printf("%hu", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(int a) {
+  printf("%d", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(unsigned int a) {
+  printf("%u", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(long a) {
+  printf("%ld", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(unsigned long a) {
+  printf("%lu", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(long long a) {
+  printf("%lld", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(unsigned long long a) {
+  printf("%llu", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(float a) {
+  printf("%f", a);
+}
+
+MUTE_HOST_DEVICE
+void
+print(double a) {
+  printf("%f", a);
+}
+
+template <class... T>
+MUTE_HOST_DEVICE
+void
+print(char const* format, T const&... t) {
+  printf(format, t...);
+}
+
+MUTE_HOST_DEVICE
+void
+print(char const* format) {
+  printf("%s", format);
+}
+
+//
+// pretty printing
+//
+
+MUTE_HOST_DEVICE void
+pretty_print(bool v) {
+  printf("%*d", 3, int(v));
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(int32_t v) {
+  printf("%*d", 5, v);
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(uint32_t v) {
+  printf("%*d", 5, v);
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(int64_t v) {
+  printf("%*lld", 5, static_cast<long long>(v));
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(uint64_t v) {
+  printf("%*llu", 5, static_cast<unsigned long long>(v));
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(float v) {
+  printf("%*.2e", 10, v);
+}
+
+MUTE_HOST_DEVICE void
+pretty_print(double v) {
+  printf("%*.3e", 11, v);
+}
+
+template <class T>
+MUTE_HOST_DEVICE void
+pretty_print(T t) {
+  printf("  "); print(t);
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mute/util/type_traits.hpp b/vllm_musa/csrc/flash_attn/include/mute/util/type_traits.hpp
new file mode 100644
index 000000000..8a08e3d63
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mute/util/type_traits.hpp
@@ -0,0 +1,299 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/type_traits>
+#include <musa/std/utility>
+#include <musa/std/cstddef>
+#include <musa/std/cstdint>
+#include <musa/std/limits>
+#else
+#include <type_traits>
+#include <utility>      // tuple_size, tuple_element
+#include <cstddef>      // ptrdiff_t
+#include <cstdint>      // uintptr_t
+#include <limits>       // numeric_limits
+#endif
+
+#include <mute/config.hpp> // MUTE_STL_NAMESPACE
+
+namespace mute
+{
+  using MUTE_STL_NAMESPACE::enable_if;
+  using MUTE_STL_NAMESPACE::enable_if_t;
+}
+
+#define __MUTE_REQUIRES(...)   typename mute::enable_if<(__VA_ARGS__)>::type* = nullptr
+#define __MUTE_REQUIRES_V(...) typename mute::enable_if<decltype((__VA_ARGS__))::value>::type* = nullptr
+
+namespace mute
+{
+
+// <type_traits>
+using MUTE_STL_NAMESPACE::conjunction;
+using MUTE_STL_NAMESPACE::conjunction_v;
+
+using MUTE_STL_NAMESPACE::disjunction;
+using MUTE_STL_NAMESPACE::disjunction_v;
+
+using MUTE_STL_NAMESPACE::negation;
+using MUTE_STL_NAMESPACE::negation_v;
+
+using MUTE_STL_NAMESPACE::void_t;
+using MUTE_STL_NAMESPACE::is_void_v;
+
+using MUTE_STL_NAMESPACE::is_base_of;
+using MUTE_STL_NAMESPACE::is_base_of_v;
+
+using MUTE_STL_NAMESPACE::is_const;
+using MUTE_STL_NAMESPACE::is_const_v;
+using MUTE_STL_NAMESPACE::is_volatile;
+using MUTE_STL_NAMESPACE::is_volatile_v;
+
+// Defined in mute/numeric/integral_constant.hpp
+// using MUTE_STL_NAMESPACE::true_type;
+// using MUTE_STL_NAMESPACE::false_type;
+
+using MUTE_STL_NAMESPACE::conditional;
+using MUTE_STL_NAMESPACE::conditional_t;
+
+using MUTE_STL_NAMESPACE::add_const_t;
+
+using MUTE_STL_NAMESPACE::remove_const_t;
+using MUTE_STL_NAMESPACE::remove_cv_t;
+using MUTE_STL_NAMESPACE::remove_reference_t;
+
+using MUTE_STL_NAMESPACE::extent;
+using MUTE_STL_NAMESPACE::remove_extent;
+
+using MUTE_STL_NAMESPACE::decay;
+using MUTE_STL_NAMESPACE::decay_t;
+
+using MUTE_STL_NAMESPACE::is_lvalue_reference;
+using MUTE_STL_NAMESPACE::is_lvalue_reference_v;
+
+using MUTE_STL_NAMESPACE::is_reference;
+using MUTE_STL_NAMESPACE::is_trivially_copyable;
+
+using MUTE_STL_NAMESPACE::is_convertible;
+using MUTE_STL_NAMESPACE::is_convertible_v;
+
+using MUTE_STL_NAMESPACE::is_same;
+using MUTE_STL_NAMESPACE::is_same_v;
+
+using MUTE_STL_NAMESPACE::is_constructible;
+using MUTE_STL_NAMESPACE::is_constructible_v;
+using MUTE_STL_NAMESPACE::is_default_constructible;
+using MUTE_STL_NAMESPACE::is_default_constructible_v;
+using MUTE_STL_NAMESPACE::is_standard_layout;
+using MUTE_STL_NAMESPACE::is_standard_layout_v;
+
+using MUTE_STL_NAMESPACE::is_arithmetic;
+using MUTE_STL_NAMESPACE::is_unsigned;
+using MUTE_STL_NAMESPACE::is_unsigned_v;
+using MUTE_STL_NAMESPACE::is_signed;
+using MUTE_STL_NAMESPACE::is_signed_v;
+
+using MUTE_STL_NAMESPACE::make_signed;
+using MUTE_STL_NAMESPACE::make_signed_t;
+
+// using MUTE_STL_NAMESPACE::is_integral;
+template <class T>
+using is_std_integral = MUTE_STL_NAMESPACE::is_integral<T>;
+
+using MUTE_STL_NAMESPACE::is_empty;
+using MUTE_STL_NAMESPACE::is_empty_v;
+
+using MUTE_STL_NAMESPACE::invoke_result_t;
+
+using MUTE_STL_NAMESPACE::common_type;
+using MUTE_STL_NAMESPACE::common_type_t;
+
+using MUTE_STL_NAMESPACE::remove_pointer;
+using MUTE_STL_NAMESPACE::remove_pointer_t;
+
+using MUTE_STL_NAMESPACE::add_pointer;
+using MUTE_STL_NAMESPACE::add_pointer_t;
+
+using MUTE_STL_NAMESPACE::alignment_of;
+using MUTE_STL_NAMESPACE::alignment_of_v;
+
+using MUTE_STL_NAMESPACE::is_pointer;
+using MUTE_STL_NAMESPACE::is_pointer_v;
+
+// <utility>
+using MUTE_STL_NAMESPACE::declval;
+
+template <class T>
+constexpr T&& forward(remove_reference_t<T>& t) noexcept
+{
+  return static_cast<T&&>(t);
+}
+
+template <class T>
+constexpr T&& forward(remove_reference_t<T>&& t) noexcept
+{
+  static_assert(! is_lvalue_reference_v<T>, "T cannot be an lvalue reference (e.g., U&).");
+  return static_cast<T&&>(t);
+}
+
+template <class T>
+constexpr remove_reference_t<T>&& move(T&& t) noexcept
+{
+  return static_cast<remove_reference_t<T>&&>(t);
+}
+
+// <limits>
+using MUTE_STL_NAMESPACE::numeric_limits;
+
+// <cstddef>
+using MUTE_STL_NAMESPACE::ptrdiff_t;
+
+// <cstdint>
+using MUTE_STL_NAMESPACE::uintptr_t;
+
+// C++20
+// using std::remove_cvref;
+template <class T>
+struct remove_cvref {
+  using type = remove_cv_t<remove_reference_t<T>>;
+};
+
+// C++20
+// using std::remove_cvref_t;
+template <class T>
+using remove_cvref_t = typename remove_cvref<T>::type;
+
+//
+// dependent_false
+//
+// @brief An always-false value that depends on one or more template parameters.
+// See
+// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2019/p1830r1.pdf
+// https://github.com/cplusplus/papers/issues/572
+// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2022/p2593r0.html
+template <class... Args>
+inline constexpr bool dependent_false = false;
+
+//
+// tuple_size, tuple_element
+//
+// @brief MuTe-local tuple-traits to prevent conflicts with other libraries.
+// For mute:: types, we specialize std::tuple-traits, which is explicitly allowed.
+//   mute::tuple, mute::array, mute::array_subbyte, etc
+// But MuTe wants to treat some external types as tuples as well. For those,
+// we specialize mute::tuple-traits to avoid polluting external traits.
+//   dim3, uint3, etc
+
+template <class T, class = void>
+struct tuple_size;
+
+template <class T>
+struct tuple_size<T,void_t<typename MUTE_STL_NAMESPACE::tuple_size<T>::type>> : MUTE_STL_NAMESPACE::integral_constant<size_t, MUTE_STL_NAMESPACE::tuple_size<T>::value> {};
+
+// S =  : std::integral_constant<std::size_t, std::tuple_size<T>::value> {};
+
+template <class T>
+constexpr size_t tuple_size_v = tuple_size<T>::value;
+
+template <size_t I, class T, class = void>
+struct tuple_element;
+
+template <size_t I, class T>
+struct tuple_element<I,T,void_t<typename MUTE_STL_NAMESPACE::tuple_element<I,T>::type>> : MUTE_STL_NAMESPACE::tuple_element<I,T> {};
+
+template <size_t I, class T>
+using tuple_element_t = typename tuple_element<I,T>::type;
+
+//
+// is_valid
+//
+
+namespace detail {
+
+template <class F, class... Args, class = decltype(declval<F&&>()(declval<Args&&>()...))>
+MUTE_HOST_DEVICE constexpr auto
+is_valid_impl(int) { return MUTE_STL_NAMESPACE::true_type{}; }
+
+template <class F, class... Args>
+MUTE_HOST_DEVICE constexpr auto
+is_valid_impl(...) { return MUTE_STL_NAMESPACE::false_type{}; }
+
+template <class F>
+struct is_valid_fn {
+  template <class... Args>
+  MUTE_HOST_DEVICE constexpr auto
+  operator()(Args&&...) const { return is_valid_impl<F, Args&&...>(int{}); }
+};
+
+} // end namespace detail
+
+template <class F>
+MUTE_HOST_DEVICE constexpr auto
+is_valid(F&&) {
+  return detail::is_valid_fn<F&&>{};
+}
+
+template <class F, class... Args>
+MUTE_HOST_DEVICE constexpr auto
+is_valid(F&&, Args&&...) {
+  return detail::is_valid_impl<F&&, Args&&...>(int{});
+}
+
+template <bool B, template<class...> class True, template<class...> class False>
+struct conditional_template {
+  template <class... U>
+  using type = True<U...>;
+};
+
+template <template<class...> class True, template<class...> class False>
+struct conditional_template<false, True, False> {
+  template <class... U>
+  using type = False<U...>;
+};
+
+//
+// is_any_of
+//
+
+// Member `value` is true if and only if T is same as (is_same_v) at least one of the types in Us
+template <class T, class... Us>
+struct is_any_of {
+  constexpr static bool value = (... || MUTE_STL_NAMESPACE::is_same_v<T, Us>);
+};
+
+// Is true if and only if T is same as (is_same_v) at least one of the types in Us
+template <class T, class... Us>
+inline constexpr bool is_any_of_v = is_any_of<T, Us...>::value;
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/aligned_buffer.h b/vllm_musa/csrc/flash_attn/include/mutlass/aligned_buffer.h
new file mode 100644
index 000000000..09cd37200
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/aligned_buffer.h
@@ -0,0 +1,129 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 AlignedBuffer is a container for trivially copyable elements suitable for use in
+      unions and shared memory.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Modifies semantics of mutlass::Array<> to provide guaranteed alignment.
+template <
+  typename T,
+  int N,
+  int Align = 16
+>
+struct AlignedBuffer {
+
+  /// Internal storage type
+  using Storage = uint8_t;
+
+  /// Number of logical elements held in buffer
+  static int const kCount = N;
+
+  /// Alignment requirement in bytes
+  static int const kAlign = Align;
+
+  /// Number of storage elements
+  static int const kBytes =
+    (sizeof_bits<T>::value * N + 7) / 8;
+
+private:
+
+  /// Internal storage
+  alignas(Align) Storage storage[kBytes];
+
+public:
+
+  //
+  // C++ standard members
+  //
+
+  typedef T value_type;
+  typedef size_t size_type;
+  typedef ptrdiff_t difference_type;
+  typedef value_type *pointer;
+  typedef value_type const * const_pointer;
+
+  using Array = Array<T, N>;
+  using reference = typename Array::reference;
+  using const_reference = typename Array::const_reference;
+
+public:
+
+  MUTLASS_HOST_DEVICE
+  pointer data() {
+    return reinterpret_cast<pointer>(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_pointer data() const {
+    return reinterpret_cast<pointer>(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Storage * raw_data() {
+    return storage;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Storage const * raw_data() const {
+    return storage;
+  }
+
+
+  MUTLASS_HOST_DEVICE
+  constexpr bool empty() const {
+    return !kCount;
+  }
+
+  MUTLASS_HOST_DEVICE
+  constexpr size_type size() const {
+    return kCount;
+  }
+
+  MUTLASS_HOST_DEVICE
+  constexpr size_type max_size() const {
+    return kCount;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/arch/arch.h b/vllm_musa/csrc/flash_attn/include/mutlass/arch/arch.h
new file mode 100644
index 000000000..ee1807937
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/arch/arch.h
@@ -0,0 +1,60 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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
+
+#include "mutlass/mutlass.h"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace arch {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Mp22 {
+  static int const kMinComputeCapability = 22;
+};
+
+struct Mp31 {
+  static int const kMinComputeCapability = 31;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace arch
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/arch/barrier.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/arch/barrier.hpp
new file mode 100644
index 000000000..48f423006
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/arch/barrier.hpp
@@ -0,0 +1,182 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mutlass/mutlass.h>
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ >= 310
+#define MUSA_BARRIER_ENABLED 1
+#else
+#define MUSA_BARRIER_ENABLED 0
+#endif
+
+namespace mutlass {
+namespace arch {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+MUTLASS_DEVICE
+void allocate_async_barriers(const uint32_t num) {
+#if MUSA_BARRIER_ENABLED
+  __musa_async_bar_record(num);
+#endif
+}
+
+struct AsyncBarrier {
+  // Range: [0, 63]
+  uint32_t const id_;
+
+public:
+  static const uint32_t ReservedAsyncBarrierCount = 1;
+
+  static const uint32_t HardwareMaxNumAsyncTransactionBarriers = 64;
+
+  MUTLASS_DEVICE
+  AsyncBarrier(uint32_t id = 0)
+    : id_(id + ReservedAsyncBarrierCount) {}
+
+  MUTLASS_DEVICE
+  void init(uint32_t arrive_count, uint32_t init_phase = 0) const {
+    AsyncBarrier::init(id_, arrive_count, init_phase);
+  }
+
+  template <bool return_phase = false>
+  MUTLASS_DEVICE
+  auto arrive() const {
+    return AsyncBarrier::arrive<return_phase>(id_);
+  }
+
+  MUTLASS_DEVICE
+  void wait(uint32_t phase) const {
+    AsyncBarrier::wait(id_, phase);
+  }
+
+  MUTLASS_DEVICE
+  uint32_t get_barrier_id() const {
+    return id_;
+  }
+
+  //
+  //  Static Versions
+  //
+  MUTLASS_DEVICE
+  static void init(uint32_t id, uint32_t arrive_count, uint32_t init_phase) {
+#if MUSA_BARRIER_ENABLED
+    __musa_async_init_arrival(id, arrive_count, init_phase);
+#endif
+  }
+
+  template <bool return_phase = false>
+  MUTLASS_DEVICE
+  static auto arrive(uint32_t id) {
+#if MUSA_BARRIER_ENABLED
+    if constexpr (return_phase) {
+      return __musa_async_arrive(id);
+    } else {
+      __musa_async_arrive(id);
+    }
+#else
+    if constexpr (return_phase) {
+      return -1;
+    } else {
+      return;
+    }
+#endif
+  }
+
+  MUTLASS_DEVICE
+  static void wait(uint32_t id, uint32_t phase) {
+#if MUSA_BARRIER_ENABLED
+    __musa_async_wait(id, phase);
+#endif
+  }
+};
+
+struct AsyncTransactionBarrier : public AsyncBarrier {
+  using AsyncBarrier::AsyncBarrier;
+
+  template <bool return_phase = false>
+  MUTLASS_DEVICE
+  auto arrive_and_expect_tx(uint32_t transaction_bytes) const {
+    return AsyncTransactionBarrier::arrive_and_expect_tx<return_phase>(id_, transaction_bytes);
+  }
+
+  MUTLASS_DEVICE
+  void expect_transaction(uint32_t transaction_bytes) const {
+    AsyncTransactionBarrier::expect_transaction(id_, transaction_bytes);
+  }
+
+  MUTLASS_DEVICE
+  void complete_transaction(uint32_t transaction_bytes) const {
+    AsyncTransactionBarrier::complete_transaction(id_, transaction_bytes);
+  }
+
+  //
+  //  Static Versions
+  //
+
+  template <bool return_phase = false>
+  MUTLASS_DEVICE
+  static auto arrive_and_expect_tx(uint32_t id, uint32_t transaction_bytes) {
+#if MUSA_BARRIER_ENABLED
+    __musa_async_add_trans(id, transaction_bytes);
+    uint32_t phase = __musa_async_arrive(id);
+
+    if constexpr (return_phase) {
+      return phase;
+    } else {
+      (void)phase;
+      return;
+    }
+#endif
+  }
+
+  MUTLASS_DEVICE
+  static void expect_transaction(uint32_t id, uint32_t transaction_bytes) {
+#if MUSA_BARRIER_ENABLED
+    __musa_async_add_trans(id, transaction_bytes);
+#endif
+  }
+
+  MUTLASS_DEVICE
+  static void complete_transaction(uint32_t id, uint32_t transaction_bytes) {
+#if MUSA_BARRIER_ENABLED
+    __musa_async_decrease_trans(id, transaction_bytes);
+#endif
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+}  // end namespace arch
+}  // end namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/arch/cache_operation.h b/vllm_musa/csrc/flash_attn/include/mutlass/arch/cache_operation.h
new file mode 100644
index 000000000..711056e88
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/arch/cache_operation.h
@@ -0,0 +1,67 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Directives related to cache operations
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+
+namespace mutlass {
+namespace arch {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Controls Intrinsic cache operations
+struct CacheOperation {
+  enum Kind {
+    /// Cache at all levels - accessed again
+    Always,
+    /// Cache at global level
+    Global,
+    /// Streaming - likely to be accessed once
+    Streaming,
+    /// Indicates the line will not be used again
+    LastUse,
+    /// Don't cache, and fetch again
+    Volatile,
+    /// Write back at all coherent levels
+    WriteBack,
+    /// Write through to system memory
+    WriteThrough
+  };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace arch
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/arch/memory.h b/vllm_musa/csrc/flash_attn/include/mutlass/arch/memory.h
new file mode 100644
index 000000000..a4debce25
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/arch/memory.h
@@ -0,0 +1,355 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Architecture-specific operators on memory
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/arch/cache_operation.h"
+
+namespace mutlass {
+namespace arch {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// AddressSpace
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+enum class AddressSpace {
+  Generic = 0,
+  Global  = 1,
+  Shared  = 3,
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// Fragment type to store loaded data
+    typename AccessType,
+    /// The bytes of loading
+    int LoadBytes,
+    /// Cache operation
+    CacheOperation::Kind cache_op = CacheOperation::Always
+    >
+struct global_load;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Specializations
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename AccessType>
+struct global_load<AccessType,
+                   32
+                  > {
+  MUTLASS_DEVICE
+  global_load(AccessType &D, void const *ptr, bool pred_guard) {
+    using Type = uint32_t __attribute__((vector_size(32)));
+    Type &data = reinterpret_cast<Type &>(D);
+    if (pred_guard) {
+      data = *reinterpret_cast<Type const *>(ptr);
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_load<AccessType,
+                   16
+                  > {
+  MUTLASS_DEVICE
+  global_load(AccessType &D, void const *ptr, bool pred_guard) {
+    using Type = uint32_t __attribute__((vector_size(16)));
+    Type &data = reinterpret_cast<Type &>(D);
+    if (pred_guard) {
+      data = *reinterpret_cast<Type const *>(ptr);
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_load<AccessType,
+                   8
+                  > {
+  MUTLASS_DEVICE
+  global_load(AccessType &D, void const *ptr, bool pred_guard) {
+    uint64_t &data = reinterpret_cast<uint64_t &>(D);
+    if (pred_guard) {
+      data = *reinterpret_cast<uint64_t const *>(ptr);
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_load<AccessType,
+                   4
+                  > {
+  MUTLASS_DEVICE
+  global_load(AccessType &D, void const *ptr, bool pred_guard) {
+    uint32_t &data = reinterpret_cast<uint32_t &>(D);
+    if (pred_guard) {
+      data = *reinterpret_cast<uint32_t const *>(ptr);
+    }
+
+  }
+};
+
+template <typename AccessType>
+struct global_load<AccessType,
+                   2
+                  > {
+  MUTLASS_DEVICE
+  global_load(AccessType &D, void const *ptr, bool pred_guard) {
+    uint16_t &data = reinterpret_cast<uint16_t &>(D);
+    if (pred_guard) {
+      data = *reinterpret_cast<uint16_t const *>(ptr);
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_load<AccessType,
+                   1
+                  > {
+  MUTLASS_DEVICE
+  global_load(AccessType &D, void const *ptr, bool pred_guard) {
+    uint8_t &data = reinterpret_cast<uint8_t &>(D);
+    if (pred_guard) {
+      data = *reinterpret_cast<uint8_t const *>(ptr);
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+    /// Fragment type to store data
+    typename AccessType,
+    /// The bytes of storing
+    int StoreBytes
+    >
+struct global_store;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Specializations
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+template <typename AccessType>
+struct global_store<AccessType, 64> {
+  MUTLASS_DEVICE
+  global_store(AccessType const &D, void *ptr, bool pred_guard) {
+    using Type = uint32_t __attribute__((vector_size(64)));
+    Type const &data = reinterpret_cast<Type const &>(D);
+    if (pred_guard) {
+      *(reinterpret_cast<Type *>(ptr)) = data;
+    }
+  }
+};
+
+
+template <typename AccessType>
+struct global_store<AccessType, 32> {
+  MUTLASS_DEVICE
+  global_store(AccessType const &D, void *ptr, bool pred_guard) {
+    using Type = uint32_t __attribute__((vector_size(32)));
+    Type const &data = reinterpret_cast<Type const &>(D);
+    if (pred_guard) {
+      *(reinterpret_cast<Type *>(ptr)) = data;
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_store<AccessType, 16> {
+  MUTLASS_DEVICE
+  global_store(AccessType const &D, void *ptr, bool pred_guard) {
+    using Type = uint32_t __attribute__((vector_size(16)));
+    Type const &data = reinterpret_cast<Type const &>(D);
+    if (pred_guard) {
+      *(reinterpret_cast<Type *>(ptr)) = data;
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_store<AccessType, 8> {
+  MUTLASS_DEVICE
+  global_store(AccessType const &D, void *ptr, bool pred_guard) {
+    uint64_t const &data = reinterpret_cast<uint64_t const &>(D);
+    if (pred_guard) {
+      *(reinterpret_cast<uint64_t *>(ptr)) = data;
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_store<AccessType, 4> {
+  MUTLASS_DEVICE
+  global_store(AccessType const &D, void *ptr, bool pred_guard) {
+    uint32_t const &data = reinterpret_cast<uint32_t const &>(D);
+    if (pred_guard) {
+      *(reinterpret_cast<uint32_t *>(ptr)) = data;
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_store<AccessType, 2> {
+  MUTLASS_DEVICE
+  global_store(AccessType const &D, void *ptr, bool pred_guard) {
+    uint16_t const &data = reinterpret_cast<uint16_t const &>(D);
+    if (pred_guard) {
+      *(reinterpret_cast<uint16_t *>(ptr)) = data;
+    }
+  }
+};
+
+template <typename AccessType>
+struct global_store<AccessType, 1> {
+  MUTLASS_DEVICE
+  global_store(AccessType const &D, void *ptr, bool pred_guard) {
+    uint8_t const &data = reinterpret_cast<uint8_t const &>(D);
+    if (pred_guard) {
+      *(reinterpret_cast<uint8_t *>(ptr)) = data;
+    }
+  }
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// ld.shared
+template <int Bytes>
+MUTLASS_DEVICE
+void shared_load(void *dst, uint32_t ptr);
+
+/// ld.shared - 16b
+template <>
+MUTLASS_DEVICE
+void shared_load<2>(void *dst, uint32_t ptr) {
+  auto smem_ptr = reinterpret_cast<
+                    uint16_t __attribute__((address_space(static_cast<int>(AddressSpace::Shared))))*>(ptr);
+  uint16_t *dst_u16  = reinterpret_cast<uint16_t*>(dst);
+  *dst_u16 = *smem_ptr;
+}
+
+/// ld.shared - 32b
+template <>
+MUTLASS_DEVICE
+void shared_load<4>(void *dst, uint32_t ptr) {
+  auto smem_ptr = reinterpret_cast<
+                    uint32_t __attribute__((address_space(static_cast<int>(AddressSpace::Shared))))*>(ptr);
+  uint32_t *dst_u32  = reinterpret_cast<uint32_t*>(dst);
+  *dst_u32 = *smem_ptr;
+}
+
+/// ld.shared - 64b
+template <>
+MUTLASS_DEVICE
+void shared_load<8>(void *dst, uint32_t ptr) {
+  auto smem_ptr = reinterpret_cast<
+                    uint64_t __attribute__((address_space(static_cast<int>(AddressSpace::Shared))))*>(ptr);
+  uint64_t *dst_u64 = reinterpret_cast<uint64_t *>(dst);
+  *dst_u64 = *smem_ptr;
+}
+
+/// ld.shared - 128b
+template <>
+MUTLASS_DEVICE
+void shared_load<16>(void *dst, uint32_t ptr) {
+  using Type = uint32_t __attribute__((vector_size(16)));
+  auto smem_ptr = reinterpret_cast<
+                    Type __attribute__((address_space(static_cast<int>(AddressSpace::Shared))))*>(ptr);
+  Type *dst_u128 = reinterpret_cast<Type *>(dst);
+  *dst_u128 = *smem_ptr;
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// st.shared
+template <int Bytes>
+MUTLASS_DEVICE
+void shared_store(uint32_t ptr, void const *src);
+
+/// st.shared - 16b
+template <>
+MUTLASS_DEVICE
+void shared_store<2>(uint32_t ptr, void const *src) {
+  auto smem_ptr = reinterpret_cast<
+                    uint16_t __attribute__((address_space(static_cast<int>(AddressSpace::Shared)))*>(ptr);
+  *smem_ptr = *reinterpret_cast<uint16_t const*>(src);
+}
+
+/// st.shared - 32b
+template <>
+MUTLASS_DEVICE
+void shared_store<4>(uint32_t ptr, void const *src) {
+  auto smem_ptr = reinterpret_cast<
+                    uint32_t __attribute__((address_space(static_cast<int>(AddressSpace::Shared))))*>(ptr);
+  *smem_ptr = *reinterpret_cast<uint32_t const*>(src);
+}
+
+/// st.shared - 64b
+template <>
+MUTLASS_DEVICE
+void shared_store<8>(uint32_t ptr, void const *src) {
+  auto smem_ptr = reinterpret_cast<
+                    uint64_t __attribute__((address_space(static_cast<int>(AddressSpace::Shared))))*>(ptr);
+  *smem_ptr = *reinterpret_cast<uint64_t const*>(src);
+}
+
+/// st.shared - 128b
+template <>
+MUTLASS_DEVICE
+void shared_store<16>(uint32_t ptr, void const *src) {
+  using Type = uint32_t __attribute__((vector_size(16)));
+  auto smem_ptr = reinterpret_cast<
+                    Type __attribute__((address_space(static_cast<int>(AddressSpace::Shared))))*>(ptr);
+  *smem_ptr = *reinterpret_cast<Type const*>(src);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace arch
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/arch/mma.h b/vllm_musa/csrc/flash_attn/include/mutlass/arch/mma.h
new file mode 100644
index 000000000..5d12d23b9
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/arch/mma.h
@@ -0,0 +1,136 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Templates exposing architecture support for multiply-add operations
+*/
+
+#pragma once
+
+#include "mutlass/array.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/functional.h"
+
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/arch/arch.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace arch {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the operation implied by MMA.
+struct OpMultiplyAdd {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the result is saturated to MAX_FLOAT|MIN_FLOAT or MAX_INT|MIN_INT
+struct OpMultiplyAddSaturate {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the input is converted to a narrower type (BF16)
+struct OpMultiplyAddFastBF16 {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the input is converted to a narrower type (F16)
+struct OpMultiplyAddFastF16 {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the input data types are mixed and the narrower type is
+/// upcasted to the wider type
+struct OpMultiplyAddMixedInputUpcast {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the input is converted to 2 (big and small) TF32 components
+//  Perform 3xTF32 or 4xTF32 for every F32 output element
+struct OpMultiplyAddFastF32 {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the input is converted to 2 (big and small) TF32 components
+//  Perform 3xTF32 or 4xTF32 for every complex<F32> output element
+struct OpMultiplyAddComplexFastF32 {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper for determining whether staged accumulation should be used for a given operator
+template <typename Operator>
+struct UseStagedAccumulation {
+  static bool const value = platform::is_same<Operator, OpMultiplyAddFastF32>::value ||
+                            platform::is_same<Operator, OpMultiplyAddComplexFastF32>::value;
+};
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the complex multiply-add operation
+struct OpMultiplyAddComplex {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the gaussian complex multiply-add operation
+struct OpMultiplyAddGaussianComplex {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the inner product is defined by (XOR, POPC)
+struct OpXorPopc {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag indicating the inner product is defined by (AND, POPC)
+struct OpAndPopc {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag classifying math operators as thread-level operations.
+struct OpClassSimt {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag classifying operators as Tensor Core operations.
+struct OpClassTensorOp {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag classifying operators as Tensor Core with structure sparse operations.
+struct OpClassSparseTensorOp {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace arch
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/arch/simd.h b/vllm_musa/csrc/flash_attn/include/mutlass/arch/simd.h
new file mode 100644
index 000000000..ab88b82c3
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/arch/simd.h
@@ -0,0 +1,121 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Templates exposing SIMD operators
+*/
+
+#pragma once
+
+#include "../array.h"
+#include "../numeric_types.h"
+
+namespace mutlass {
+namespace arch {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// Element-wise operators
+//
+
+MUTLASS_HOST_DEVICE
+template <typename T, int N>
+Array<T, N> operator*(Array<T, N> const &a, Array<T, N> const &b) {
+  Array<T, N> d;
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < N; ++i) {
+    d[i] = a[i] * b[i];
+  }
+  return d;
+}
+
+MUTLASS_HOST_DEVICE
+template <typename T, int N>
+Array<T, N> operator+(Array<T, N> const &a, Array<T, N> const &b) {
+  Array<T, N> d;
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < N; ++i) {
+    d[i] = a[i] + b[i];
+  }
+  return d;
+}
+
+MUTLASS_HOST_DEVICE
+template <typename T, int N>
+Array<T, N> operator-(Array<T, N> const &a, Array<T, N> const &b) {
+  Array<T, N> d;
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < N; ++i) {
+    d[i] = a[i] - b[i];
+  }
+  return d;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// Multiply-accumulate operators
+//
+
+MUTLASS_HOST_DEVICE
+template <typename T, int N>
+Array<T, N> mac(Array<T, N> const &a, Array<T, N> const &b, Array<T, N> const &c) {
+  Array<T, N> d;
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < N; ++i) {
+    d[i] = a[i] * b[i] + c[i];
+  }
+  return d;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// Dot product operator
+//
+
+MUTLASS_HOST_DEVICE
+template <typename Element, typename Accumulator, int N>
+Accumulator dot(Array<T, N> const &a, Array<T, N> const &b, Accumulator accum) {
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < N; ++i) {
+    accum += a[i] * b[i];
+  }
+  return accum;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace arch
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/array.h b/vllm_musa/csrc/flash_attn/include/mutlass/array.h
new file mode 100644
index 000000000..15e60df51
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/array.h
@@ -0,0 +1,2180 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Statically sized array of elements that accommodates all MUTLASS-supported numeric types
+           and is safe to use in a union.
+*/
+
+#pragma once
+#include "mutlass/mutlass.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_types.h"
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Statically sized array for any data type
+template <
+  typename T,
+  int N,
+  bool RegisterSized = sizeof_bits<T>::value >= 32
+>
+struct Array;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines the size of an Array<> in bits
+template <typename T, int N, bool RegisterSized>
+struct sizeof_bits<Array<T, N, RegisterSized> > {
+  static constexpr int value = sizeof(Array<T, N, RegisterSized>) * 8;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns true if the argument is a power of 2
+MUTLASS_HOST_DEVICE
+constexpr bool ispow2(unsigned x) {
+  return x && (!(x & (x - 1)));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns the largest power of two not greater than the argument.
+MUTLASS_HOST_DEVICE
+constexpr unsigned floor_pow_2(unsigned x) {
+  return (x == 0 || ispow2(x)) ? x : ((floor_pow_2(x >> 1)) << 1);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Statically sized array for any data type
+template <
+  typename T,
+  int N
+>
+struct Array<T, N, true> {
+
+  /// Storage type
+  using Storage = T;
+
+  /// Element type
+  using Element = T;
+
+  /// Number of storage elements
+  //static std::size_t const kStorageElements = N;
+  static constexpr size_t kStorageElements = N;
+
+  /// Number of logical elements
+  static constexpr size_t kElements = N;
+
+  //
+  // C++ standard members
+  //
+
+  typedef T value_type;
+  typedef size_t size_type;
+  typedef ptrdiff_t difference_type;
+  typedef value_type &reference;
+  typedef value_type const & const_reference;
+  typedef value_type *pointer;
+  typedef value_type const * const_pointer;
+
+  //
+  // Iterators
+  //
+
+  /// Bidirectional iterator over elements
+  class iterator {
+
+    /// Pointer to object
+    T *ptr_;
+
+  public:
+
+    MUTLASS_HOST_DEVICE
+    iterator(): ptr_(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    iterator(T *_ptr): ptr_(_ptr) { }
+
+    MUTLASS_HOST_DEVICE
+    iterator &operator++() {
+      ++ptr_;
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator &operator--() {
+      --ptr_;
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator operator++(int) {
+      iterator ret(*this);
+      ++ptr_;
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator operator--(int) {
+      iterator ret(*this);
+      --ptr_;
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    T &operator*() const {
+      return *ptr_;
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator==(iterator const &other) const {
+      return ptr_ == other.ptr_;
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator!=(iterator const &other) const {
+      return ptr_ != other.ptr_;
+    }
+  };
+
+  /// Bidirectional constant iterator over elements
+  class const_iterator {
+
+    /// Pointer to object
+    const T *ptr_;
+
+  public:
+
+    MUTLASS_HOST_DEVICE
+    const_iterator(): ptr_(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    const_iterator(T const *_ptr): ptr_(_ptr) { }
+
+    MUTLASS_HOST_DEVICE
+    const_iterator &operator++() {
+      ++ptr_;
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    const_iterator &operator--() {
+      --ptr_;
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    const_iterator operator++(int) {
+      const_iterator ret(*this);
+      ++ptr_;
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    const_iterator operator--(int) {
+      const_iterator ret(*this);
+      --ptr_;
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    T const &operator*() const {
+      return *ptr_;
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator==(const_iterator const &other) const {
+      return ptr_ == other.ptr_;
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator!=(const_iterator const &other) const {
+      return ptr_ != other.ptr_;
+    }
+  };
+
+  /// Bidirectional iterator over elements
+  class reverse_iterator {
+
+    /// Pointer to object
+    T *ptr_;
+
+  public:
+
+    MUTLASS_HOST_DEVICE
+    reverse_iterator(): ptr_(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    reverse_iterator(T *_ptr): ptr_(_ptr) { }
+
+    MUTLASS_HOST_DEVICE
+    reverse_iterator &operator++() {
+      --ptr_;
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    reverse_iterator &operator--() {
+      ++ptr_;
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    reverse_iterator operator++(int) {
+      iterator ret(*this);
+      --ptr_;
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    reverse_iterator operator--(int) {
+      iterator ret(*this);
+      ++ptr_;
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    T &operator*() const {
+      return *(ptr_ - 1);
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator==(reverse_iterator const &other) const {
+      return ptr_ == other.ptr_;
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator!=(reverse_iterator const &other) const {
+      return ptr_ != other.ptr_;
+    }
+  };
+
+  /// Bidirectional constant iterator over elements
+  class const_reverse_iterator {
+
+    /// Pointer to object
+    T const *ptr_;
+
+  public:
+
+    MUTLASS_HOST_DEVICE
+    const_reverse_iterator(): ptr_(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    const_reverse_iterator(T const *_ptr): ptr_(_ptr) { }
+
+    MUTLASS_HOST_DEVICE
+    const_reverse_iterator &operator++() {
+      --ptr_;
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    const_reverse_iterator &operator--() {
+      ++ptr_;
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    const_reverse_iterator operator++(int) {
+      const_reverse_iterator ret(*this);
+      --ptr_;
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    const_reverse_iterator operator--(int) {
+      const_reverse_iterator ret(*this);
+      ++ptr_;
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    T const &operator*() const {
+      return *(ptr_ - 1);
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator==(const_iterator const &other) const {
+      return ptr_ == other.ptr_;
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator!=(const_iterator const &other) const {
+      return ptr_ != other.ptr_;
+    }
+  };
+
+  /// Internal storage
+  Storage storage[kElements];
+
+  /// Efficient clear method
+  MUTLASS_HOST_DEVICE
+  void clear() {
+    fill(T(0));
+  }
+
+  MUTLASS_HOST_DEVICE
+  reference at(size_type pos) {
+    return reinterpret_cast<reference>(storage[pos]);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reference at(size_type pos) const {
+    return reinterpret_cast<const_reference>(storage[pos]);
+  }
+
+  MUTLASS_HOST_DEVICE
+  reference operator[](size_type pos) {
+    return reinterpret_cast<reference>(storage[pos]);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reference operator[](size_type pos) const {
+    return reinterpret_cast<const_reference>(storage[pos]);
+  }
+
+  MUTLASS_HOST_DEVICE
+  reference front() {
+    return reinterpret_cast<reference>(storage[0]);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reference front() const {
+    return reinterpret_cast<const_reference>(storage[0]);
+  }
+
+  MUTLASS_HOST_DEVICE
+  reference back() {
+    return reinterpret_cast<reference>(storage[kStorageElements - 1]);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reference back() const {
+    return reinterpret_cast<const_reference>(storage[kStorageElements - 1]);
+  }
+
+  MUTLASS_HOST_DEVICE
+  pointer data() {
+    return reinterpret_cast<pointer>(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_pointer data() const {
+    return reinterpret_cast<const_pointer>(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  pointer raw_data() {
+    return reinterpret_cast<pointer>(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_pointer raw_data() const {
+    return reinterpret_cast<const_pointer>(storage);
+  }
+
+
+  MUTLASS_HOST_DEVICE
+  constexpr bool empty() const {
+    return !kElements;
+  }
+
+  MUTLASS_HOST_DEVICE
+  constexpr size_type size() const {
+    return kElements;
+  }
+
+  MUTLASS_HOST_DEVICE
+  constexpr size_type max_size() const {
+    return kElements;
+  }
+
+  MUTLASS_HOST_DEVICE
+  void fill(T const &value) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < int(kElements); ++i) {
+      storage[i] = static_cast<Storage>(value);
+    }
+  }
+
+  MUTLASS_HOST_DEVICE
+  iterator begin() {
+    return iterator(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_iterator begin() const {
+    return cbegin();
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_iterator cbegin() const {
+    return const_iterator(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  iterator end() {
+    return iterator(reinterpret_cast<pointer>(storage + kStorageElements));
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_iterator end() const {
+    return cend();
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_iterator cend() const {
+    return const_iterator(reinterpret_cast<const_pointer>(storage + kStorageElements));
+  }
+
+  MUTLASS_HOST_DEVICE
+  reverse_iterator rbegin() {
+    return reverse_iterator(reinterpret_cast<pointer>(storage + kStorageElements));
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reverse_iterator rbegin() const {
+    return crbegin();
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reverse_iterator crbegin() const {
+    return const_reverse_iterator(reinterpret_cast<const_pointer>(storage + kStorageElements));
+  }
+
+  MUTLASS_HOST_DEVICE
+  reverse_iterator rend() {
+    return reverse_iterator(reinterpret_cast<pointer>(storage));
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reverse_iterator rend() const {
+    return crend();
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reverse_iterator crend() const {
+    return const_reverse_iterator(reinterpret_cast<const_pointer>(storage));
+  }
+
+  //
+  // Comparison operators
+  //
+
+};
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// Factories
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Array<Element, 1> make_Array(Element x) {
+  return {x};
+}
+
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Array<Element, 2> make_Array(Element x, Element y) {
+  return {x,y};
+}
+
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Array<Element, 3> make_Array(Element x, Element y, Element z) {
+  return {x,y,z};
+}
+
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Array<Element, 4> make_Array(Element x, Element y, Element z, Element w) {
+  return {x,y,z,w};
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// functional.h numeric specializations
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T, int N>
+struct absolute_value_op< Array<T, N> > {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs) const {
+
+    Array<T, N> result;
+    absolute_value_op<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct plus<Array<T, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    plus<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], rhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
+
+    Array<T, N> result;
+    plus<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], scalar);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()( T const &scalar, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    plus<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, rhs[i]);
+    }
+
+    return result;
+  }
+};
+template <typename T, int N>
+struct minus<Array<T, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    minus<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], rhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
+
+    Array<T, N> result;
+    minus<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], scalar);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()( T const &scalar, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    minus<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, rhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct multiplies<Array<T, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    multiplies<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], rhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
+
+    Array<T, N> result;
+    multiplies<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], scalar);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()( T const &scalar, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    multiplies<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, rhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N, bool PropogateNaN>
+struct maximum_absolute_value_reduction<Array<T, N>, PropogateNaN> {
+
+  MUTLASS_HOST_DEVICE
+  T operator() (T const& scalar, Array<T, N> const& rhs) const {
+
+    T result = scalar;
+    maximum_absolute_value_reduction<T, PropogateNaN> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result = scalar_op(result, rhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct scale<Array<T, N>> {
+  T const scaling_factor_;
+
+  MUTLASS_HOST_DEVICE
+  scale(T scaling_factor) : scaling_factor_(scaling_factor) {
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const & rhs) const {
+    Array<T, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = rhs[i] * scaling_factor_;
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct divides<Array<T, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    divides<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], rhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
+
+    Array<T, N> result;
+    divides<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], scalar);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()( T const &scalar, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    divides<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, rhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct reciprocal_approximate<Array<T, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs) const {
+
+    Array<T, N> result;
+    reciprocal_approximate<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct maximum<Array<T, N>, false> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    maximum<T, false> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], rhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
+
+    Array<T, N> result;
+    maximum<T, false> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], scalar);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    maximum<T, false> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, rhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct maximum<Array<T, N>, true> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    maximum<T, true> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], rhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
+
+    Array<T, N> result;
+    maximum<T, true> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], scalar);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    maximum<T, true> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, rhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct minimum<Array<T, N>, false> {
+
+  MUTLASS_HOST_DEVICE
+  static T scalar_op(T const &lhs, T const &rhs) {
+    return (rhs < lhs ? rhs : lhs);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    minimum<T, false> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], rhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
+
+    Array<T, N> result;
+    minimum<T, false> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], scalar);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    minimum<T, false> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, rhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct minimum<Array<T, N>, true> {
+
+  MUTLASS_HOST_DEVICE
+  static T scalar_op(T const &lhs, T const &rhs) {
+    return (rhs < lhs ? rhs : lhs);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    minimum<T, true> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], rhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &scalar) const {
+
+    Array<T, N> result;
+    minimum<T, true> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i], scalar);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(T const &scalar, Array<T, N> const &rhs) const {
+
+    Array<T, N> result;
+    minimum<T, true> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, rhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <typename T, int N>
+struct negate<Array<T, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs) const {
+
+    Array<T, N> result;
+    negate<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(lhs[i]);
+    }
+
+    return result;
+  }
+};
+
+/// Fused multiply-add
+template <typename T, int N>
+struct multiply_add<Array<T, N>, Array<T, N>, Array<T, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &a, Array<T, N> const &b, Array<T, N> const &c) const {
+
+    Array<T, N> result;
+    multiply_add<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(a[i], b[i], c[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &a, T const &scalar, Array<T, N> const &c) const {
+
+    Array<T, N> result;
+    multiply_add<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(a[i], scalar, c[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(T const &scalar, Array<T, N> const &b, Array<T, N> const &c) const {
+
+    Array<T, N> result;
+    multiply_add<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(scalar, b[i], c[i]);
+    }
+
+    return result;
+  }
+};
+
+/// Fused square-and-plus
+template <typename T, int N>
+struct square_and_plus<Array<T, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, Array<T, N> const &rhs) const {
+    multiply_add<Array<T, N>, Array<T, N>, Array<T, N>> ma_op;
+    return ma_op(rhs, rhs, lhs);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &lhs, T const &rhs) const {
+    plus<Array<T, N>> plus_op;
+    multiplies<T> multiplies_op;
+    return plus_op(multiplies_op(rhs, rhs), lhs);
+  }
+};
+
+/// Inverse-square-root
+template <typename T, int N>
+struct inverse_square_root<Array<T, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &a) const {
+    Array<T, N> result;
+    inverse_square_root<T> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(a[i]);
+    }
+    return result;
+  }
+};
+
+template <int N>
+struct inverse_square_root<Array<half_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & a) const {
+    Array<half_t, N> result;
+
+    inverse_square_root<half_t> scalar_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = scalar_op(a[i]);
+    }
+
+    return result;
+  }
+};
+
+/// Fused multiply-add-relu0
+template <typename T, int N>
+struct multiply_add_relu0<Array<T, N>, Array<T, N>, Array<T, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &a, Array<T, N> const &b, Array<T, N> const &c) const {
+
+    Array<T, N> result;
+    multiply_add<T> scalar_op;
+    maximum<T> mx;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = mx(scalar_op(a[i], b[i], c[i]), T(0));
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &a, T const &scalar, Array<T, N> const &c) const {
+
+    Array<T, N> result;
+    multiply_add<T> scalar_op;
+    maximum<T> mx;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = mx(scalar_op(a[i], scalar, c[i]), T(0));
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(T const &scalar, Array<T, N> const &b, Array<T, N> const &c) const {
+
+    Array<T, N> result;
+    multiply_add<T> scalar_op;
+    maximum<T> mx;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = mx(scalar_op(scalar, b[i], c[i]), T(0));
+    }
+
+    return result;
+  }
+};
+
+
+template <typename T, int N>
+struct conjugate<Array<T, N> >  {
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &a) const {
+
+    conjugate<T> conj_op;
+
+    Array<T, N> ca;
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      ca[i] = conj_op(a[i]);
+    }
+    return ca;
+  }
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// functional.h numeric specializations targeting SIMD instructions in device code.
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+struct plus<Array<half_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs);
+    __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hadd2(lhs_ptr[i], rhs_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
+      __half d_residual = __hadd(a_residual_ptr[N - 1], b_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs[i] + rhs[i];
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(half_t const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs));
+    __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hadd2(lhs_pair, rhs_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
+      __half d_residual = __hadd(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs + rhs[i];
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, half_t const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs);
+    __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs));
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hadd2(lhs_ptr[i], rhs_pair);
+    }
+
+    if constexpr (N % 2) {
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
+      __half d_residual = __hadd(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs));
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs[i] + rhs;
+    }
+    #endif
+
+    return result;
+  }
+};
+
+template <int N>
+struct minus<Array<half_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs);
+    __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hsub2(lhs_ptr[i], rhs_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
+      __half d_residual = __hsub(a_residual_ptr[N - 1], b_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs[i] - rhs[i];
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(half_t const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs));
+    __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hsub2(lhs_pair, rhs_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
+      __half d_residual = __hsub(reinterpret_cast<__half const &>(lhs), b_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs - rhs[i];
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, half_t const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs);
+    __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs));
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hsub2(lhs_ptr[i], rhs_pair);
+    }
+
+    if constexpr (N % 2) {
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
+      __half d_residual = __hsub(a_residual_ptr[N - 1], reinterpret_cast<__half const &>(rhs));
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs[i] - rhs;
+    }
+    #endif
+
+    return result;
+  }
+};
+
+template <int N>
+struct multiplies<Array<half_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs);
+    __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hmul2(lhs_ptr[i], rhs_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
+      __half d_residual = __hmul(a_residual_ptr[N - 1], b_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs[i] * rhs[i];
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(half_t const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 lhs_pair = __half2half2(reinterpret_cast<__half const &>(lhs));
+    __half2 const *rhs_ptr = reinterpret_cast<__half2 const *>(&rhs);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hmul2(lhs_pair, rhs_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&rhs);
+
+      __half d_residual = __hmul(
+        reinterpret_cast<__half const &>(lhs),
+        b_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs * rhs[i];
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, half_t const &rhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *lhs_ptr = reinterpret_cast<__half2 const *>(&lhs);
+    __half2 rhs_pair = __half2half2(reinterpret_cast<__half const &>(rhs));
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hmul2(lhs_ptr[i], rhs_pair);
+    }
+
+    if constexpr (N % 2) {
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&lhs);
+
+      __half d_residual = __hmul(
+        a_residual_ptr[N - 1],
+        reinterpret_cast<__half const &>(rhs));
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs[i] * rhs;
+    }
+    #endif
+
+    return result;
+  }
+};
+
+template <int N>
+struct divides<Array<half_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs[i] / rhs[i];
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(half_t const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs / rhs[i];
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, half_t const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = lhs[i] / rhs;
+    }
+
+    return result;
+  }
+};
+
+template <int N>
+struct negate<Array<half_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs) const {
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *source_ptr = reinterpret_cast<__half2 const *>(&lhs);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hneg2(source_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+      half_t x = -lhs[N - 1];
+      __half lhs_val = reinterpret_cast<__half const &>(x);
+      result[N - 1] = reinterpret_cast<half_t const &>(lhs_val);
+    }
+
+    #else
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = -lhs[i];
+    }
+    #endif
+
+    return result;
+  }
+};
+
+/// Fused multiply-add
+template <int N>
+struct multiply_add<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(
+    Array<half_t, N> const &a,
+    Array<half_t, N> const &b,
+    Array<half_t, N> const &c) const {
+
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a);
+    __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b);
+    __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hfma2(a_ptr[i], b_ptr[i], c_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
+      __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c);
+
+      __half d_residual = __hfma(
+        a_residual_ptr[N - 1],
+        b_residual_ptr[N - 1],
+        c_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    multiply_add<half_t> op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = op(a[i], b[i], c[i]);
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(
+    half_t const &a,
+    Array<half_t, N> const &b,
+    Array<half_t, N> const &c) const {
+
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 a_pair = __half2half2(reinterpret_cast<__half const &>(a));
+    __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b);
+    __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hfma2(a_pair, b_ptr[i], c_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
+      __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c);
+      __half d_residual = __hfma(
+        reinterpret_cast<__half const &>(a),
+        b_residual_ptr[N - 1],
+        c_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    multiply_add<half_t> op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = op(a, b[i], c[i]);
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(
+    Array<half_t, N> const &a,
+    half_t const &b,
+    Array<half_t, N> const &c) const {
+
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a);
+    __half2 b_pair = __half2half2(reinterpret_cast<__half const &>(b));
+    __half2 const *c_ptr = reinterpret_cast<__half2 const *>(&c);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hfma2(a_ptr[i], b_pair, c_ptr[i]);
+    }
+
+    if constexpr (N % 2) {
+
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
+      __half const *c_residual_ptr = reinterpret_cast<__half const *>(&c);
+
+      __half d_residual = __hfma(
+        a_residual_ptr[N - 1],
+        reinterpret_cast<__half const &>(b),
+        c_residual_ptr[N - 1]);
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    multiply_add<half_t> op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = op(a[i], b, c[i]);
+    }
+    #endif
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(
+    Array<half_t, N> const &a,
+    Array<half_t, N> const &b,
+    half_t const &c) const {
+
+    Array<half_t, N> result;
+    #if defined(__MUSA_ARCH__)
+
+    __half2 *result_ptr = reinterpret_cast<__half2 *>(&result);
+    __half2 const *a_ptr = reinterpret_cast<__half2 const *>(&a);
+    __half2 const *b_ptr = reinterpret_cast<__half2 const *>(&b);
+    __half2 c_pair = __half2half2(reinterpret_cast<__half const &>(c));
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = __hfma2(a_ptr[i], b_ptr[i], c_pair);
+    }
+
+    if constexpr (N % 2) {
+
+      __half const *a_residual_ptr = reinterpret_cast<__half const *>(&a);
+      __half const *b_residual_ptr = reinterpret_cast<__half const *>(&b);
+
+      __half d_residual = __hfma(
+        a_residual_ptr[N - 1],
+        b_residual_ptr[N - 1],
+        reinterpret_cast<__half const &>(c));
+
+      result[N - 1] = reinterpret_cast<half_t const &>(d_residual);
+    }
+
+    #else
+
+    multiply_add<half_t> op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = op(a[i], b[i], c);
+    }
+    #endif
+
+    return result;
+  }
+};
+
+/// Fused multiply-add-relu0
+template <int N>
+struct multiply_add_relu0<Array<half_t, N>, Array<half_t, N>, Array<half_t, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(
+    Array<half_t, N> const &a,
+    Array<half_t, N> const &b,
+    Array<half_t, N> const &c) const {
+
+    Array<half_t, N> result;
+
+    multiply_add<half_t> op;
+    maximum<half_t> mx;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = mx(op(a[i], b[i], c[i]), (half_t)0);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(
+    half_t const &a,
+    Array<half_t, N> const &b,
+    Array<half_t, N> const &c) const {
+
+    Array<half_t, N> result;
+
+    multiply_add<half_t> op;
+    maximum<half_t> mx;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = mx(op(a, b[i], c[i]), half_t(0));
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(
+    Array<half_t, N> const &a,
+    half_t const &b,
+    Array<half_t, N> const &c) const {
+
+    Array<half_t, N> result;
+
+    multiply_add<half_t> op;
+    maximum<half_t> mx;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = mx(op(a[i], b, c[i]), half_t(0));
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(
+    Array<half_t, N> const &a,
+    Array<half_t, N> const &b,
+    half_t const &c) const {
+
+    Array<half_t, N> result;
+
+    multiply_add<half_t> op;
+    maximum<half_t> mx;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = mx(op(a[i], b[i], c), half_t(0));
+    }
+
+    return result;
+  }
+};
+
+template <int N>
+struct minimum<Array<half_t, N>, false> {
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = (rhs[i] < lhs[i] ? rhs[i] : lhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(half_t const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = (rhs[i] < lhs ? rhs[i] : lhs);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, half_t const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = (rhs < lhs[i] ? rhs : lhs[i]);
+    }
+
+    return result;
+  }
+};
+
+template <int N>
+struct maximum<Array<half_t, N>, false> {
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = (lhs[i] < rhs[i] ? rhs[i] : lhs[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(half_t const & lhs, Array<half_t, N> const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = (lhs < rhs[i] ? rhs[i] : lhs);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const & lhs, half_t const &rhs) const {
+    Array<half_t, N> result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = (lhs[i] < rhs ? rhs : lhs[i]);
+    }
+
+    return result;
+  }
+};
+
+/// Fused multiply-add
+template <int N>
+struct multiply_add<Array<bfloat16_t, N>, Array<bfloat16_t, N>, Array<bfloat16_t, N>> {
+
+  MUTLASS_HOST_DEVICE
+  Array<bfloat16_t, N> operator()(
+    Array<bfloat16_t, N> const &a,
+    Array<bfloat16_t, N> const &b,
+    Array<bfloat16_t, N> const &c) const {
+
+    Array<bfloat16_t, N> result;
+
+    multiply_add<bfloat16_t> op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = op(a[i], b[i], c[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<bfloat16_t, N> operator()(
+    bfloat16_t const &a,
+    Array<bfloat16_t, N> const &b,
+    Array<bfloat16_t, N> const &c) const {
+
+    Array<bfloat16_t, N> result;
+
+    multiply_add<bfloat16_t> op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = op(a, b[i], c[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<bfloat16_t, N> operator()(
+    Array<bfloat16_t, N> const &a,
+    bfloat16_t const &b,
+    Array<bfloat16_t, N> const &c) const {
+
+    Array<bfloat16_t, N> result;
+
+    multiply_add<bfloat16_t> op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = op(a[i], b, c[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<bfloat16_t, N> operator()(
+    Array<bfloat16_t, N> const &a,
+    Array<bfloat16_t, N> const &b,
+    bfloat16_t const &c) const {
+
+    Array<bfloat16_t, N> result;
+
+    multiply_add<bfloat16_t> op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      result[i] = op(a[i], b[i], c);
+    }
+
+    return result;
+  }
+};
+
+
+/// bit_and
+template <int N>
+struct bit_and<Array<uint1b_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<uint1b_t, N> operator()(Array<uint1b_t, N> const &a, Array<uint1b_t, N> const &b) const {
+    using ArrayType = Array<uint1b_t, N>;
+    using Storage = typename ArrayType::Storage;
+    ArrayType result;
+
+    Storage *result_data = result.raw_data();
+    Storage const *a_data = a.raw_data();
+    Storage const *b_data = b.raw_data();
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < ArrayType::kStorageElements; ++i) {
+      result_data[i] = (a_data[i] & b_data[i]);
+    }
+
+    return result;
+  }
+};
+
+
+/// bit_or
+template <int N>
+struct bit_or<Array<uint1b_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<uint1b_t, N> operator()(Array<uint1b_t, N> const &a, Array<uint1b_t, N> const &b) const {
+    using ArrayType = Array<uint1b_t, N>;
+    using Storage = typename ArrayType::Storage;
+    ArrayType result;
+
+    Storage *result_data = result.raw_data();
+    Storage const *a_data = a.raw_data();
+    Storage const *b_data = b.raw_data();
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < ArrayType::kStorageElements; ++i) {
+      result_data[i] = (a_data[i] | b_data[i]);
+    }
+
+    return result;
+  }
+};
+
+
+/// bit_not
+template <int N>
+struct bit_not<Array<uint1b_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<uint1b_t, N> operator()(Array<uint1b_t, N> const &a) const {
+    using ArrayType = Array<uint1b_t, N>;
+    using Storage = typename ArrayType::Storage;
+    ArrayType result;
+
+    Storage *result_data = result.raw_data();
+    Storage const *a_data = a.raw_data();
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < ArrayType::kStorageElements; ++i) {
+      result_data[i] = (~a_data[i]);
+    }
+
+    return result;
+  }
+};
+
+
+/// bit_xor
+template <int N>
+struct bit_xor<Array<uint1b_t, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<uint1b_t, N> operator()(Array<uint1b_t, N> const &a, Array<uint1b_t, N> const &b) const {
+    using ArrayType = Array<uint1b_t, N>;
+    using Storage = typename ArrayType::Storage;
+    ArrayType result;
+
+    Storage *result_data = result.raw_data();
+    Storage const *a_data = a.raw_data();
+    Storage const *b_data = b.raw_data();
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < ArrayType::kStorageElements; ++i) {
+      result_data[i] = (a_data[i] ^ b_data[i]);
+    }
+
+    return result;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// Operator overloads
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator+(Array<T, N> const &lhs, Array<T, N> const &rhs) {
+  plus<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator+(T const &lhs, Array<T, N> const &rhs) {
+  plus<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator+(Array<T, N> const &lhs, T const &rhs) {
+  plus<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator-(Array<T, N> const &lhs, Array<T, N> const &rhs) {
+  minus<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator-(Array<T, N> const &lhs) {
+  negate<Array<T, N>> op;
+  return op(lhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator*(Array<T, N> const &lhs, Array<T, N> const &rhs) {
+  multiplies<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator*(T lhs, Array<T, N> const &rhs) {
+  multiplies<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator*(Array<T, N> const &lhs, T rhs) {
+  multiplies<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> operator/(Array<T, N> const &lhs, Array<T, N> const &rhs) {
+  divides<Array<T, N>> op;
+  return op(lhs, rhs);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> fma(Array<T, N> const &a, Array<T, N> const &b, Array<T, N> const &c) {
+  multiply_add<Array<T, N>> op;
+  return op(a, b, c);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> fma(T a, Array<T, N> const &b, Array<T, N> const &c) {
+  multiply_add<Array<T, N>> op;
+  return op(a, b, c);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> fma(Array<T, N> const &a, T b, Array<T, N> const &c) {
+  multiply_add<Array<T, N>> op;
+  return op(a, b, c);
+}
+
+template <typename T, int N>
+MUTLASS_HOST_DEVICE
+Array<T, N> fma(Array<T, N> const &a, Array<T, N> const &b, T c) {
+  multiply_add<Array<T, N>> op;
+  return op(a, b, c);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// AlignedArray
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Aligned array type
+template <
+  /// Element type
+  typename T,
+  /// Number of elements in the array
+  int N,
+  /// Alignment requirement in bytes
+  int Alignment = ( sizeof_bits<T>::value * N + 7 ) / 8
+>
+class alignas(Alignment) AlignedArray: public Array<T, N> {
+public:
+
+};
+
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "mutlass/array_subbyte.h"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/array_planar_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/array_planar_complex.h
new file mode 100644
index 000000000..687577cac
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/array_planar_complex.h
@@ -0,0 +1,90 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Templates implementing warp-level matrix multiply-accumulate operations.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Array holding planar complex elements
+template <typename Element_, int N>
+struct ArrayPlanarComplex {
+
+  /// Underlying real element
+  using Element = Element_;
+
+  /// Number of logical elements
+  static constexpr size_t kElements = N;
+
+  /// Underlying Fragment of real-valued elemenets
+  using ArrayReal = mutlass::Array<Element, N>;
+
+public:
+  /// Fragment of real-valued elements representing the real part
+  ArrayReal real;
+
+  /// Fragment of real-valued elements representing the imaginary part
+  ArrayReal imag;
+
+public:
+  /// Sets the array to zero efficiently
+  MUTLASS_HOST_DEVICE
+  void clear() {
+    real.clear();
+    imag.clear();
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper to deduce template arguments
+template <typename Element, int N>
+MUTLASS_HOST_DEVICE
+ArrayPlanarComplex<Element, N>
+make_ArrayPlanarComplex(Array<Element, N> const &real, Array<Element, N> const &imag) {
+  return ArrayPlanarComplex<Element, N>{real, imag};
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/array_subbyte.h b/vllm_musa/csrc/flash_attn/include/mutlass/array_subbyte.h
new file mode 100644
index 000000000..2279f1593
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/array_subbyte.h
@@ -0,0 +1,560 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Statically sized array of elements that accommodates all MUTLASS-supported numeric types
+           and is safe to use in a union.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+#include "mutlass/platform/platform.h"
+
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Statically sized array for any data type
+template <
+  typename T,
+  int N
+>
+struct Array<T, N, false> {
+  static constexpr int kSizeBits = sizeof_bits<T>::value * N;
+
+  /// Storage type
+  using Storage = typename platform::conditional<
+    ((kSizeBits % 32) != 0),
+    typename platform::conditional<
+      ((kSizeBits % 16) != 0),
+      uint8_t,
+      uint16_t
+    >::type,
+    uint32_t
+  >::type;
+
+  /// Element type
+  using Element = T;
+
+  /// Number of logical elements per stored object
+  static constexpr int kElementsPerStoredItem = int(sizeof(Storage) * 8) / sizeof_bits<T>::value;
+
+  /// Number of storage elements
+  static constexpr size_t kStorageElements = (N + kElementsPerStoredItem - 1) / kElementsPerStoredItem;
+
+  /// Number of logical elements
+  static constexpr size_t kElements = N;
+
+  /// Bitmask for covering one item
+  static constexpr Storage kMask = ((Storage(1) << sizeof_bits<T>::value) - 1);
+
+  //
+  // C++ standard members with pointer types removed
+  //
+
+  typedef T value_type;
+  typedef size_t size_type;
+  typedef ptrdiff_t difference_type;
+  typedef value_type *pointer;
+  typedef value_type const *const_pointer;
+
+  //
+  // References
+  //
+
+  /// Reference object inserts or extracts sub-byte items
+  class reference {
+    /// Pointer to storage element
+    Storage *ptr_{nullptr};
+
+    /// Index into elements packed into Storage object
+    int idx_{0};
+
+  public:
+
+    reference() = default;
+
+    /// Ctor
+    MUTLASS_HOST_DEVICE
+    reference(Storage *ptr, int idx = 0): ptr_(ptr), idx_(idx) { }
+
+    /// Assignment
+    MUTLASS_HOST_DEVICE
+    reference &operator=(T x) {
+    // `*ptr_ & kUpdateMask` will read ptr_ before write to it
+    // This means code pattern like
+    //
+    // ```cpp
+    // Array<half_t, N> result;
+    // result[0] = xxx;
+    // ```
+    //
+    // Will leads to compiler warning on use of unintialized member variable. Although we know
+    //      this read of uninitialized member variable is harmeless.
+
+#if defined(__clang__)
+#  pragma clang diagnostic push
+#  pragma clang diagnostic ignored "-Wuninitialized"
+#elif defined(__GNUC__)
+#  pragma GCC diagnostic push
+#  pragma GCC diagnostic ignored "-Wuninitialized"
+#  pragma GCC diagnostic ignored "-Wmaybe-uninitialized"
+#endif
+
+      Storage item = (reinterpret_cast<Storage const &>(x) & kMask);
+
+      Storage kUpdateMask = Storage(~(kMask << (idx_ * sizeof_bits<T>::value)));
+
+      *ptr_ = Storage(((*ptr_ & kUpdateMask) | (item << idx_ * sizeof_bits<T>::value)));
+
+#if defined(__clang__)
+#  pragma clang diagnostic pop
+#elif defined(__GNUC__)
+#  pragma GCC diagnostic pop
+#endif
+
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    T get() const {
+      Storage item = Storage((*ptr_ >> (idx_ * sizeof_bits<T>::value)) & kMask);
+      return reinterpret_cast<T const &>(item);
+    }
+
+    /// Extract
+    MUTLASS_HOST_DEVICE
+    operator T() const {
+      return get();
+    }
+
+    /// Explicit cast to int
+    MUTLASS_HOST_DEVICE
+    explicit operator int() const {
+      return int(get());
+    }
+
+    /// Explicit cast to float
+    MUTLASS_HOST_DEVICE
+    explicit operator float() const {
+      return float(get());
+    }
+  };
+
+  /// Reference object extracts sub-byte items
+  class const_reference {
+
+    /// Pointer to storage element
+    Storage const *ptr_{nullptr};
+
+    /// Index into elements packed into Storage object
+    int idx_{0};
+
+  public:
+
+    const_reference() = default;
+
+    /// Ctor
+    MUTLASS_HOST_DEVICE
+    const_reference(Storage const *ptr, int idx = 0): ptr_(ptr), idx_(idx) { }
+
+    MUTLASS_HOST_DEVICE
+    const T get() const {
+      Storage item = (*ptr_ >> (idx_ * sizeof_bits<T>::value)) & kMask;
+      return reinterpret_cast<T const &>(item);
+    }
+
+    /// Extract
+    MUTLASS_HOST_DEVICE
+    operator T() const {
+      Storage item = Storage(Storage(*ptr_ >> Storage(idx_ * sizeof_bits<T>::value)) & kMask);
+      return reinterpret_cast<T const &>(item);
+    }
+
+    /// Explicit cast to int
+    MUTLASS_HOST_DEVICE
+    explicit operator int() const {
+      return int(get());
+    }
+
+    /// Explicit cast to float
+    MUTLASS_HOST_DEVICE
+    explicit operator float() const {
+      return float(get());
+    }
+  };
+
+  //
+  // Iterators
+  //
+
+  /// Bidirectional iterator over elements
+  class iterator {
+
+    /// Pointer to storage element
+    Storage *ptr_{nullptr};
+
+    /// Index into elements packed into Storage object
+    int idx_{0};
+
+  public:
+
+    iterator() = default;
+
+    MUTLASS_HOST_DEVICE
+    iterator(Storage *ptr, int idx = 0): ptr_(ptr), idx_(idx) { }
+
+    MUTLASS_HOST_DEVICE
+    iterator &operator++() {
+      ++idx_;
+      if (idx_ == kElementsPerStoredItem) {
+        ++ptr_;
+        idx_ = 0;
+      }
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator &operator--() {
+      if (!idx_) {
+        --ptr_;
+        idx_ = kElementsPerStoredItem - 1;
+      }
+      else {
+        --idx_;
+      }
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator operator++(int) {
+      iterator ret(*this);
+      ++idx_;
+      if (idx_ == kElementsPerStoredItem) {
+        ++ptr_;
+        idx_ = 0;
+      }
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator operator--(int) {
+      iterator ret(*this);
+      if (!idx_) {
+        --ptr_;
+        idx_ = kElementsPerStoredItem - 1;
+      }
+      else {
+        --idx_;
+      }
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    reference operator*() const {
+      return reference(ptr_, idx_);
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator==(iterator const &other) const {
+      return ptr_ == other.ptr_ && idx_ == other.idx_;
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator!=(iterator const &other) const {
+      return !(*this == other);
+    }
+  };
+
+  /// Bidirectional constant iterator over elements
+  class const_iterator {
+
+    /// Pointer to storage element
+    Storage const *ptr_{nullptr};
+
+    /// Index into elements packed into Storage object
+    int idx_{0};
+
+  public:
+
+    const_iterator() = default;
+
+    MUTLASS_HOST_DEVICE
+    const_iterator(Storage const *ptr, int idx = 0): ptr_(ptr), idx_(idx) { }
+
+    MUTLASS_HOST_DEVICE
+    iterator &operator++() {
+      ++idx_;
+      if (idx_ == kElementsPerStoredItem) {
+        ++ptr_;
+        idx_ = 0;
+      }
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator &operator--() {
+      if (!idx_) {
+        --ptr_;
+        idx_ = kElementsPerStoredItem - 1;
+      }
+      else {
+        --idx_;
+      }
+      return *this;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator operator++(int) {
+      iterator ret(*this);
+      ++idx_;
+      if (idx_ == kElementsPerStoredItem) {
+        ++ptr_;
+        idx_ = 0;
+      }
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    iterator operator--(int) {
+      iterator ret(*this);
+      if (!idx_) {
+        --ptr_;
+        idx_ = kElementsPerStoredItem - 1;
+      }
+      else {
+        --idx_;
+      }
+      return ret;
+    }
+
+    MUTLASS_HOST_DEVICE
+    const_reference operator*() const {
+      return const_reference(ptr_, idx_);
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator==(iterator const &other) const {
+      return ptr_ == other.ptr_ && idx_ == other.idx_;
+    }
+
+    MUTLASS_HOST_DEVICE
+    bool operator!=(iterator const &other) const {
+      return !(*this == other);
+    }
+  };
+
+  /// Bidirectional iterator over elements
+  class reverse_iterator {
+
+    /// Pointer to storage element
+    Storage *ptr_{nullptr};
+
+    /// Index into elements packed into Storage object
+    int idx_{0};
+
+  public:
+
+    reverse_iterator() = default;
+
+    MUTLASS_HOST_DEVICE
+    reverse_iterator(Storage *ptr, int idx = 0): ptr_(ptr), idx_(idx) { }
+  };
+
+  /// Bidirectional constant iterator over elements
+  class const_reverse_iterator {
+
+    /// Pointer to storage element
+    Storage const *ptr_{nullptr};
+
+    /// Index into elements packed into Storage object
+    int idx_{0};
+
+  public:
+
+    const_reverse_iterator() = default;
+
+    MUTLASS_HOST_DEVICE
+    const_reverse_iterator(Storage const *ptr, int idx = 0): ptr_(ptr), idx_(idx) { }
+  };
+
+  /// Efficient clear method
+  MUTLASS_HOST_DEVICE
+  void clear() {
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < int(kStorageElements); ++i) {
+      storage[i] = Storage(0);
+    }
+  }
+
+  MUTLASS_HOST_DEVICE
+  reference at(size_type pos) {
+    return reference(storage + pos / kElementsPerStoredItem, pos % kElementsPerStoredItem);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reference at(size_type pos) const {
+    return const_reference(storage + pos / kElementsPerStoredItem, pos % kElementsPerStoredItem);
+  }
+
+  MUTLASS_HOST_DEVICE
+  reference operator[](size_type pos) {
+    return at(pos);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reference operator[](size_type pos) const {
+    return at(pos);
+  }
+
+  MUTLASS_HOST_DEVICE
+  reference front() {
+    return at(0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reference front() const {
+    return at(0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  reference back() {
+    return reference(storage + kStorageElements - 1, kElementsPerStoredItem - 1);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reference back() const {
+    return const_reference(storage + kStorageElements - 1, kElementsPerStoredItem - 1);
+  }
+
+  MUTLASS_HOST_DEVICE
+  pointer data() {
+    return reinterpret_cast<pointer>(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_pointer data() const {
+    return reinterpret_cast<const_pointer>(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Storage * raw_data() {
+    return storage;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Storage const * raw_data() const {
+    return storage;
+  }
+
+  MUTLASS_HOST_DEVICE
+  constexpr bool empty() const {
+    return !kElements;
+  }
+
+  MUTLASS_HOST_DEVICE
+  constexpr size_type size() const {
+    return kElements;
+  }
+
+  MUTLASS_HOST_DEVICE
+  constexpr size_type max_size() const {
+    return kElements;
+  }
+
+  MUTLASS_HOST_DEVICE
+  void fill(T const &value) {
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kElementsPerStoredItem; ++i) {
+      reference ref(storage, i);
+      ref = value;
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < kStorageElements; ++i) {
+      storage[i] = storage[0];
+    }
+  }
+
+  MUTLASS_HOST_DEVICE
+  iterator begin() {
+    return iterator(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_iterator cbegin() const {
+    return const_iterator(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  iterator end() {
+    return iterator(storage + kStorageElements);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_iterator cend() const {
+    return const_iterator(storage + kStorageElements);
+  }
+
+  MUTLASS_HOST_DEVICE
+  reverse_iterator rbegin() {
+    return reverse_iterator(storage + kStorageElements);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reverse_iterator crbegin() const {
+    return const_reverse_iterator(storage + kStorageElements);
+  }
+
+  MUTLASS_HOST_DEVICE
+  reverse_iterator rend() {
+    return reverse_iterator(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  const_reverse_iterator crend() const {
+    return const_reverse_iterator(storage);
+  }
+
+private:
+  /// Internal storage
+  Storage storage[kStorageElements];
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/bfloat16.h b/vllm_musa/csrc/flash_attn/include/mutlass/bfloat16.h
new file mode 100644
index 000000000..2684c6bc4
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/bfloat16.h
@@ -0,0 +1,530 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a proxy class for storing non-standard 16-bit floating point values with
+          8 bits of exponent and 7 bit of mantissa.
+*/
+
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include "mutlass/floating_point_mtrtc.h"
+#else
+#include <cmath>
+#include <limits>
+#include <cstdint>
+#include <cstring>
+#endif
+
+#include <musa_bf16.h>
+#include "mutlass/mutlass.h"
+#include "mutlass/platform/platform.h"
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Floating-point type with 8 bits of exponent and 7 bits of mantissa.
+struct alignas(2) bfloat16_t {
+
+  //
+  // Data members
+  //
+
+  /// Storage type
+  uint16_t storage;
+
+  //
+  // Methods
+  //
+
+  /// Constructs from an unsigned short
+  MUTLASS_HOST_DEVICE
+  static bfloat16_t bitcast(uint16_t x) {
+    bfloat16_t h;
+    h.storage = x;
+    return h;
+  }
+
+private:
+  struct from_32_bit_integer_t {};
+  static constexpr from_32_bit_integer_t from_32_bit_integer{};
+
+  template<class T>
+  MUTLASS_HOST_DEVICE
+  explicit bfloat16_t(from_32_bit_integer_t, T x) {
+    static_assert(mutlass::platform::is_integral<T>::value && sizeof(T) == 4, "Requires 32-bit integer");
+
+    float flt = static_cast<float>(x);
+    uint32_t bits;
+
+    #if defined(__MUSA_ARCH__)
+    bits = reinterpret_cast<uint32_t &>(flt);
+    #else
+    std::memcpy(&bits, &flt, sizeof(bits));
+    #endif
+
+    storage = uint16_t(bits >> 16);
+  }
+
+public:
+  /// Default constructor
+  bfloat16_t() = default;
+
+  /// Reinterpret cast from MUSA's bfloat16 type
+  MUTLASS_HOST_DEVICE
+  explicit bfloat16_t(__mt_bfloat16 const & x) {
+    #if defined(__MUSA_ARCH__)
+    storage = reinterpret_cast<uint16_t const &>(x);
+    #else
+    __mt_bfloat16_raw raw(x);
+    std::memcpy(&storage, &raw.x, sizeof(storage));
+    #endif
+  }
+
+
+  /// Floating-point conversion - round toward nearest
+  MUTLASS_HOST_DEVICE
+  explicit bfloat16_t(float x) {
+
+    #if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 220)
+
+    auto value = __float2bfloat16_rn(x);
+    storage = *reinterpret_cast<uint16_t*>(&value);
+
+    #else
+    uint32_t bits;
+
+    #if defined(__MUSA_ARCH__)
+    bits = reinterpret_cast<uint32_t &>(x);
+    #else
+    std::memcpy(&bits, &x, sizeof(bits));
+    #endif
+
+    if ((bits & 0x7f800000) != 0x7f800000) {
+
+      bool mantissa_bit = ((bits & (1 << 16)) != 0);
+      bool round_bit = ((bits & (1 << 15)) != 0);
+      bool sticky_bit = ((bits & ((1 << 15) - 1)) != 0);
+
+      if ((round_bit && sticky_bit) || (round_bit && mantissa_bit)) {
+        bits += uint32_t(1 << 16);
+      }
+    }
+    else if (bits & ~0xff800000) {
+      bits = 0x7fffffff;
+    }
+
+    storage = uint16_t((bits >> 16) & 0xffff);
+    #endif
+  }
+
+  /// Floating-point conversion - round toward nearest
+  MUTLASS_HOST_DEVICE
+  explicit bfloat16_t(double x): bfloat16_t(float(x)) {
+
+  }
+
+  /// Integer conversion - round toward nearest
+  MUTLASS_HOST_DEVICE
+  explicit bfloat16_t(int x) : bfloat16_t(from_32_bit_integer, x) {}
+
+  MUTLASS_HOST_DEVICE
+  explicit bfloat16_t(uint32_t x) : bfloat16_t(from_32_bit_integer, x) {}
+
+  /// Converts to float
+  MUTLASS_HOST_DEVICE
+  operator float() const {
+    unsigned bits = (unsigned(storage) << 16);
+    #if defined(__MUSA_ARCH__)
+    return reinterpret_cast<float const &>(bits);
+    #else
+    float flt;
+    std::memcpy(&flt, &bits, sizeof(flt));
+    return flt;
+    #endif
+  }
+
+  /// Converts to float
+  MUTLASS_HOST_DEVICE
+  explicit operator double() const {
+    return double(float(*this));
+  }
+
+  /// Converts to int
+  MUTLASS_HOST_DEVICE
+  explicit operator int() const {
+    return int(float(*this));
+  }
+
+  /// Casts to bool
+  MUTLASS_HOST_DEVICE
+  explicit operator bool() const {
+    return (float(*this) != 0.0f);
+  }
+
+  /// Obtains raw bits
+  MUTLASS_HOST_DEVICE
+  uint16_t raw() const {
+    return storage;
+  }
+    /// Returns the sign bit
+  MUTLASS_HOST_DEVICE
+  bool signbit() const {
+    return ((raw() & 0x8000) != 0);
+  }
+
+  /// Returns the biased exponent
+  MUTLASS_HOST_DEVICE
+  int exponent_biased() const {
+    return int((raw() >> 7) & 0x0ff);
+  }
+
+  /// Returns the unbiased exponent
+  MUTLASS_HOST_DEVICE
+  int exponent() const {
+    return exponent_biased() - 127;
+  }
+
+  /// Returns the mantissa
+  MUTLASS_HOST_DEVICE
+  int mantissa() const {
+    return int(raw() & 0x7f);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+MUTLASS_HOST_DEVICE
+bool signbit(mutlass::bfloat16_t const& h) {
+  return h.signbit();
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::bfloat16_t abs(mutlass::bfloat16_t const& h) {
+  return mutlass::bfloat16_t::bitcast(h.raw() & 0x7fff);
+}
+
+MUTLASS_HOST_DEVICE
+bool isnan(mutlass::bfloat16_t const& h) {
+  return (h.exponent_biased() == 0x0ff) && h.mantissa();
+}
+
+MUTLASS_HOST_DEVICE
+bool isfinite(mutlass::bfloat16_t const& h) {
+  return (h.exponent_biased() != 0x0ff);
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::bfloat16_t nan_bf16(const char*) {
+  return mutlass::bfloat16_t::bitcast(0x7fff);
+}
+
+MUTLASS_HOST_DEVICE
+bool isinf(mutlass::bfloat16_t const& h) {
+  return (h.exponent_biased() == 0x0ff) && !h.mantissa();
+}
+
+MUTLASS_HOST_DEVICE
+bool isnormal(mutlass::bfloat16_t const& h) {
+  return h.exponent_biased() && h.exponent_biased() != 0x0ff;
+}
+
+MUTLASS_HOST_DEVICE
+int fpclassify(mutlass::bfloat16_t const& h) {
+  int exp = h.exponent_biased();
+  int mantissa = h.mantissa();
+  if (exp == 0x0ff) {
+    if (mantissa) {
+      return FP_NAN;
+    }
+    else {
+      return FP_INFINITE;
+    }
+  }
+  else if (!exp) {
+    if (mantissa) {
+      return FP_SUBNORMAL;
+    }
+    else {
+      return FP_ZERO;
+    }
+  }
+  return FP_NORMAL;
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::bfloat16_t sqrt(mutlass::bfloat16_t const& h) {
+#if defined(__MUSACC_RTC__)
+  return mutlass::bfloat16_t(sqrtf(float(h)));
+#else
+  return mutlass::bfloat16_t(std::sqrt(float(h)));
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t copysign(bfloat16_t const& a, bfloat16_t const& b) {
+
+  uint16_t a_bits;
+  uint16_t b_bits;
+
+  #if defined(__MUSA_ARCH__)
+  a_bits = reinterpret_cast<uint16_t const &>(a);
+  b_bits = reinterpret_cast<uint16_t const &>(b);
+  #else
+  std::memcpy(&a_bits, &a, sizeof(a_bits));
+  std::memcpy(&b_bits, &b, sizeof(b_bits));
+  #endif
+
+  uint16_t a_mag = (a_bits & 0x7fff);
+  uint16_t b_sign = (b_bits & 0x8000);
+  uint16_t result = (a_mag | b_sign);
+
+  return bfloat16_t::bitcast(result);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Standard Library operations and definitions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace std {
+
+#if !defined(__MUSACC_RTC__)
+/// Numeric limits
+template <>
+struct numeric_limits<mutlass::bfloat16_t> {
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_infinity = true;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+  static std::float_denorm_style const has_denorm = std::denorm_present;
+  static bool const has_denorm_loss = true;
+  static std::float_round_style const round_style = std::round_to_nearest;
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = 7;
+
+  /// Least positive value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t min() { return mutlass::bfloat16_t::bitcast(0x01); }
+
+  /// Minimum finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t lowest() { return mutlass::bfloat16_t::bitcast(0xff7f); }
+
+  /// Maximum finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t max() { return mutlass::bfloat16_t::bitcast(0x7f7f); }
+
+  /// Returns smallest finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t epsilon() { return mutlass::bfloat16_t::bitcast(0x1000); }
+
+  /// Returns smallest finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t round_error() { return mutlass::bfloat16_t(0.5f); }
+
+  /// Returns smallest finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t infinity() { return mutlass::bfloat16_t::bitcast(0x7f80); }
+
+  /// Returns smallest finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t quiet_NaN() { return mutlass::bfloat16_t::bitcast(0x7fff); }
+
+  /// Returns smallest finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t signaling_NaN() { return mutlass::bfloat16_t::bitcast(0x7fff); }
+
+  /// Returns smallest finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::bfloat16_t denorm_min() { return mutlass::bfloat16_t::bitcast(0x1); }
+};
+#endif
+
+} // namespace std
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Arithmetic operators
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+MUTLASS_HOST_DEVICE
+bool operator==(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return float(lhs) == float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator!=(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return float(lhs) != float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return float(lhs) < float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<=(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return float(lhs) <= float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return float(lhs) > float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>=(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return float(lhs) >= float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t operator+(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return bfloat16_t(float(lhs) + float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t operator-(bfloat16_t const& lhs) {
+  return bfloat16_t(-float(lhs));
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t operator-(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return bfloat16_t(float(lhs) - float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t operator*(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return bfloat16_t(float(lhs) * float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t operator/(bfloat16_t const& lhs, bfloat16_t const& rhs) {
+  return bfloat16_t(float(lhs) / float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t& operator+=(bfloat16_t & lhs, bfloat16_t const& rhs) {
+  lhs = bfloat16_t(float(lhs) + float(rhs));
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t& operator-=(bfloat16_t & lhs, bfloat16_t const& rhs) {
+  lhs = bfloat16_t(float(lhs) - float(rhs));
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t& operator*=(bfloat16_t & lhs, bfloat16_t const& rhs) {
+  lhs = bfloat16_t(float(lhs) * float(rhs));
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t& operator/=(bfloat16_t & lhs, bfloat16_t const& rhs) {
+  lhs = bfloat16_t(float(lhs) / float(rhs));
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t& operator++(bfloat16_t & lhs) {
+  float tmp(lhs);
+  ++tmp;
+  lhs = bfloat16_t(tmp);
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t& operator--(bfloat16_t & lhs) {
+  float tmp(lhs);
+  --tmp;
+  lhs = bfloat16_t(tmp);
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t operator++(bfloat16_t & lhs, int) {
+  bfloat16_t ret(lhs);
+  float tmp(lhs);
+  tmp++;
+  lhs = bfloat16_t(tmp);
+  return ret;
+}
+
+MUTLASS_HOST_DEVICE
+bfloat16_t operator--(bfloat16_t & lhs, int) {
+  bfloat16_t ret(lhs);
+  float tmp(lhs);
+  tmp--;
+  lhs = bfloat16_t(tmp);
+  return ret;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// User-defined literals
+//
+
+MUTLASS_HOST_DEVICE
+mutlass::bfloat16_t operator "" _bf16(long double x) {
+  return mutlass::bfloat16_t(float(x));
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::bfloat16_t operator "" _bf16(unsigned long long int x) {
+  return mutlass::bfloat16_t(int(x));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/blas3.h b/vllm_musa/csrc/flash_attn/include/mutlass/blas3.h
new file mode 100644
index 000000000..df8442020
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/blas3.h
@@ -0,0 +1,143 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Basic include for MUTLASS BLAS3/HPC code.
+
+
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+#include "mutlass/blas3_types.h"
+#include "mutlass/coord.h"
+#include "mutlass/complex.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_types.h"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines FillMode inversions
+template <FillMode kFillMode>
+struct InvertFillMode;
+
+/// Invert FillMode lower to upper
+template <>
+struct InvertFillMode<FillMode::kLower> {
+  static FillMode const mode = FillMode::kUpper;
+};
+
+/// Invert FillMode upper to lower
+template <>
+struct InvertFillMode<FillMode::kUpper> {
+  static FillMode const mode = FillMode::kLower;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// Defines SideMode inversions
+template <SideMode kSideMode>
+struct InvertSideMode;
+
+/// Invert SideMode left to right
+template <>
+struct InvertSideMode<SideMode::kLeft> {
+  static SideMode const mode = SideMode::kRight;
+};
+
+/// Invert SideMode right to left
+template <>
+struct InvertSideMode<SideMode::kRight> {
+  static SideMode const mode = SideMode::kLeft;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// Defines correct compare operation for Triangular matrix boundary
+template <FillMode kFillMode, DiagType kDiagType = DiagType::kNonUnit>
+struct TrMatrixCompareOp {
+  using Index = int32_t;
+  using Type = typename platform::conditional<
+                        (kFillMode == FillMode::kLower),
+                        greater_equal<Index>,
+                        less_equal<Index>>::type;
+};
+
+template <FillMode kFillMode>
+struct TrMatrixCompareOp <kFillMode, DiagType::kUnit> {
+   using Index = int32_t;
+   using Type = typename platform::conditional<
+                        (kFillMode == FillMode::kLower),
+                        greater_equal<Index>,
+                        less_equal<Index>>::type;
+};
+
+template <FillMode kFillMode>
+struct TrMatrixCompareOp <kFillMode, DiagType::kZero> {
+   using Index = int32_t;
+   using Type = typename platform::conditional<
+                        (kFillMode == FillMode::kLower),
+                        greater<Index>,
+                        less<Index>>::type;
+};
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// Returns precision in terms of bits (based on datatype) to fill tensors with.
+// Defaults to 5 bits of mantissa for TF32 and FP32 (with implicit round-offs).
+// Also defines acceptable mantissa result variance/error.
+template <typename Element>
+struct MantissaInBits {
+  static int constexpr bits = 5;
+  static double constexpr error = 1.0e-7;
+};
+
+// Full precision is supported for FP64
+template <>
+struct MantissaInBits<double> {
+  static int constexpr bits = 30;
+  static double constexpr error = 1.0e-14;
+};
+
+template <>
+struct MantissaInBits<mutlass::complex<double>> {
+  static int constexpr bits = 30;
+  static double constexpr error = 1.0e-15;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/blas3_types.h b/vllm_musa/csrc/flash_attn/include/mutlass/blas3_types.h
new file mode 100644
index 000000000..5f3f2cf88
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/blas3_types.h
@@ -0,0 +1,79 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Enumerated type describing the type of kernel (based on input or output matrices).
+enum class BlasMode {
+  kGemm,
+  kSymmetric,
+  kHermitian,
+  kTriangular,
+  kInvalid
+};
+
+/// Enumerated type describing the fill mode for matrices for BLAS functions.
+enum class FillMode {
+  kFull,              /// The entire tensor is covered.
+  kLower,             /// The 'lower' part of a tensor is covered including diagonal
+  kUpper,             /// The 'upper' part of a tensor is covered including diaognal
+  kDiagonal,          /// Only diagonal elements are covered.
+  kNone,              /// No element is covered.
+  kInvalid
+};
+
+/// Enumerated type describing the diagonal property of matrices for BLAS functions.
+enum class DiagType {
+  kNonUnit,
+  kUnit,
+  kZero, // Only used internally for computing SYMM/HEMM
+  kInvalid
+};
+
+/// Enumerated type describing the side dense matrix is in matrix equation for BLAS functions.
+enum class SideMode {
+  kLeft,
+  kRight,
+  kInvalid
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/complex.h
new file mode 100644
index 000000000..514962ccf
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/complex.h
@@ -0,0 +1,736 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <muComplex.h>
+
+#include <musa_fp16.h>
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <cstdint>
+#endif
+
+#include "mutlass/mutlass.h"
+#include "mutlass/functional.h"
+#include "mutlass/real.h"
+
+#include "mutlass/numeric_types.h"
+
+#include "mutlass/fast_math.h"
+
+#if !defined(__MUSACC_RTC__)
+#include <iosfwd>
+#endif
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// Enumeraed type describing a transformation on a complex value.
+enum class ComplexTransform {
+  kNone,
+  kConjugate
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// Defines ComplexTransform inversions
+template <ComplexTransform kTransform>
+struct InvertComplexTransform;
+
+/// Invert ComplexTransform from kNone to kConjugate
+template <>
+struct InvertComplexTransform<ComplexTransform::kNone> {
+  static ComplexTransform const transform = ComplexTransform::kConjugate;
+};
+
+/// Invert ComplexTransform from kConjugate to kNone
+template <>
+struct InvertComplexTransform<ComplexTransform::kConjugate> {
+  static ComplexTransform const transform = ComplexTransform::kNone;
+};
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// Accessors for MUSA complex types
+//
+
+#if !defined(__MUSACC_RTC__)
+/// Returns the real part of the complex number
+MUTLASS_HOST_DEVICE
+float const &real(muFloatComplex const &z) { return z.x; }
+
+/// Returns the real part of the complex number
+MUTLASS_HOST_DEVICE
+float &real(muFloatComplex &z) { return z.x; }
+
+/// Returns the real part of the complex number
+MUTLASS_HOST_DEVICE
+double const &real(muDoubleComplex const &z) { return z.x; }
+
+/// Returns the real part of the complex number
+MUTLASS_HOST_DEVICE
+double &real(muDoubleComplex &z) { return z.x; }
+
+/// Returns the imaginary part of the complex number
+MUTLASS_HOST_DEVICE
+float const &imag(muFloatComplex const &z) { return z.y; }
+
+/// Returns the imaginary part of the complex number
+MUTLASS_HOST_DEVICE
+float &imag(muFloatComplex &z) { return z.y; }
+
+/// Returns the imaginary part of the complex number
+MUTLASS_HOST_DEVICE
+double const &imag(muDoubleComplex const &z) { return z.y; }
+
+/// Returns the imaginary part of the complex number
+MUTLASS_HOST_DEVICE
+double &imag(muDoubleComplex &z) { return z.y; }
+#endif
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Class for representing and manipulating complex numbers with conversions from built-in MUSA
+/// complex types.
+
+template <typename T>
+class complex
+{
+ public:
+  /// Type alias for scalar type
+  using value_type = T;
+
+ private:
+  //
+  // Data members
+  //
+
+  /// Real part
+  T _real;
+
+  /// Imaginary part
+  T _imag;
+
+ public:
+
+//
+// Methods
+//
+
+  /// Default constructor
+  complex() = default;
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  complex(T r) : _real(r), _imag(T(0)) {}
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  complex(T r, T i) : _real(r), _imag(i) {}
+
+  /// Constructor
+  template<typename A>
+  MUTLASS_HOST_DEVICE
+  complex(complex<A> const &z) : _real(static_cast<T>(z.real())), _imag(static_cast<T>(z.imag())) {}
+
+
+  #if !defined(__MUSACC_RTC__)
+  /// Conversion from muFloatComplex
+  MUTLASS_HOST_DEVICE
+  complex(muFloatComplex const &z) : _real(static_cast<T>(muCrealf(z))), _imag(static_cast<T>(muCimagf(z))) {}
+
+  /// Conversion from muDoubleComplex
+  MUTLASS_HOST_DEVICE
+  complex(muDoubleComplex const &z) : _real(static_cast<T>(muCreal(z))), _imag(static_cast<T>(muCimag(z))) {}
+  #endif
+
+  /// Equality operator
+  MUTLASS_HOST_DEVICE bool operator==(complex<T> const &rhs) const {
+    return this->real() == rhs.real() && this->imag() == rhs.imag();
+  }
+
+  /// Inequality operator
+  MUTLASS_HOST_DEVICE bool operator!=(complex<T> const &rhs) const {
+    return !(*this == rhs);
+  }
+
+  /// Addition
+    template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> operator+(complex<A> const &rhs) const {
+    return complex<T>(this->real() + rhs.real(), this->imag() + rhs.imag());
+  }
+
+  /// Reduction into memory address.  Components may update out of order.
+  template <typename OtherT>
+  MUTLASS_DEVICE void red(complex<OtherT> *ptr) const {
+    static_assert(platform::is_same<T, OtherT>::value, "Component type must match");
+    mutlass::atomic_add<T> reduce;
+    reduce(&ptr->_real, _real);
+    reduce(&ptr->_imag, _imag);
+  }
+
+  /// Reduction into memory address.  Components may update out of order.  (Half specialization)
+  MUTLASS_DEVICE void red(complex<half_t> *ptr) const {
+    static_assert(platform::is_same<T, half_t>::value, "Component type must match");
+    half2 *h2_ptr = reinterpret_cast<half2*>(ptr);
+    half2 h2_data = reinterpret_cast<half2&>(*this);
+    mutlass::atomic_add<half2> reduce;
+    reduce(h2_ptr, h2_data);
+  }
+
+  /// Subtraction
+    template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> operator-(complex<A> const &rhs) const {
+    return complex<T>(this->real() - rhs.real(), this->imag() - rhs.imag());
+  }
+
+  /// Multiplication
+    template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> operator*(complex<A> const &rhs) const {
+    return complex<T>(this->real() * rhs.real() - this->imag() * rhs.imag(),
+                      this->real() * rhs.imag() + this->imag() * rhs.real());
+  }
+
+  /// Scalar Multiplication
+    template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> operator*(A const &s) const {
+    return complex<T>(this->real() * s, this->imag() * s);
+  }
+
+  /// Division
+    template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> operator/(complex<A> const &rhs) const {
+    T d = T(rhs.real() * rhs.real() + rhs.imag() * rhs.imag());
+
+    return complex<T>(
+      (real() * rhs.real() + imag() * rhs.imag()) / d,
+      (imag() * rhs.real() - real() * rhs.imag()) / d
+    );
+  }
+
+  /// Scalar Division
+    template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> operator/(A const &s) const {
+    return complex<T>(this->real() / s, this->imag() / s);
+  }
+
+  /// Addition
+    template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> &operator+=(complex<A> const &rhs) {
+      *this = *this + rhs;
+      return *this;
+  }
+
+  /// Subtraction
+  template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> &operator-=(complex<A> const &rhs) {
+      *this = *this - rhs;
+      return *this;
+  }
+
+  /// Multiplication
+  template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> &operator*=(complex<A> const &rhs) {
+      *this = *this * rhs;
+      return *this;
+  }
+
+  /// Scalar multiplication
+  template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> &operator*=(A s) {
+      *this = *this * s;
+      return *this;
+  }
+
+  /// Division
+  template <typename A>
+  MUTLASS_HOST_DEVICE complex<T> &operator/=(complex<A> const &rhs) {
+      *this = *this / rhs;
+      return *this;
+  }
+
+  /// Accesses the real part of the complex number
+  MUTLASS_HOST_DEVICE
+  T const &real() const { return _real; }
+
+  /// Accesses the real part of the complex number
+  MUTLASS_HOST_DEVICE
+  T &real() { return _real; }
+
+  /// Accesses the imaginary part of the complex number
+  MUTLASS_HOST_DEVICE
+  T const &imag() const { return _imag; }
+
+  /// Accesses the imaginary part of the complex number
+  MUTLASS_HOST_DEVICE
+  T &imag() { return _imag; }
+
+  /// Set the real part of the complex number
+  MUTLASS_HOST_DEVICE
+  void real(T real) { _real = real; }
+
+  /// Set the imaginary part of the complex number
+  MUTLASS_HOST_DEVICE
+  void imag(T imag) { _imag = imag; }
+
+  #if !defined(__MUSACC_RTC__)
+  /// Converts to muFloatComplex
+  MUTLASS_HOST_DEVICE
+  explicit operator muFloatComplex() const { return make_muFloatComplex(float(real()), float(imag())); }
+
+  /// Converts to muDoubleComplex
+  MUTLASS_HOST_DEVICE
+  explicit operator muDoubleComplex() const { return make_muDoubleComplex(real(), imag()); }
+  #endif
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// Accessors for complex template
+//
+
+/// Returns the real part of the complex number
+template <typename T>
+MUTLASS_HOST_DEVICE T const &real(complex<T> const &z) {
+  return z.real();
+}
+
+/// Returns the real part of the complex number
+template <typename T>
+MUTLASS_HOST_DEVICE T &real(complex<T> &z) {
+  return z.real();
+}
+
+/// Returns the imaginary part of the complex number
+template <typename T>
+MUTLASS_HOST_DEVICE T const &imag(complex<T> const &z) {
+  return z.imag();
+}
+
+/// Returns the imaginary part of the complex number
+template <typename T>
+MUTLASS_HOST_DEVICE T &imag(complex<T> &z) {
+  return z.imag();
+}
+
+/// Returns the real part of the real number
+template <typename T>
+MUTLASS_HOST_DEVICE T const &real(T const &r) {
+  return r;
+}
+
+/// Returns the real part of the real number
+template <typename T>
+MUTLASS_HOST_DEVICE T &real(T &r) {
+  return r;
+}
+
+/// Returns the imaginary part of the real number
+template <typename T>
+MUTLASS_HOST_DEVICE T const &imag(T const &r) {
+  return T();
+}
+
+/// Returns the imaginary part of the complex number
+template <typename T>
+MUTLASS_HOST_DEVICE T &imag(T &r) {
+  return T();
+}
+
+//
+// Output operators
+//
+
+#if !defined(__MUSACC_RTC__)
+template <typename T>
+std::ostream &operator<<(std::ostream &out, complex<T> const &z) {
+  T _r = real(z);
+  T _i = imag(z);
+
+  if (bool(_i)) {
+    return out << _r << "+i" << _i;
+  }
+  return out << _r;
+}
+#endif
+
+//
+// Non-member operators defined for complex types
+//
+
+
+//
+// Non-member functions defined for complex numbers
+//
+
+/// Returns the magnitude of the complex number
+template <typename T>
+MUTLASS_HOST_DEVICE T abs(complex<T> const &z) {
+  return sqrt(norm(z));
+}
+
+/// Returns the magnitude of the complex number
+template <typename T>
+MUTLASS_HOST_DEVICE T arg(complex<T> const &z) {
+  return atan2(imag(z), real(z));
+}
+
+/// Returns the squared magnitude of a real number
+template <typename T>
+MUTLASS_HOST_DEVICE T norm(T const &z) {
+    return z * z;
+}
+
+/// Returns the squared magnitude of a real number
+template <>
+MUTLASS_HOST_DEVICE int8_t norm(int8_t const &z) {
+    return static_cast<int8_t>(z * z);
+}
+
+/// Returns the squared magnitude of a complex number
+template <typename T>
+MUTLASS_HOST_DEVICE double norm(complex<T> const &z) {
+  return real(z) * real(z) + imag(z) * imag(z);
+}
+
+/// Norm-accumulate calculation
+template <typename T, typename R>
+MUTLASS_HOST_DEVICE R norm_accumulate(T const &x, R const & accumulator) {
+  return accumulator + static_cast<R>(x) * static_cast<R>(x);
+}
+
+/// Norm accumulate specialized for complex types
+template <typename T, typename R>
+MUTLASS_HOST_DEVICE R norm_accumulate(complex<T> const &z, R const &accumulator) {
+  return accumulator + static_cast<R>(real(z)) * static_cast<R>(real(z)) +
+    static_cast<R>(imag(z)) * static_cast<R>(imag(z));
+}
+
+MUTLASS_HOST_DEVICE float conj(float const &z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE double conj(double const &z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE half_t conj(half_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE int32_t conj(int32_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE uint32_t conj(uint32_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE int64_t conj(int64_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE uint64_t conj(uint64_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE int4b_t conj(int4b_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE uint4b_t conj(uint4b_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE bfloat16_t conj(bfloat16_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE uint1b_t conj(uint1b_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE tfloat32_t conj(tfloat32_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE float_e4m3_t conj(float_e4m3_t const& z) {
+  return z;
+}
+
+MUTLASS_HOST_DEVICE float_e5m2_t conj(float_e5m2_t const& z) {
+  return z;
+}
+
+
+/// Returns the complex conjugate
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> conj(complex<T> const &z) {
+  return complex<T>(real(z), -imag(z));
+}
+
+/// Projects the complex number z onto the Riemann sphere
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> proj(complex<T> const &z) {
+  T d = real(z) * real(z) + imag(z) * imag(z) + T(1);
+  return complex<T>((T(2) * real(z)) / d, (T(2) * imag(z)) / d);
+}
+
+/// Returns a complex number with magnitude r and phase theta
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> polar(T const &r, T const &theta = T()) {
+  return complex<T>(r * cos(theta), r * sin(theta));
+}
+
+/// Computes the complex exponential of z.
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> exp(complex<T> const &z) {
+  return complex<T>(fast_exp(real(z)) * fast_cos(imag(z)), fast_exp(real(z)) * fast_sin(imag(z)));
+}
+
+/// Computes the log of z
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> log(complex<T> const &z) {
+  return complex<T>(log(abs(z)), arg(z));
+}
+
+/// Computes the log base 10 of z
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> log10(complex<T> const &z) {
+  return log(z) / T(log(T(10)));
+}
+
+/// Computes the square root of complex number z
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> sqrt(complex<T> const &z) {
+  return sqrt(T(2)) / T(2) *
+         complex<T>(sqrt(sqrt(norm(z)) + real(z)),
+                    (imag(z) < 0 ? T(-1) : T(1)) * sqrt(sqrt(norm(z)) - real(z)));
+}
+
+/// Computes the cosine of complex z.
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> cos(complex<T> const &z) {
+  return (exp(z) + exp(-z)) / T(2);
+}
+
+/// Computes the sin of complex z.
+template <typename T>
+MUTLASS_HOST_DEVICE complex<T> sin(complex<T> const &z) {
+  return (exp(-z) - exp(z)) * complex<T>(T(0), T(1) / T(2));
+}
+
+/// Comparison
+template <typename T>
+MUTLASS_HOST_DEVICE bool operator<(complex<T> const &lhs, complex<T> const &rhs) {
+  return true;
+}
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for complex-valued type.
+template <typename T>
+struct RealType< complex<T> >
+{
+  using Type = T;
+
+  /// Number of elements
+  static int const kExtent = 2;
+
+  MUTLASS_HOST_DEVICE
+  static complex<T> from_real(double x) {
+    return complex<T>(static_cast<T>(x));
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <>
+MUTLASS_HOST_DEVICE
+mutlass::complex<half_t> from_real<mutlass::complex<half_t> >(double r) {
+  return mutlass::complex<half_t>(half_t(r));
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+mutlass::complex<float> from_real<mutlass::complex<float> >(double r) {
+  return mutlass::complex<float>(float(r));
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+mutlass::complex<double> from_real<mutlass::complex<double> >(double r) {
+  return mutlass::complex<double>(r);
+}
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct is_complex {
+  static bool const value = false;
+};
+
+template <typename T>
+struct is_complex<complex<T>> {
+  static bool const value = true;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// functional.h numeric specializations
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Squares with optional conversion
+template <typename T, typename Output>
+struct magnitude_squared<complex<T>, Output> {
+  MUTLASS_HOST_DEVICE
+  Output operator()(complex<T> lhs) const {
+    multiplies<Output> mul_op;
+
+    Output y_r = Output(lhs.real());
+    Output y_i = Output(lhs.imag());
+
+    return mul_op(y_r, y_r) + mul_op(y_i, y_i);
+  }
+};
+
+/// Fused multiply-add
+template <typename T>
+struct multiply_add<complex<T>, complex<T>, complex<T>> {
+  MUTLASS_HOST_DEVICE
+  complex<T> operator()(
+    complex<T> const &a,
+    complex<T> const &b,
+    complex<T> const &c) const {
+
+    T real = c.real();
+    T imag = c.imag();
+
+    real += a.real() * b.real();
+    real += -a.imag() * b.imag();
+    imag += a.real() * b.imag();
+    imag += a.imag () * b.real();
+
+    return complex<T>{
+      real,
+      imag
+    };
+  }
+};
+
+/// Fused multiply-add
+template <typename T>
+struct multiply_add<complex<T>, T, complex<T>> {
+  MUTLASS_HOST_DEVICE
+  complex<T> operator()(
+    complex<T> const &a,
+    T const &b,
+    complex<T> const &c) const {
+
+    T real = c.real();
+    T imag = c.imag();
+
+    real += a.real() * b;
+    imag += a.imag () * b;
+
+    return complex<T>{
+      real,
+      imag
+    };
+  }
+};
+
+/// Fused multiply-add
+template <typename T>
+struct multiply_add<T, complex<T>, complex<T>> {
+  MUTLASS_HOST_DEVICE
+  complex<T> operator()(
+    T const &a,
+    complex<T> const &b,
+    complex<T> const &c) const {
+
+    T real = c.real();
+    T imag = c.imag();
+
+    real += a * b.real();
+    imag += a * b.imag();
+
+    return complex<T>{
+      real,
+      imag
+    };
+  }
+};
+
+/// Conjugate
+template <typename T>
+struct conjugate<complex<T>>  {
+  MUTLASS_HOST_DEVICE
+  complex<T> operator()(complex<T> const &a) const {
+    return conj(a);
+  }
+};
+
+/// Computes the square of a difference with optional conversion
+template <typename T, typename Output>
+struct magnitude_squared_difference<complex<T>, Output> {
+  MUTLASS_HOST_DEVICE
+  Output operator()(complex<T> lhs, complex<T> rhs) const {
+    multiplies<Output> mul_op;
+
+    Output y_r = Output(lhs.real()) - Output(rhs.real());
+    Output y_i = Output(lhs.imag()) - Output(rhs.imag());
+
+    return mul_op(y_r, y_r) + mul_op(y_i, y_i);
+  }
+};
+
+/// Reduces value into the data pointed to by ptr (complex<T> specialization)
+template <typename T>
+struct atomic_add<complex<T>> {
+  MUTLASS_DEVICE
+  void operator()(complex<T> *ptr, const complex<T> &data)
+  {
+    data.red(ptr);
+  }
+};
+
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/constants.h b/vllm_musa/csrc/flash_attn/include/mutlass/constants.h
new file mode 100644
index 000000000..2d1757a20
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/constants.h
@@ -0,0 +1,1240 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Boost-style constant definitions for floating-point types.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/numeric_types.h"
+
+#include "mutlass/complex.h"
+
+///////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace constants {
+
+///////////////////////////////////////////////////////////////////////////////////
+
+//
+// Primary templates
+//
+
+/// Returns 1, the multiplicative identity element
+template <typename T> MUTLASS_HOST_DEVICE T one();
+
+/// Returns 0, the additive identity element
+template <typename T> MUTLASS_HOST_DEVICE T zero();
+
+/// Returns 2
+template <typename T> MUTLASS_HOST_DEVICE T two();
+
+/// Returns pi, approximately 3.141
+template <typename T> MUTLASS_HOST_DEVICE T pi();
+
+/// Returns 2 * pi
+template <typename T> MUTLASS_HOST_DEVICE T two_pi();
+
+/// Returns pi / 2
+template <typename T> MUTLASS_HOST_DEVICE T half_pi();
+
+/// Returns sqrt(pi)
+template <typename T> MUTLASS_HOST_DEVICE T root_pi();
+
+/// Returns sqrt(pi / 2)
+template <typename T> MUTLASS_HOST_DEVICE T root_half_pi();
+
+/// Returns sqrt(2 * pi)
+template <typename T> MUTLASS_HOST_DEVICE T root_two_pi();
+
+/// Returns sqrt(ln(4))
+template <typename T> MUTLASS_HOST_DEVICE T root_ln_four();
+
+/// Returns e, approximately 2.718...
+template <typename T> MUTLASS_HOST_DEVICE T e();
+
+/// Returns (1/2)
+template <typename T> MUTLASS_HOST_DEVICE T half();
+
+/// Returns sqrt(2), approximately 1.414...
+template <typename T> MUTLASS_HOST_DEVICE T root_two();
+
+/// Returns sqrt(2)/2, approximately 0.707...
+template <typename T> MUTLASS_HOST_DEVICE T half_root_two();
+
+/// Returns ln(2), approximately 0.693...
+template <typename T> MUTLASS_HOST_DEVICE T ln_two();
+
+/// Returns ln(ln(2)), approximately -0.3665...
+template <typename T> MUTLASS_HOST_DEVICE T ln_ln_two();
+
+/// Returns 1/3, approximately 0.333...
+template <typename T> MUTLASS_HOST_DEVICE T third();
+
+/// Returns 2/3, approximately 0.666...
+template <typename T> MUTLASS_HOST_DEVICE T twothirds();
+
+/// Returns pi - 3, approximately 0.1416...
+template <typename T> MUTLASS_HOST_DEVICE T pi_minus_three();
+
+/// Returns 4 - pi, approximately 0.858...
+template <typename T> MUTLASS_HOST_DEVICE T four_minus_pi();
+
+
+/////////////////////////////////////////////////////////////////////////////////////
+
+// Specialization for double
+
+/// Returns 1, the multiplicative identity element  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double one<double>() {
+  uint64_t bits = 0x3ff0000000000000ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns 1, the multiplicative identity element  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> one< complex<double> >() {
+  return complex<double>(one<double>(), double());
+}
+
+/// Returns 0, the additive identity element  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double zero<double>() {
+  uint64_t bits = 0x0ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns 0, the additive identity element  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> zero< complex<double> >() {
+  return complex<double>(zero<double>(), double());
+}
+
+/// Returns 2  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double two<double>() {
+  uint64_t bits = 0x4000000000000000ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns 2  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> two< complex<double> >() {
+  return complex<double>(two<double>(), double());
+}
+
+/// Returns pi, approximately 3.141  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double pi<double>() {
+  uint64_t bits = 0x400921fb54442d18ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns pi, approximately 3.141  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> pi< complex<double> >() {
+  return complex<double>(pi<double>(), double());
+}
+
+/// Returns 2 * pi  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double two_pi<double>() {
+  uint64_t bits = 0x401921fb54442d18ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns 2 * pi  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> two_pi< complex<double> >() {
+  return complex<double>(two_pi<double>(), double());
+}
+
+/// Returns pi / 2  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double half_pi<double>() {
+  uint64_t bits = 0x3ff921fb54442d18ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns pi / 2  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> half_pi< complex<double> >() {
+  return complex<double>(half_pi<double>(), double());
+}
+
+/// Returns sqrt(pi)  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double root_pi<double>() {
+  uint64_t bits = 0x3ffc5bf891b4ef6aull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns sqrt(pi)  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> root_pi< complex<double> >() {
+  return complex<double>(root_pi<double>(), double());
+}
+
+/// Returns sqrt(pi / 2)  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double root_half_pi<double>() {
+  uint64_t bits = 0x3ff40d931ff62705ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns sqrt(pi / 2)  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> root_half_pi< complex<double> >() {
+  return complex<double>(root_half_pi<double>(), double());
+}
+
+/// Returns sqrt(2 * pi)  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double root_two_pi<double>() {
+  uint64_t bits = 0x40040d931ff62705ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns sqrt(2 * pi)  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> root_two_pi< complex<double> >() {
+  return complex<double>(root_two_pi<double>(), double());
+}
+
+/// Returns sqrt(ln(4))  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double root_ln_four<double>() {
+  uint64_t bits = 0x3ff2d6abe44afc43ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns sqrt(ln(4))  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> root_ln_four< complex<double> >() {
+  return complex<double>(root_ln_four<double>(), double());
+}
+
+/// Returns e, approximately 2.718...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double e<double>() {
+  uint64_t bits = 0x4005bf0a8b145769ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns e, approximately 2.718...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> e< complex<double> >() {
+  return complex<double>(e<double>(), double());
+}
+
+/// Returns (1/2)  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double half<double>() {
+  uint64_t bits = 0x3fe0000000000000ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns (1/2)  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> half< complex<double> >() {
+  return complex<double>(half<double>(), double());
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double root_two<double>() {
+  uint64_t bits = 0x3ff6a09e667f3bcdull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> root_two< complex<double> >() {
+  return complex<double>(root_two<double>(), double());
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double half_root_two<double>() {
+  uint64_t bits = 0x3fe6a09e667f3bcdull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> half_root_two< complex<double> >() {
+  return complex<double>(half_root_two<double>(), double());
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double ln_two<double>() {
+  uint64_t bits = 0x3fe62e42fefa39efull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> ln_two< complex<double> >() {
+  return complex<double>(ln_two<double>(), double());
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double ln_ln_two<double>() {
+  uint64_t bits = 0xbfd774f29bdd6b9full;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> ln_ln_two< complex<double> >() {
+  return complex<double>(ln_ln_two<double>(), double());
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double third<double>() {
+  uint64_t bits = 0x3fd5555555555555ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> third< complex<double> >() {
+  return complex<double>(third<double>(), double());
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double twothirds<double>() {
+  uint64_t bits = 0x3fe5555555555555ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> twothirds< complex<double> >() {
+  return complex<double>(twothirds<double>(), double());
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double pi_minus_three<double>() {
+  uint64_t bits = 0x3fc21fb54442d180ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> pi_minus_three< complex<double> >() {
+  return complex<double>(pi_minus_three<double>(), double());
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for double)
+template <> MUTLASS_HOST_DEVICE double four_minus_pi<double>() {
+  uint64_t bits = 0x3feb7812aeef4ba0ull;
+  return reinterpret_cast<double const &>(bits);
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for complex<double>)
+template <> MUTLASS_HOST_DEVICE complex<double> four_minus_pi< complex<double> >() {
+  return complex<double>(four_minus_pi<double>(), double());
+}
+
+/////////////////////////////////////////////////////////////////////////////////////
+
+// Specialization for float
+
+/// Returns 1, the multiplicative identity element  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float one<float>() {
+  uint32_t bits = 0x3f800000u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns 1, the multiplicative identity element  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> one< complex<float> >() {
+  return complex<float>(one<float>(), float());
+}
+
+/// Returns 0, the additive identity element  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float zero<float>() {
+  uint32_t bits = 0x0u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns 0, the additive identity element  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> zero< complex<float> >() {
+  return complex<float>(zero<float>(), float());
+}
+
+/// Returns 2  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float two<float>() {
+  uint32_t bits = 0x40000000u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns 2  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> two< complex<float> >() {
+  return complex<float>(two<float>(), float());
+}
+
+/// Returns pi, approximately 3.141  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float pi<float>() {
+  uint32_t bits = 0x40490fdbu;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns pi, approximately 3.141  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> pi< complex<float> >() {
+  return complex<float>(pi<float>(), float());
+}
+
+/// Returns 2 * pi  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float two_pi<float>() {
+  uint32_t bits = 0x40c90fdbu;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns 2 * pi  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> two_pi< complex<float> >() {
+  return complex<float>(two_pi<float>(), float());
+}
+
+/// Returns pi / 2  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float half_pi<float>() {
+  uint32_t bits = 0x3fc90fdbu;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns pi / 2  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> half_pi< complex<float> >() {
+  return complex<float>(half_pi<float>(), float());
+}
+
+/// Returns sqrt(pi)  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float root_pi<float>() {
+  uint32_t bits = 0x3fe2dfc5u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns sqrt(pi)  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> root_pi< complex<float> >() {
+  return complex<float>(root_pi<float>(), float());
+}
+
+/// Returns sqrt(pi / 2)  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float root_half_pi<float>() {
+  uint32_t bits = 0x3fa06c99u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns sqrt(pi / 2)  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> root_half_pi< complex<float> >() {
+  return complex<float>(root_half_pi<float>(), float());
+}
+
+/// Returns sqrt(2 * pi)  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float root_two_pi<float>() {
+  uint32_t bits = 0x40206c99u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns sqrt(2 * pi)  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> root_two_pi< complex<float> >() {
+  return complex<float>(root_two_pi<float>(), float());
+}
+
+/// Returns sqrt(ln(4))  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float root_ln_four<float>() {
+  uint32_t bits = 0x3f96b55fu;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns sqrt(ln(4))  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> root_ln_four< complex<float> >() {
+  return complex<float>(root_ln_four<float>(), float());
+}
+
+/// Returns e, approximately 2.718...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float e<float>() {
+  uint32_t bits = 0x402df854u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns e, approximately 2.718...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> e< complex<float> >() {
+  return complex<float>(e<float>(), float());
+}
+
+/// Returns (1/2)  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float half<float>() {
+  uint32_t bits = 0x3f000000u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns (1/2)  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> half< complex<float> >() {
+  return complex<float>(half<float>(), float());
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float root_two<float>() {
+  uint32_t bits = 0x3fb504f3u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> root_two< complex<float> >() {
+  return complex<float>(root_two<float>(), float());
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float half_root_two<float>() {
+  uint32_t bits = 0x3f3504f3u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> half_root_two< complex<float> >() {
+  return complex<float>(half_root_two<float>(), float());
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float ln_two<float>() {
+  uint32_t bits = 0x3f317218u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> ln_two< complex<float> >() {
+  return complex<float>(ln_two<float>(), float());
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float ln_ln_two<float>() {
+  uint32_t bits = 0xbebba795u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> ln_ln_two< complex<float> >() {
+  return complex<float>(ln_ln_two<float>(), float());
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float third<float>() {
+  uint32_t bits = 0x3eaaaaabu;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> third< complex<float> >() {
+  return complex<float>(third<float>(), float());
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float twothirds<float>() {
+  uint32_t bits = 0x3f2aaaabu;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> twothirds< complex<float> >() {
+  return complex<float>(twothirds<float>(), float());
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float pi_minus_three<float>() {
+  uint32_t bits = 0x3e10fdaau;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> pi_minus_three< complex<float> >() {
+  return complex<float>(pi_minus_three<float>(), float());
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for float)
+template <> MUTLASS_HOST_DEVICE float four_minus_pi<float>() {
+  uint32_t bits = 0x3f5bc095u;
+  return reinterpret_cast<float const &>(bits);
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for complex<float>)
+template <> MUTLASS_HOST_DEVICE complex<float> four_minus_pi< complex<float> >() {
+  return complex<float>(four_minus_pi<float>(), float());
+}
+
+/////////////////////////////////////////////////////////////////////////////////////
+
+// Specialization for tfloat32_t
+
+/// Returns 1, the multiplicative identity element  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t one<tfloat32_t>() {
+  uint32_t bits = 0x3f801000u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns 1, the multiplicative identity element  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> one< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(one<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns 0, the additive identity element  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t zero<tfloat32_t>() {
+  uint32_t bits = 0x1000u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns 0, the additive identity element  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> zero< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(zero<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns 2  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t two<tfloat32_t>() {
+  uint32_t bits = 0x40001000u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns 2  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> two< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(two<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns pi, approximately 3.141  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t pi<tfloat32_t>() {
+  uint32_t bits = 0x40491fdbu;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns pi, approximately 3.141  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> pi< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(pi<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns 2 * pi  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t two_pi<tfloat32_t>() {
+  uint32_t bits = 0x40c91fdbu;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns 2 * pi  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> two_pi< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(two_pi<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns pi / 2  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t half_pi<tfloat32_t>() {
+  uint32_t bits = 0x3fc91fdbu;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns pi / 2  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> half_pi< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(half_pi<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns sqrt(pi)  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t root_pi<tfloat32_t>() {
+  uint32_t bits = 0x3fe2efc5u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns sqrt(pi)  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> root_pi< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(root_pi<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns sqrt(pi / 2)  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t root_half_pi<tfloat32_t>() {
+  uint32_t bits = 0x3fa07c99u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns sqrt(pi / 2)  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> root_half_pi< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(root_half_pi<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns sqrt(2 * pi)  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t root_two_pi<tfloat32_t>() {
+  uint32_t bits = 0x40207c99u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns sqrt(2 * pi)  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> root_two_pi< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(root_two_pi<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns sqrt(ln(4))  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t root_ln_four<tfloat32_t>() {
+  uint32_t bits = 0x3f96c55fu;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns sqrt(ln(4))  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> root_ln_four< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(root_ln_four<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns e, approximately 2.718...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t e<tfloat32_t>() {
+  uint32_t bits = 0x402e0854u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns e, approximately 2.718...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> e< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(e<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns (1/2)  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t half<tfloat32_t>() {
+  uint32_t bits = 0x3f001000u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns (1/2)  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> half< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(half<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t root_two<tfloat32_t>() {
+  uint32_t bits = 0x3fb514f3u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> root_two< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(root_two<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t half_root_two<tfloat32_t>() {
+  uint32_t bits = 0x3f3514f3u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> half_root_two< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(half_root_two<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t ln_two<tfloat32_t>() {
+  uint32_t bits = 0x3f318218u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> ln_two< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(ln_two<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t ln_ln_two<tfloat32_t>() {
+  uint32_t bits = 0xbebbb795u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> ln_ln_two< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(ln_ln_two<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t third<tfloat32_t>() {
+  uint32_t bits = 0x3eaabaabu;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> third< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(third<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t twothirds<tfloat32_t>() {
+  uint32_t bits = 0x3f2abaabu;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> twothirds< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(twothirds<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t pi_minus_three<tfloat32_t>() {
+  uint32_t bits = 0x3e110daau;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> pi_minus_three< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(pi_minus_three<tfloat32_t>(), tfloat32_t());
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for tfloat32_t)
+template <> MUTLASS_HOST_DEVICE tfloat32_t four_minus_pi<tfloat32_t>() {
+  uint32_t bits = 0x3f5bd095u;
+  return reinterpret_cast<tfloat32_t const &>(bits);
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for complex<tfloat32_t>)
+template <> MUTLASS_HOST_DEVICE complex<tfloat32_t> four_minus_pi< complex<tfloat32_t> >() {
+  return complex<tfloat32_t>(four_minus_pi<tfloat32_t>(), tfloat32_t());
+}
+
+/////////////////////////////////////////////////////////////////////////////////////
+
+// Specialization for half_t
+
+/// Returns 1, the multiplicative identity element  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t one<half_t>() {
+  uint16_t bits = 0x3c00u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns 1, the multiplicative identity element  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> one< complex<half_t> >() {
+  return complex<half_t>(one<half_t>(), half_t());
+}
+
+/// Returns 0, the additive identity element  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t zero<half_t>() {
+  uint16_t bits = 0x0u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns 0, the additive identity element  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> zero< complex<half_t> >() {
+  return complex<half_t>(zero<half_t>(), half_t());
+}
+
+/// Returns 2  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t two<half_t>() {
+  uint16_t bits = 0x4000u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns 2  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> two< complex<half_t> >() {
+  return complex<half_t>(two<half_t>(), half_t());
+}
+
+/// Returns pi, approximately 3.141  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t pi<half_t>() {
+  uint16_t bits = 0x4248u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns pi, approximately 3.141  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> pi< complex<half_t> >() {
+  return complex<half_t>(pi<half_t>(), half_t());
+}
+
+/// Returns 2 * pi  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t two_pi<half_t>() {
+  uint16_t bits = 0x4648u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns 2 * pi  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> two_pi< complex<half_t> >() {
+  return complex<half_t>(two_pi<half_t>(), half_t());
+}
+
+/// Returns pi / 2  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t half_pi<half_t>() {
+  uint16_t bits = 0x3e48u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns pi / 2  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> half_pi< complex<half_t> >() {
+  return complex<half_t>(half_pi<half_t>(), half_t());
+}
+
+/// Returns sqrt(pi)  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t root_pi<half_t>() {
+  uint16_t bits = 0x3f17u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns sqrt(pi)  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> root_pi< complex<half_t> >() {
+  return complex<half_t>(root_pi<half_t>(), half_t());
+}
+
+/// Returns sqrt(pi / 2)  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t root_half_pi<half_t>() {
+  uint16_t bits = 0x3d03u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns sqrt(pi / 2)  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> root_half_pi< complex<half_t> >() {
+  return complex<half_t>(root_half_pi<half_t>(), half_t());
+}
+
+/// Returns sqrt(2 * pi)  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t root_two_pi<half_t>() {
+  uint16_t bits = 0x4103u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns sqrt(2 * pi)  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> root_two_pi< complex<half_t> >() {
+  return complex<half_t>(root_two_pi<half_t>(), half_t());
+}
+
+/// Returns sqrt(ln(4))  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t root_ln_four<half_t>() {
+  uint16_t bits = 0x3cb6u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns sqrt(ln(4))  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> root_ln_four< complex<half_t> >() {
+  return complex<half_t>(root_ln_four<half_t>(), half_t());
+}
+
+/// Returns e, approximately 2.718...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t e<half_t>() {
+  uint16_t bits = 0x4170u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns e, approximately 2.718...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> e< complex<half_t> >() {
+  return complex<half_t>(e<half_t>(), half_t());
+}
+
+/// Returns (1/2)  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t half<half_t>() {
+  uint16_t bits = 0x3800u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns (1/2)  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> half< complex<half_t> >() {
+  return complex<half_t>(half<half_t>(), half_t());
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t root_two<half_t>() {
+  uint16_t bits = 0x3da8u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> root_two< complex<half_t> >() {
+  return complex<half_t>(root_two<half_t>(), half_t());
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t half_root_two<half_t>() {
+  uint16_t bits = 0x39a8u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> half_root_two< complex<half_t> >() {
+  return complex<half_t>(half_root_two<half_t>(), half_t());
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t ln_two<half_t>() {
+  uint16_t bits = 0x398cu;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> ln_two< complex<half_t> >() {
+  return complex<half_t>(ln_two<half_t>(), half_t());
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t ln_ln_two<half_t>() {
+  uint16_t bits = 0xb5ddu;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> ln_ln_two< complex<half_t> >() {
+  return complex<half_t>(ln_ln_two<half_t>(), half_t());
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t third<half_t>() {
+  uint16_t bits = 0x3555u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> third< complex<half_t> >() {
+  return complex<half_t>(third<half_t>(), half_t());
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t twothirds<half_t>() {
+  uint16_t bits = 0x3955u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> twothirds< complex<half_t> >() {
+  return complex<half_t>(twothirds<half_t>(), half_t());
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t pi_minus_three<half_t>() {
+  uint16_t bits = 0x3088u;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> pi_minus_three< complex<half_t> >() {
+  return complex<half_t>(pi_minus_three<half_t>(), half_t());
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for half_t)
+template <> MUTLASS_HOST_DEVICE half_t four_minus_pi<half_t>() {
+  uint16_t bits = 0x3adeu;
+  return reinterpret_cast<half_t const &>(bits);
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for complex<half_t>)
+template <> MUTLASS_HOST_DEVICE complex<half_t> four_minus_pi< complex<half_t> >() {
+  return complex<half_t>(four_minus_pi<half_t>(), half_t());
+}
+
+/////////////////////////////////////////////////////////////////////////////////////
+
+// Specialization for bfloat16_t
+
+/// Returns 1, the multiplicative identity element  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t one<bfloat16_t>() {
+  uint16_t bits = 0x3f80u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns 1, the multiplicative identity element  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> one< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(one<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns 0, the additive identity element  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t zero<bfloat16_t>() {
+  uint16_t bits = 0x0u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns 0, the additive identity element  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> zero< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(zero<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns 2  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t two<bfloat16_t>() {
+  uint16_t bits = 0x4000u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns 2  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> two< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(two<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns pi, approximately 3.141  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t pi<bfloat16_t>() {
+  uint16_t bits = 0x4049u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns pi, approximately 3.141  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> pi< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(pi<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns 2 * pi  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t two_pi<bfloat16_t>() {
+  uint16_t bits = 0x40c9u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns 2 * pi  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> two_pi< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(two_pi<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns pi / 2  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t half_pi<bfloat16_t>() {
+  uint16_t bits = 0x3fc9u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns pi / 2  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> half_pi< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(half_pi<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns sqrt(pi)  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t root_pi<bfloat16_t>() {
+  uint16_t bits = 0x3fe3u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns sqrt(pi)  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> root_pi< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(root_pi<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns sqrt(pi / 2)  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t root_half_pi<bfloat16_t>() {
+  uint16_t bits = 0x3fa0u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns sqrt(pi / 2)  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> root_half_pi< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(root_half_pi<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns sqrt(2 * pi)  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t root_two_pi<bfloat16_t>() {
+  uint16_t bits = 0x4020u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns sqrt(2 * pi)  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> root_two_pi< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(root_two_pi<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns sqrt(ln(4))  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t root_ln_four<bfloat16_t>() {
+  uint16_t bits = 0x3f97u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns sqrt(ln(4))  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> root_ln_four< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(root_ln_four<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns e, approximately 2.718...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t e<bfloat16_t>() {
+  uint16_t bits = 0x402eu;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns e, approximately 2.718...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> e< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(e<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns (1/2)  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t half<bfloat16_t>() {
+  uint16_t bits = 0x3f00u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns (1/2)  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> half< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(half<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t root_two<bfloat16_t>() {
+  uint16_t bits = 0x3fb5u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns sqrt(2), approximately 1.414...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> root_two< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(root_two<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t half_root_two<bfloat16_t>() {
+  uint16_t bits = 0x3f35u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns sqrt(2)/2, approximately 0.707...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> half_root_two< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(half_root_two<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t ln_two<bfloat16_t>() {
+  uint16_t bits = 0x3f31u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns ln(2), approximately 0.693...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> ln_two< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(ln_two<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t ln_ln_two<bfloat16_t>() {
+  uint16_t bits = 0xbebcu;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns ln(ln(2)), approximately -0.3665...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> ln_ln_two< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(ln_ln_two<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t third<bfloat16_t>() {
+  uint16_t bits = 0x3eabu;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns 1/3, approximately 0.333...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> third< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(third<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t twothirds<bfloat16_t>() {
+  uint16_t bits = 0x3f2bu;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns 2/3, approximately 0.666...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> twothirds< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(twothirds<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t pi_minus_three<bfloat16_t>() {
+  uint16_t bits = 0x3e11u;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns pi - 3, approximately 0.1416...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> pi_minus_three< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(pi_minus_three<bfloat16_t>(), bfloat16_t());
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for bfloat16_t)
+template <> MUTLASS_HOST_DEVICE bfloat16_t four_minus_pi<bfloat16_t>() {
+  uint16_t bits = 0x3f5cu;
+  return reinterpret_cast<bfloat16_t const &>(bits);
+}
+
+/// Returns 4 - pi, approximately 0.858...  (specialization for complex<bfloat16_t>)
+template <> MUTLASS_HOST_DEVICE complex<bfloat16_t> four_minus_pi< complex<bfloat16_t> >() {
+  return complex<bfloat16_t>(four_minus_pi<bfloat16_t>(), bfloat16_t());
+}
+///////////////////////////////////////////////////////////////////////////////////
+
+} // namespace constants
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/conv/conv2d_problem_size.h b/vllm_musa/csrc/flash_attn/include/mutlass/conv/conv2d_problem_size.h
new file mode 100644
index 000000000..3cc69dd3a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/conv/conv2d_problem_size.h
@@ -0,0 +1,654 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 This file contains definitions and utility functions for describing convolution problem sizes.
+
+  Conv2dProblem desciption:
+    activation (NHWC),
+    filter (KRSC),
+    output (NPQK),
+    pading (pad_h, pad_w),
+    stride (stride_h, stride_w),
+    dilation (dilation_h, dilation_w).
+
+  Free functions to map:
+    Map tensor extents (Conv2d -> ImplicitGemm)      : implicit_gemm_tensor_[a|b|c]_extent(ConvolutionOperator)
+    Map tensor sizes (Conv2d -> ImplicitGemm)        : implicit_gemm_tensor_[a|b|c]_size(ConvolutionOperator)
+    Map tensor problem sizes (Conv2d -> ImplicitGemm): implicit_gemm_problem_size(ConvolutionOperator)
+*/
+/*
+  Note:  MUTLASS 3x increases the host compiler requirements to C++17. However, certain
+         existing integrations of MUTLASS require C++11 host compilers.
+
+         Until this requirement can be lifted, certain headers with this annotation are required
+         to be remain consistent with C++11 syntax.
+
+         C++11 compatibility is enforced by `mutlass_test_unit_core_cpp11`.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/fast_math.h"
+#include "mutlass/gemm/gemm_enumerated_types.h"
+#include "mutlass/matrix_coord.h"
+#include "mutlass/conv/convolution.h"
+#include "mutlass/functional.h"
+
+namespace mutlass {
+namespace conv {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Problem size structure
+struct Conv2dProblemSize {
+
+  // Conv2d strictly problem size parameters
+  int N, H, W, C, P, Q, K, R, S;
+  int pad_h, pad_w;
+  int stride_h, stride_w;
+  int dilation_h, dilation_w;
+  Mode mode;
+
+  // Conv2d implementation-related parameters
+  int split_k_slices;
+  int groups;
+
+  //
+  // Methods
+  //
+
+public:
+  MUTLASS_HOST_DEVICE
+  Conv2dProblemSize():
+    N(0), H(0), W(0), C(0), P(0), Q(0), K(0), R(0), S(0),
+    pad_h(0), pad_w(0), stride_h(1), stride_w(1), dilation_h(1), dilation_w(1),
+    mode(Mode::kConvolution), split_k_slices(1), groups(1) { }
+
+  /// Constructor for default padding, stride, dilation, and split-K
+  MUTLASS_HOST_DEVICE
+  Conv2dProblemSize(
+    int N,
+    int H,
+    int W,
+    int C,
+    int P,
+    int Q,
+    int K,
+    int R,
+    int S,
+    Mode mode
+  ):
+    N(N), H(H), W(W), C(C), P(P), Q(Q), K(K), R(R), S(S),
+    pad_h(R / 2), pad_w(S / 2), stride_h(1), stride_w(1), dilation_h(1), dilation_w(1),
+    mode(mode), split_k_slices(1), groups (1) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Conv2dProblemSize(
+    int N,
+    int H,
+    int W,
+    int C,
+    int K,
+    int R,
+    int S,
+    int P,
+    int Q,
+    int pad_h,
+    int pad_w,
+    int stride_h,
+    int stride_w,
+    int dilation_h,
+    int dilation_w,
+    Mode mode,
+    int split_k_slices = 1,
+    int groups = 1
+  ):
+    N(N), H(H), W(W), C(C), P(P), Q(Q), K(K), R(R), S(S),
+    pad_h(pad_h), pad_w(pad_w), stride_h(stride_h), stride_w(stride_w),
+    dilation_h(dilation_h), dilation_w(dilation_w),
+    mode(mode), split_k_slices(split_k_slices), groups (groups) { }
+
+  /// Constructs convolution problem size from mutlass Tensor4DCoord and MatrixCoord
+  // set user-defined output size and sets P and Q (include all data members in ctor)
+  MUTLASS_HOST_DEVICE
+  Conv2dProblemSize(
+    mutlass::Tensor4DCoord input_size,    // NHWC
+    mutlass::Tensor4DCoord filter_size,   // KRSC
+    mutlass::Tensor4DCoord padding,       // pad_h, _, pad_w, _
+    mutlass::MatrixCoord stride,          // stride_h, stride_w
+    mutlass::MatrixCoord dilation,        // dilation_h, dilation_w
+    mutlass::Tensor4DCoord output_size,   // NPQK
+    mutlass::conv::Mode mode = mutlass::conv::Mode::kCrossCorrelation,
+    int split_k_slices = 1,
+    int groups = 1
+  ):
+    N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()),
+    P(output_size.h()), Q(output_size.w()),
+    K(filter_size.n()), R(filter_size.h()), S(filter_size.w()),
+    pad_h(padding[0]), pad_w(padding[2]),
+    stride_h(stride.row()), stride_w(stride.column()),
+    dilation_h(dilation.row()), dilation_w(dilation.column()),
+    mode(mode), split_k_slices(split_k_slices), groups(groups) {}
+
+  /// Constructs convolution problem size from mutlass Tensor4DCoord and MatrixCoord
+  // computes output size and sets P and Q (skip output from ctor arguments)
+  MUTLASS_HOST_DEVICE
+  Conv2dProblemSize(
+    mutlass::Tensor4DCoord input_size,   // NHWC
+    mutlass::Tensor4DCoord filter_size,  // KRSC
+    mutlass::Tensor4DCoord padding,      // pad_h, _, pad_w, _
+    mutlass::MatrixCoord stride,         // stride_h, stride_w
+    mutlass::MatrixCoord dilation,       // dilation_h, dilation_w
+    mutlass::conv::Mode mode = mutlass::conv::Mode::kCrossCorrelation,
+    int split_k_slices = 1,
+    int groups = 1
+  ):
+    N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()),
+    K(filter_size.n()), R(filter_size.h()), S(filter_size.w()),
+    pad_h(padding[0]), pad_w(padding[2]),
+    stride_h(stride.row()), stride_w(stride.column()),
+    dilation_h(dilation.row()), dilation_w(dilation.column()),
+    mode(mode), split_k_slices(split_k_slices), groups(groups) {
+      // set output P and Q
+      P = ((H + pad_h * 2 - R * dilation_h) / stride_h) + 1;
+      Q = ((W + pad_w * 2 - S * dilation_w) / stride_w) + 1;
+    }
+
+  /// Constructs convolution problem size from mutlass Tensor4DCoord and MatrixCoord
+  // set user-defined output size and sets P and Q (skip padding, striding, and dilation)
+  MUTLASS_HOST_DEVICE
+  Conv2dProblemSize(
+    mutlass::Tensor4DCoord input_size,    // NHWC
+    mutlass::Tensor4DCoord filter_size,   // KRSC
+    mutlass::Tensor4DCoord output_size,   // NPQK
+    mutlass::conv::Mode mode = mutlass::conv::Mode::kCrossCorrelation,
+    int split_k_slices = 1,
+    int groups = 1
+  ):
+    N(input_size.n()), H(input_size.h()), W(input_size.w()), C(input_size.c()),
+    P(output_size.h()), Q(output_size.w()),
+    K(filter_size.n()), R(filter_size.h()), S(filter_size.w()),
+    pad_h(R / 2), pad_w(S / 2), stride_h(1), stride_w(1),
+    dilation_h(1), dilation_w(1),
+    mode(mode), split_k_slices(split_k_slices), groups(groups) {}
+
+  // Reset covolution mode in the problem
+  MUTLASS_HOST_DEVICE
+  Conv2dProblemSize reset_mode(mutlass::conv::Mode mode_) {
+    Conv2dProblemSize tmp(*this);
+    tmp.mode = mode_;
+    return tmp;
+  }
+
+  // Reset covolution mode in the problem
+  MUTLASS_HOST_DEVICE
+  Conv2dProblemSize reset_split_k_slices(int split_k_slices_) {
+    Conv2dProblemSize tmp(*this);
+    tmp.split_k_slices = split_k_slices_;
+    return tmp;
+  }
+
+  /// Equality operator (ignores mode and split_k_slice)
+  MUTLASS_HOST_DEVICE
+  bool operator==(Conv2dProblemSize const &conv) const {
+    return (
+      (N == conv.N) && (H == conv.H) && (W == conv.W) && (C == conv.C) &&
+      (K == conv.K) && (R == conv.R) && (S == conv.S) &&
+      (P == conv.P) && (Q == conv.Q) &&
+      (pad_h == conv.pad_h) && (pad_w == conv.pad_w) &&
+      (stride_h == conv.stride_h) && (stride_w == conv.stride_w) &&
+      (dilation_h == conv.dilation_h) && (dilation_w == conv.dilation_w)
+    );
+  }
+
+  /// Inequality operator
+  MUTLASS_HOST_DEVICE
+  bool operator!=(Conv2dProblemSize const &rhs) const {
+    return !(*this == rhs);
+  }
+
+  /// Returns activation extent as Tensor4DCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::Tensor4DCoord activation_extent() const {
+
+    return mutlass::Tensor4DCoord ({N, H, W, C});
+  }
+
+  /// Returns filter extent as Tensor4DCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::Tensor4DCoord filter_extent() const {
+
+    return mutlass::Tensor4DCoord ({K, R, S, C / groups});
+  }
+
+  /// Returns output extent as Tensor4DCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::Tensor4DCoord output_extent() const {
+
+    return mutlass::Tensor4DCoord ({N, P, Q, K});
+  }
+
+  /// Returns activation size in number of elements
+  MUTLASS_HOST_DEVICE
+  int64_t activation_size() const {
+
+    return (N * H * W * C);
+  }
+
+  /// Returns filter size in number of elements
+  MUTLASS_HOST_DEVICE
+  int64_t filter_size() const {
+
+    return (K * R * S * C / groups);
+  }
+
+  /// Returns output size in number of elements
+  MUTLASS_HOST_DEVICE
+  int64_t output_size() const {
+
+    return (N * P * Q * K);
+  }
+
+  /// Returns padding as Tensor4DCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::Tensor4DCoord padding() const {
+
+    return mutlass::Tensor4DCoord ({pad_h, pad_h, pad_w, pad_w});
+  }
+
+  /// Returns stride as MatrixCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::MatrixCoord stride() const {
+
+    return mutlass::MatrixCoord ({stride_h, stride_w});
+  }
+
+  /// Returns dilation as MatrixCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::MatrixCoord dilation() const {
+
+    return mutlass::MatrixCoord ({dilation_h, dilation_w});
+  }
+
+  /////////////////////////////////////////////////////////////////
+  //        Methods used for strided dgrad implementation
+  /////////////////////////////////////////////////////////////////
+  /// Number of filter r positions to accumulate in gemm-k dim
+  MUTLASS_HOST_DEVICE
+  int num_gemm_k_filter_r(int r) const {
+    return ((R - r + stride_h - 1) / stride_h);
+  }
+
+  /// Number of filter s positions to accumulate in gemm-k dim
+  MUTLASS_HOST_DEVICE
+  int num_gemm_k_filter_s(int s) const {
+    return ((S - s + stride_w - 1) / stride_w);
+  }
+
+  /// Number of filter positions to accumulate in gemm-k dim
+  MUTLASS_HOST_DEVICE
+  int num_gemm_k_filter_positions(int r, int s) const {
+    return num_gemm_k_filter_r(r) * num_gemm_k_filter_s(s);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//                                  ImplicitGemm helper functions                                 //
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Determine the problem size of the implicit GEMM operation
+MUTLASS_HOST_DEVICE
+mutlass::gemm::GemmCoord implicit_gemm_problem_size(
+  Operator conv_operator,
+  Conv2dProblemSize const &problem_size) {
+  // Compute problem size
+  switch (conv_operator) {
+  case Operator::kFprop:
+    return gemm::GemmCoord(
+      problem_size.N * problem_size.P * problem_size.Q,
+      problem_size.K,
+      problem_size.R * problem_size.S * problem_size.C / problem_size.groups
+    );
+  case Operator::kDgrad:
+    return gemm::GemmCoord(
+      problem_size.N * problem_size.H * problem_size.W,
+      problem_size.C,
+      problem_size.R * problem_size.S * problem_size.K
+    );
+  case Operator::kWgrad:
+    return gemm::GemmCoord(
+      problem_size.K,
+      problem_size.R * problem_size.S * problem_size.C,
+      problem_size.N * problem_size.P * problem_size.Q
+    );
+  default:
+    break;
+  }
+  return gemm::GemmCoord();
+}
+
+// Determine the number of gemm_k iterations for conv2d problem using implicit gemm algorithm
+MUTLASS_HOST_DEVICE
+int implicit_gemm_k_iterations(
+  Operator conv_operator,
+  int threadblock_K,
+  Conv2dProblemSize const &problem_size,
+  IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic,
+  GroupMode group_mode = GroupMode::kNone,
+  int threadblock_N = 0) {
+
+  int iterations = 0;
+
+  if (group_mode == GroupMode::kNone) {
+
+    if (algorithm == IteratorAlgorithm::kFixedChannels) {
+
+      int positions_per_iteration = threadblock_K / problem_size.C;
+      switch (conv_operator) {
+      case Operator::kFprop:
+        iterations = (problem_size.R * problem_size.S + positions_per_iteration - 1 ) / positions_per_iteration;
+        break;
+
+      default:
+        break;
+      }
+    }
+    else if (algorithm == IteratorAlgorithm::kFewChannels) {
+
+      switch (conv_operator) {
+      case Operator::kFprop:
+        iterations = (problem_size.R * problem_size.S * problem_size.C + threadblock_K - 1 ) / threadblock_K;
+        break;
+
+      default:
+        break;
+      }
+    }
+    else {
+      int elements_per_split_k_slice = 0;
+
+      switch (conv_operator) {
+      case Operator::kFprop:
+        elements_per_split_k_slice = (problem_size.C + problem_size.split_k_slices - 1) / problem_size.split_k_slices;
+        iterations = problem_size.R * problem_size.S * ((elements_per_split_k_slice + threadblock_K - 1) / threadblock_K);
+        break;
+
+      case Operator::kDgrad:
+        elements_per_split_k_slice = (problem_size.K + problem_size.split_k_slices - 1) / problem_size.split_k_slices;
+        iterations = problem_size.R * problem_size.S * ((elements_per_split_k_slice + threadblock_K - 1) / threadblock_K);
+        break;
+
+      case Operator::kWgrad:
+        elements_per_split_k_slice = (problem_size.N * problem_size.P * problem_size.Q + problem_size.split_k_slices - 1) / problem_size.split_k_slices;
+        iterations = (elements_per_split_k_slice + threadblock_K - 1) / threadblock_K;
+        break;
+
+      default:
+        break;
+      }
+    }
+
+  } else if (group_mode == GroupMode::kDepthwise) {
+    int channels_per_cta = threadblock_N;
+
+    if (algorithm == IteratorAlgorithm::kAnalytic) {
+      switch (conv_operator) {
+        case Operator::kFprop:
+          iterations = problem_size.R * problem_size.S *
+                       ((channels_per_cta + threadblock_K - 1) / threadblock_K);
+          break;
+
+        default:
+          break;
+      }
+    }
+  } else {  // Group conv
+
+    int channels_per_group = problem_size.C / problem_size.groups;
+    int k_per_group = problem_size.K / problem_size.groups;
+
+    if (algorithm == IteratorAlgorithm::kAnalytic) {
+      switch (conv_operator) {
+        case Operator::kFprop:
+          iterations = problem_size.R * problem_size.S * ((channels_per_group + threadblock_K - 1) / threadblock_K);
+          // In group conv, if k_per_group < threadblock_N, one Threadblock will calculate multiple groups
+          if (problem_size.groups != 1) {
+            if (k_per_group < threadblock_N) {
+              iterations *= threadblock_N / k_per_group;
+            }
+          }
+          break;
+
+        default:
+          break;
+      }
+    } else if (algorithm == IteratorAlgorithm::kOptimized) {
+      // Current optimized iterator only support GroupMode::kSingleGroup
+      if (group_mode == GroupMode::kSingleGroup) {
+        switch (conv_operator) {
+          case Operator::kFprop:
+            iterations = problem_size.R * problem_size.S * ((channels_per_group + threadblock_K - 1) / threadblock_K);
+            break;
+
+          default:
+            break;
+        }
+      }
+    }
+
+  }
+
+  return iterations;
+}
+
+
+template <int N = 1, int Output_P = 1, int Output_Q = 1>
+MUTLASS_HOST_DEVICE
+int depthwise_gemm_k_iterations(
+  Operator conv_operator,
+  int threadblock_K,
+  Conv2dProblemSize const &problem_size,
+  IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic,
+  GroupMode group_mode = GroupMode::kNone,
+  int threadblock_N = 0) {
+
+    int n =  problem_size.N;
+    int p = (problem_size.P + Output_P - 1) /  Output_P;
+    int q = (problem_size.Q + Output_Q - 1) /  Output_Q;
+
+    int iterations = (n * p * q + problem_size.split_k_slices - 1) / problem_size.split_k_slices;
+    return iterations;
+}
+
+
+MUTLASS_HOST_DEVICE
+int implicit_gemm_k_iterations_per_channel(
+    Operator conv_operator,
+    Conv2dProblemSize const &problem_size,
+    IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic) {
+
+  int iterations = 0; //0 means not applicable
+  if (algorithm == IteratorAlgorithm::kAnalytic || algorithm == IteratorAlgorithm::kOptimized) {
+    switch (conv_operator) {
+      case Operator::kFprop:
+        iterations = problem_size.R * problem_size.S;
+        break;
+
+      case Operator::kDgrad:
+        iterations = problem_size.R * problem_size.S;
+        break;
+
+      default:
+        break;
+    }
+  }
+  return iterations;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+//  Mapping function (ImplicitGemm A, B, C -> Conv Activation, Filter, Output)
+////////////////////////////////////////////////////////////////////////////////
+/// Returns ImplicitGemm tensor A extent as Tensor4DCoord
+MUTLASS_HOST_DEVICE
+mutlass::Tensor4DCoord implicit_gemm_tensor_a_extent(
+  Operator conv_operator,
+  Conv2dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.activation_extent();
+    case mutlass::conv::Operator::kDgrad: return problem_size.output_extent();
+    case mutlass::conv::Operator::kWgrad: return problem_size.output_extent();
+    default : break;
+  }
+  return mutlass::Tensor4DCoord();
+}
+
+/// Returns ImplicitGemm tensor B extent as Tensor4DCoord
+MUTLASS_HOST_DEVICE
+mutlass::Tensor4DCoord implicit_gemm_tensor_b_extent(
+  Operator conv_operator,
+  Conv2dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.filter_extent();
+    case mutlass::conv::Operator::kDgrad: return problem_size.filter_extent();
+    case mutlass::conv::Operator::kWgrad: return problem_size.activation_extent();
+    default : break;
+  }
+  return mutlass::Tensor4DCoord();
+}
+
+/// Returns ImplicitGemm tensor C extent as Tensor4DCoord
+MUTLASS_HOST_DEVICE
+mutlass::Tensor4DCoord implicit_gemm_tensor_c_extent(
+  Operator conv_operator,
+  Conv2dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.output_extent();
+    case mutlass::conv::Operator::kDgrad: return problem_size.activation_extent();
+    case mutlass::conv::Operator::kWgrad: return problem_size.filter_extent();
+    default : break;
+  }
+  return mutlass::Tensor4DCoord();
+}
+
+/// Returns ImplicitGemm tensor A size in number of elements
+MUTLASS_HOST_DEVICE
+int64_t implicit_gemm_tensor_a_size(
+  Operator conv_operator,
+  Conv2dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.activation_size();
+    case mutlass::conv::Operator::kDgrad: return problem_size.output_size();
+    case mutlass::conv::Operator::kWgrad: return problem_size.output_size();
+    default : break;
+  }
+  return 0;
+}
+
+/// Returns ImplicitGemm tensor B size in number of elements
+MUTLASS_HOST_DEVICE
+int64_t implicit_gemm_tensor_b_size(
+  Operator conv_operator,
+  Conv2dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.filter_size();
+    case mutlass::conv::Operator::kDgrad: return problem_size.filter_size();
+    case mutlass::conv::Operator::kWgrad: return problem_size.activation_size();
+    default : break;
+  }
+  return 0;
+}
+
+/// Returns ImplicitGemm tensor C size in number of elements
+MUTLASS_HOST_DEVICE
+int64_t implicit_gemm_tensor_c_size(
+  Operator conv_operator,
+  Conv2dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.output_size();
+    case mutlass::conv::Operator::kDgrad: return problem_size.activation_size();
+    case mutlass::conv::Operator::kWgrad: return problem_size.filter_size();
+    default : break;
+  }
+  return 0;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//                                  Strided dgrad helper functions                                 //
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// Returns number of CTAs tile M to cover valid MMAs per starting filter postion
+MUTLASS_HOST_DEVICE
+int strided_dgrad_tile_m_per_filter(
+  Conv2dProblemSize const &problem_size,
+  int tile_size_m) {
+
+  // Compute NHW rows in Dx output that needs MMA per starting filter position
+  int rows_h_per_filter = (problem_size.H + problem_size.stride_h - 1) / problem_size.stride_h;
+  int rows_w_per_filter = (problem_size.W + problem_size.stride_w - 1) / problem_size.stride_w;
+  int rows_nhw_per_filter = problem_size.N * rows_h_per_filter * rows_w_per_filter;
+
+  // Number of CTAs tile M to cover valid MMAs per starting filter postion
+  int tile_m_per_filter = (rows_nhw_per_filter + tile_size_m - 1) / tile_size_m;
+
+  return tile_m_per_filter;
+}
+
+// Computes starting Dx coord (h, w) for given starting filter postion
+MUTLASS_HOST_DEVICE
+void strided_dgrad_starting_coords(
+  Conv2dProblemSize const &problem_size,
+  FastDivmod const &stride_h_divmod, FastDivmod const &stride_w_divmod,
+  int r, int s,
+  int &start_h, int &start_w) {
+
+  // function locals for remainder by fast divmod
+  int pad_h_rem_, pad_w_rem_;
+
+  // start_h  = std::abs(problem_size.stride_h - ((problem_size.pad_h % problem_size.stride_h) - r)) % problem_size.stride_h;
+  stride_h_divmod.divmod(pad_h_rem_, problem_size.pad_h);
+  int r_ = absolute_value(problem_size.stride_h - (pad_h_rem_ - r));
+  stride_h_divmod.divmod(start_h, r_);
+
+  //start_w  = std::abs(problem_size.stride_w - ((problem_size.pad_w % problem_size.stride_w) - s)) % problem_size.stride_w;
+  stride_w_divmod.divmod(pad_w_rem_, problem_size.pad_w);
+  int s_ = absolute_value(problem_size.stride_w - (pad_w_rem_ - s));
+  stride_w_divmod.divmod(start_w, s_);
+}
+
+} // namespace conv
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/conv/conv3d_problem_size.h b/vllm_musa/csrc/flash_attn/include/mutlass/conv/conv3d_problem_size.h
new file mode 100644
index 000000000..fdcabd3ea
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/conv/conv3d_problem_size.h
@@ -0,0 +1,486 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 This file contains definitions and utility functions for describing convolution problem sizes.
+
+  Conv3dProblem desciption:
+    activation (NDHWC),
+    filter (KTRSC),
+    output (NZPQK),
+    pading (pad_d, pad_h, pad_w),
+    stride (stride_d, stride_h, stride_w),
+    dilation (dilation_d, dilation_h, dilation_w).
+
+  Free functions to map:
+    Map tensor extents (Conv3d -> ImplicitGemm)      : implicit_gemm_tensor_[a|b|c]_extent(ConvolutionOperator)
+    Map tensor sizes (Conv3d -> ImplicitGemm)        : implicit_gemm_tensor_[a|b|c]_size(ConvolutionOperator)
+    Map tensor problem sizes (Conv3d -> ImplicitGemm): implicit_gemm_problem_size(ConvolutionOperator)
+*/
+/*
+  Note:  MUTLASS 3x increases the host compiler requirements to C++17. However, certain
+         existing integrations of MUTLASS require C++11 host compilers.
+
+         Until this requirement can be lifted, certain headers with this annotation are required
+         to be remain consistent with C++11 syntax.
+
+         C++11 compatibility is enforced by `mutlass_test_unit_core_cpp11`.
+*/
+
+#pragma once
+
+#include "mutlass/conv/convolution.h"
+#include "mutlass/conv/conv2d_problem_size.h"
+
+namespace mutlass {
+namespace conv {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Problem size structure
+struct Conv3dProblemSize : public Conv2dProblemSize {
+  //
+  // Type definitions
+  //
+
+  // 3D coordinate for padding, stride, and dilation in (d, h, w) dimensions
+  using Coord3D = Coord<3>;
+
+  //
+  // Data members
+  //
+
+  // Conv3d strictly problem size parameters
+  int D, T, Z;    // input depth, filter depth, output depth
+  int pad_d;      // padding in depth dimension
+  int stride_d;   // stride in depth dimension
+  int dilation_d; // dilation in depth dimension
+
+  //
+  // Methods
+  //
+public:
+  MUTLASS_HOST_DEVICE
+  Conv3dProblemSize():
+    Conv2dProblemSize(),
+    D(0), T(0), Z(0),
+    pad_d(0),
+    stride_d(1),
+    dilation_d(1) { }
+
+  /// Constructor for default padding, stride, dilation, and split-K
+  MUTLASS_HOST_DEVICE
+  Conv3dProblemSize(
+    int N,
+    int D,
+    int H,
+    int W,
+    int C,
+    int Z,
+    int P,
+    int Q,
+    int K,
+    int T,
+    int R,
+    int S,
+    Mode mode
+  ):
+    Conv2dProblemSize(N, H, W, C, P, Q, K, R, S, mode),
+    D(D), T(T), Z(Z),
+    pad_d(T / 2), stride_d(1), dilation_d(1) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Conv3dProblemSize(
+    int N,
+    int D,
+    int H,
+    int W,
+    int C,
+    int K,
+    int T,
+    int R,
+    int S,
+    int Z,
+    int P,
+    int Q,
+    int pad_d,
+    int pad_h,
+    int pad_w,
+    int stride_d,
+    int stride_h,
+    int stride_w,
+    int dilation_d,
+    int dilation_h,
+    int dilation_w,
+    Mode mode,
+    int split_k_slices = 1,
+    int groups = 1
+  ):
+    Conv2dProblemSize(
+    N, H, W, C, K, R, S, P, Q,
+    pad_h, pad_w,
+    stride_h, stride_w,
+    dilation_h, dilation_w,
+    mode, split_k_slices, groups),
+    D(D), T(T), Z(Z),
+    pad_d(pad_d), stride_d(stride_d), dilation_d(dilation_d) { }
+
+  /// Constructs convolution problem size from mutlass Tensor5DCoord and Coord3D
+  // set *user-defined* output size and sets Z, P, and Q (include all data members in ctor)
+  MUTLASS_HOST_DEVICE
+  Conv3dProblemSize(
+    mutlass::Tensor5DCoord input_size,    // NDHWC
+    mutlass::Tensor5DCoord filter_size,   // KTRSC
+    Coord3D padding,                      // pad_d, pad_h, pad_w
+    Coord3D stride,                       // stride_d, stride_h, stride_w
+    Coord3D dilation,                     // dilation_d, dilation_h, dilation_w
+    mutlass::Tensor5DCoord output_size,   // NZPQK
+    mutlass::conv::Mode mode = mutlass::conv::Mode::kCrossCorrelation,
+    int split_k_slices = 1,
+    int groups = 1
+  ):
+    Conv2dProblemSize(
+      {input_size.n(), input_size.h(), input_size.w(), input_size.c()},
+      {filter_size.n(), filter_size.h(), filter_size.w(), filter_size.c()},
+      {padding[1], padding[1], padding[2], padding[2]},
+      {stride[1], stride[2]},
+      {dilation[1], dilation[2]},
+      {output_size.n(), output_size.h(), output_size.w(), output_size.c()},
+      mode, split_k_slices, groups),
+    D(input_size.d()), T(filter_size.d()), Z(output_size.d()),
+    pad_d(padding[0]), stride_d(stride[0]), dilation_d(dilation[0]) { }
+
+  /// Constructs convolution problem size from mutlass Tensor5DCoord and Coord3D
+  // *computes* output size and sets Z, P and Q (include all data members in ctor)
+  MUTLASS_HOST_DEVICE
+  Conv3dProblemSize(
+    mutlass::Tensor5DCoord input_size,    // NDHWC
+    mutlass::Tensor5DCoord filter_size,   // KTRSC
+    Coord3D padding,                      // pad_d, pad_h, pad_w
+    Coord3D stride,                       // stride_d, stride_h, stride_w
+    Coord3D dilation,                     // dilation_d, dilation_h, dilation_w
+    mutlass::conv::Mode mode = mutlass::conv::Mode::kCrossCorrelation,
+    int split_k_slices = 1,
+    int groups = 1
+  ):
+    Conv2dProblemSize(
+      {input_size.n(), input_size.h(), input_size.w(), input_size.c()},
+      {filter_size.n(), filter_size.h(), filter_size.w(), filter_size.c()},
+      {padding[1], padding[1], padding[2], padding[2]},
+      {stride[1], stride[2]},
+      {dilation[1], dilation[2]},
+      mode, split_k_slices, groups),
+    D(input_size.d()), T(filter_size.d()),
+    pad_d(padding[0]), stride_d(stride[0]), dilation_d(dilation[0])
+    {
+      // set output Z
+      Z = ((D + pad_d * 2 - T * dilation_d) / stride_d) + 1;
+    }
+
+  /// Equality operator (ignores mode and split_k_slice)
+  MUTLASS_HOST_DEVICE
+  bool operator==(Conv3dProblemSize const &conv) const {
+    return (
+      (N == conv.N) && (D == conv.D) && (H == conv.H) && (W == conv.W) && (C == conv.C) &&
+      (K == conv.K) && (T == conv.T) && (R == conv.R) && (S == conv.S) &&
+      (Z == conv.Z) &&(P == conv.P) && (Q == conv.Q) &&
+      (pad_d == conv.pad_d) && (pad_h == conv.pad_h) && (pad_w == conv.pad_w) &&
+      (stride_d == conv.stride_d) && (stride_h == conv.stride_h) && (stride_w == conv.stride_w) &&
+      (dilation_d == conv.dilation_d) && (dilation_h == conv.dilation_h) && (dilation_w == conv.dilation_w)
+    );
+  }
+
+  /// Inequality operator
+  MUTLASS_HOST_DEVICE
+  bool operator!=(Conv3dProblemSize const &rhs) const {
+    return !(*this == rhs);
+  }
+
+  // Reset covolution mode in the problem
+  MUTLASS_HOST_DEVICE
+  Conv3dProblemSize reset_mode(mutlass::conv::Mode mode_) {
+    Conv3dProblemSize tmp(*this);
+    tmp.mode = mode_;
+    return tmp;
+  }
+
+  // Reset covolution mode in the problem
+  MUTLASS_HOST_DEVICE
+  Conv3dProblemSize reset_split_k_slices(int split_k_slices_) {
+    Conv3dProblemSize tmp(*this);
+    tmp.split_k_slices = split_k_slices_;
+    return tmp;
+  }
+
+  /// Returns activation extent as Tensor5DCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::Tensor5DCoord activation_extent() const {
+
+    return mutlass::Tensor5DCoord ({N, D, H, W, C});
+  }
+
+  /// Returns filter extent as Tensor5DCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::Tensor5DCoord filter_extent() const {
+
+    return mutlass::Tensor5DCoord ({K, T, R, S, C});
+  }
+
+  /// Returns output extent as Tensor5DCoord
+  MUTLASS_HOST_DEVICE
+  mutlass::Tensor5DCoord output_extent() const {
+
+    return mutlass::Tensor5DCoord ({N, Z, P, Q, K});
+  }
+
+  /// Returns activation size in number of elements
+  MUTLASS_HOST_DEVICE
+  int64_t activation_size() const {
+
+    return (N * D * H * W * C);
+  }
+
+  /// Returns filter size in number of elements
+  MUTLASS_HOST_DEVICE
+  int64_t filter_size() const {
+
+    return (K * T * R * S * C);
+  }
+
+  /// Returns output size in number of elements
+  MUTLASS_HOST_DEVICE
+  int64_t output_size() const {
+
+    return (N * Z * P * Q * K);
+  }
+
+  /// Returns output extent as Tensor5DCoord
+  MUTLASS_HOST_DEVICE
+  Coord3D padding() const {
+
+    return Coord3D ({pad_d, pad_h, pad_w});
+  }
+
+  /// Returns stride as MatrixCoord
+  MUTLASS_HOST_DEVICE
+  Coord3D stride() const {
+
+    return Coord3D ({stride_d, stride_h, stride_w});
+  }
+
+  /// Returns dilation as MatrixCoord
+  MUTLASS_HOST_DEVICE
+  Coord3D dilation() const {
+
+    return Coord3D ({dilation_d, dilation_h, dilation_w});
+  }
+
+};
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//                                  ImplicitGemm helper functions                                 //
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Determine the problem size of the implicit GEMM operation
+MUTLASS_HOST_DEVICE
+mutlass::gemm::GemmCoord implicit_gemm_problem_size(
+  Operator conv_operator,
+  Conv3dProblemSize const &problem_size) {
+  // Compute problem size
+  switch (conv_operator) {
+  case Operator::kFprop:
+    return gemm::GemmCoord(
+      problem_size.N * problem_size.Z * problem_size.P * problem_size.Q,
+      problem_size.K,
+      problem_size.T * problem_size.R * problem_size.S * problem_size.C
+    );
+  case Operator::kDgrad:
+    return gemm::GemmCoord(
+      problem_size.N * problem_size.D * problem_size.H * problem_size.W,
+      problem_size.C,
+      problem_size.T * problem_size.R * problem_size.S * problem_size.K
+    );
+  case Operator::kWgrad:
+    return gemm::GemmCoord(
+      problem_size.K,
+      problem_size.T * problem_size.R * problem_size.S * problem_size.C,
+      problem_size.N * problem_size.Z * problem_size.P * problem_size.Q
+    );
+  default:
+    break;
+  }
+  return gemm::GemmCoord();
+}
+
+// Determine the number of gemm_k iterations for conv2d problem using implicit gemm algorithm
+MUTLASS_HOST_DEVICE
+int implicit_gemm_k_iterations(
+  Operator conv_operator,
+  int threadblock_K,
+  Conv3dProblemSize const &problem_size,
+  IteratorAlgorithm algorithm = IteratorAlgorithm::kAnalytic,
+  GroupMode group_mode = GroupMode::kNone,
+  int threadblock_N = 0) {
+
+  int iterations = 0;
+  int elements_per_split_k_slice = 0;
+  if (group_mode == GroupMode::kNone) {
+    switch (conv_operator) {
+      case Operator::kFprop:
+        elements_per_split_k_slice = (problem_size.C + problem_size.split_k_slices - 1) / problem_size.split_k_slices;
+        iterations = problem_size.T * problem_size.R * problem_size.S * ((elements_per_split_k_slice + threadblock_K - 1) / threadblock_K);
+        break;
+
+      case Operator::kDgrad:
+        elements_per_split_k_slice =  (problem_size.K + problem_size.split_k_slices - 1) / problem_size.split_k_slices;
+        iterations = problem_size.T * problem_size.R * problem_size.S * ((elements_per_split_k_slice + threadblock_K - 1) / threadblock_K);
+        break;
+
+      case Operator::kWgrad:
+        elements_per_split_k_slice = (problem_size.N * problem_size.Z * problem_size.P * problem_size.Q + problem_size.split_k_slices - 1) / problem_size.split_k_slices;
+        iterations = (elements_per_split_k_slice + threadblock_K - 1) / threadblock_K;
+        break;
+
+      default:
+        break;
+    }
+  } else if (group_mode == GroupMode::kDepthwise) {
+    int channels_per_cta = threadblock_N;
+
+    if (algorithm == IteratorAlgorithm::kAnalytic) {
+      switch (conv_operator) {
+        case Operator::kFprop:
+          iterations = problem_size.T * problem_size.R * problem_size.S *
+                       ((channels_per_cta + threadblock_K - 1) / threadblock_K);
+          break;
+
+        default:
+          break;
+      }
+    }
+  }
+
+  return iterations;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+//  Mapping function (ImplicitGemm A, B, C -> Conv Activation, Filter, Output)
+////////////////////////////////////////////////////////////////////////////////
+/// Returns ImplicitGemm tensor A extent as Tensor5DCoord
+MUTLASS_HOST_DEVICE
+mutlass::Tensor5DCoord implicit_gemm_tensor_a_extent(
+  Operator conv_operator,
+  Conv3dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.activation_extent();
+    case mutlass::conv::Operator::kDgrad: return problem_size.output_extent();
+    case mutlass::conv::Operator::kWgrad: return problem_size.output_extent();
+    default : break;
+  }
+  return mutlass::Tensor5DCoord();
+}
+
+/// Returns ImplicitGemm tensor B extent as Tensor5DCoord
+MUTLASS_HOST_DEVICE
+mutlass::Tensor5DCoord implicit_gemm_tensor_b_extent(
+  Operator conv_operator,
+  Conv3dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.filter_extent();
+    case mutlass::conv::Operator::kDgrad: return problem_size.filter_extent();
+    case mutlass::conv::Operator::kWgrad: return problem_size.activation_extent();
+    default : break;
+  }
+  return mutlass::Tensor5DCoord();
+}
+
+/// Returns ImplicitGemm tensor C extent as Tensor5DCoord
+MUTLASS_HOST_DEVICE
+mutlass::Tensor5DCoord implicit_gemm_tensor_c_extent(
+  Operator conv_operator,
+  Conv3dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.output_extent();
+    case mutlass::conv::Operator::kDgrad: return problem_size.activation_extent();
+    case mutlass::conv::Operator::kWgrad: return problem_size.filter_extent();
+    default : break;
+  }
+  return mutlass::Tensor5DCoord();
+}
+
+/// Returns ImplicitGemm tensor A size in number of elements
+MUTLASS_HOST_DEVICE
+int64_t implicit_gemm_tensor_a_size(
+  Operator conv_operator,
+  Conv3dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.activation_size();
+    case mutlass::conv::Operator::kDgrad: return problem_size.output_size();
+    case mutlass::conv::Operator::kWgrad: return problem_size.output_size();
+    default : break;
+  }
+  return 0;
+}
+
+/// Returns ImplicitGemm tensor B size in number of elements
+MUTLASS_HOST_DEVICE
+int64_t implicit_gemm_tensor_b_size(
+  Operator conv_operator,
+  Conv3dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.filter_size();
+    case mutlass::conv::Operator::kDgrad: return problem_size.filter_size();
+    case mutlass::conv::Operator::kWgrad: return problem_size.activation_size();
+    default : break;
+  }
+  return 0;
+}
+
+/// Returns ImplicitGemm tensor C size in number of elements
+MUTLASS_HOST_DEVICE
+int64_t implicit_gemm_tensor_c_size(
+  Operator conv_operator,
+  Conv3dProblemSize const &problem_size) {
+  switch (conv_operator) {
+    case mutlass::conv::Operator::kFprop: return problem_size.output_size();
+    case mutlass::conv::Operator::kDgrad: return problem_size.activation_size();
+    case mutlass::conv::Operator::kWgrad: return problem_size.filter_size();
+    default : break;
+  }
+  return 0;
+}
+
+} // namespace conv
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/conv/convolution.h b/vllm_musa/csrc/flash_attn/include/mutlass/conv/convolution.h
new file mode 100644
index 000000000..f0c2fc9ae
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/conv/convolution.h
@@ -0,0 +1,180 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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
+
+This file contains definitions and utility functions for describing convolution problem sizes in terms of
+activation (NHWC), filter (KRSC), output (NPQK), padding (pad_h, pad_w), stride (stride_h, stride_w), and
+dilation (dilation_h, dilation_w).  Furthermore, it defines helper functions to map MUTLASS's implicit gemm
+tensor extents, sizes, and data types to that of the convolution's extents, sizes, and data types.
+
+                        * Mapping convolutions to Gemm computation *
+
+Mutlass implements convolutions with the Implicit Gemm algorithm.  This algorithm performs a gemm
+(general matrix-matrix multiply) on the convolution tensors Activation, Filter, and Output.
+The underlying gemm operation follows the standard gemm definition:
+
+                                     C = A * B + C
+
+                               A and B are input matrices
+                            C is source and output matrix
+
+
+For the three convolutional operators (Fprop, Dgrad, Wgrad), ImplicitGemm matrices A, B, and C are mapped
+to convolution tensors Activation, Filter and Output as described in the table below.
+
+        ___________________________________________________________________________
+         ConvolutionalOperator |        A        |      B         |       C
+        ___________________________________________________________________________
+        |                      |                 |                |               |
+        |       Fprop          |    Activation   |    Filter      |     Output    |
+        |       Dgrad          |     Output      |    Filter      |   Activation  |
+        |       Wgrad          |     Output      |  Activation    |     Filter    |
+        ___________________________________________________________________________
+
+In convolution codebase, DO NOT mix using (A, B, C) with (Activation, Filter, Output).
+
+For example, it's confusing and error prone to document a convolution class or function
+as operating on "A, B, Output."  Instead, use the mapping functions below,
+and adhere to using either A, B, C or Activation, Filter, Output.
+
+Map elements' data types (ImplicitGemm -> Conv): GemmToConvElementMap
+Map elements' data types (Conv -> ImplicitGemm): ConvToGemmElementMap
+*/
+
+/*
+  Note:  MUTLASS 3x increases the host compiler requirements to C++17. However, certain
+         existing integrations of MUTLASS require C++11 host compilers.
+
+         Until this requirement can be lifted, certain headers with this annotation are required
+         to be remain consistent with C++11 syntax.
+
+         C++11 compatibility is enforced by `mutlass_test_unit_core_cpp11`.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/fast_math.h"
+#include "mutlass/gemm/gemm_enumerated_types.h"
+#include "mutlass/matrix_coord.h"
+
+namespace mutlass {
+namespace conv {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Convolutional operator
+enum class Operator {
+  kFprop,
+  kDgrad,
+  kWgrad
+};
+
+/// Distinguishes convolution from cross correlation
+enum class Mode {
+  kCrossCorrelation,
+  kConvolution
+};
+
+/// Selects among several implementation variants trading off performance with simplicity
+enum class IteratorAlgorithm {
+  kAnalytic,      ///< functionally correct in all cases but lower performance
+  kOptimized,     ///< optimized for R <= 32, S <= 32 and unity-stride dgrad
+  kFixedChannels, ///< Analytic algorithm optimized for fixed channel count (C == AccessSize)
+  kFewChannels,   ///< Analytic algorithm optimized for few channels (C divisible by AccessSize)
+  kFixedStrideDilation ///< Optimized for fixed stride and dilation
+};
+
+/// Distinguishes among partial specializations that accelerate certain problems where convolution
+/// stride is unit.
+enum class StrideSupport {
+  kStrided,       ///< arbitrary convolution stride
+  kUnity,         ///< unit convolution stride
+  kFixed          ///< fixed convolution stride
+};
+
+/// Identifies split-K mode
+enum class SplitKMode {
+  kNone,
+  kSerial,
+  kParallel
+};
+
+/// Identifies group mode
+enum class GroupMode {
+  kNone,
+  kSingleGroup,   ///< One CTA calculates one group or less
+  kMultipleGroup, ///< One CTA calculates multiple groups
+  kDepthwise      ///< One CTA calculates cta_n groups (problem_size.C == problem_size.K == problem_size.groups)
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Shape of a tensor
+template <
+  int N = 1,
+  int H = 1,
+  int W = 1,
+  int C = 1
+>
+struct TensorNHWCShape {
+  static int const kN = N;
+  static int const kH = H;
+  static int const kW = W;
+  static int const kC = C;
+
+  static int const kHW = H * W;
+  static int const kNHW = N * kHW;
+  static int const kNHWC = N * H * W * C;
+
+  static int const kCount = kNHWC;
+
+  //
+  // Static member functions
+  //
+
+  /// Returns a Coord object
+  MUTLASS_HOST_DEVICE
+  static Coord<4> toCoord() {
+    return make_Coord(kN, kH, kW, kC);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace conv
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/coord.h b/vllm_musa/csrc/flash_attn/include/mutlass/coord.h
new file mode 100644
index 000000000..056a5121c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/coord.h
@@ -0,0 +1,480 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 A Coord is a coordinate of arbitrary rank into a tensor or matrix
+*/
+
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <stdint.h>
+#endif
+
+#include "mutlass/mutlass.h"
+
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Statically-sized array specifying Coords within a tensor
+template <
+  int Rank_,                          ///< Logical rank of coordinate
+  typename Index_ = int,              ///< Index type used for each dimension
+  typename LongIndex_ = int64_t       ///< Long index type used for linear offsets
+>
+struct Coord {
+
+public:
+
+  //
+  // Type and constant definitions
+  //
+
+  /// Number of elements in Coord
+  static int const kRank = Rank_;
+
+  /// Index type used to store elements
+  using Index = Index_;
+
+  /// Type used to represent linear offsets
+  using LongIndex = LongIndex_;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Indices
+  Index idx[kRank];
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Default ctor initializes uniformly
+  MUTLASS_HOST_DEVICE
+  explicit Coord(Index value = Index(0)) {
+    for (int i = 0; i < kRank; ++i) {
+      idx[i] = value;
+    }
+  }
+
+  /// Constructs from an array of integers
+  MUTLASS_HOST_DEVICE
+  Coord(Index const (&_idx)[kRank]) {
+    for (int i = 0; i < kRank; ++i) {
+      idx[i] = _idx[i];
+    }
+  }
+
+  /// Constructs from some other Coord
+  template <int R, typename I, typename L>
+  MUTLASS_HOST_DEVICE
+  Coord(Coord<R, I, L> other) {
+    for (int i = 0; i < kRank; ++i) {
+      idx[i] = other[i];
+    }
+  }
+
+  /// Returns a slice of the Coord which may be larger or smaller in rank
+  /// than this.
+  template <int Slice>
+  MUTLASS_HOST_DEVICE
+  Coord<Slice, Index, LongIndex> slice(int start = 0, Index identity = 0) const {
+    Coord<Slice, Index, LongIndex> result;
+    for (int i = 0; i < Slice; ++i) {
+      if (i + start < kRank) {
+        result[i] = idx[i + start];
+      }
+      else {
+        result[i] = identity;
+      }
+    }
+    return result;
+  }
+
+  /// Returns the index of the dimension with least value
+  MUTLASS_HOST_DEVICE
+  int min_dim_index() const {
+    int i = 0;
+    for (int j = 1; j < kRank; ++j) {
+      if (idx[j] < idx[i]) {
+        i = j;
+      }
+    }
+    return i;
+  }
+
+  /// Returns the index of the dimension with greatest value
+  MUTLASS_HOST_DEVICE
+  int max_dim_index() const {
+    int i = 0;
+    for (int j = 1; j < kRank; ++j) {
+      if (idx[j] > idx[i]) {
+        i = j;
+      }
+    }
+    return i;
+  }
+
+  /// Returns true if Coord is non-zero.
+  MUTLASS_HOST_DEVICE
+  explicit operator bool() const {
+    for (int i = 0; i < kRank; ++i) {
+      if (idx[i]) {
+        return true;
+      }
+    }
+    return false;
+  }
+
+  /// Returns true if Coord is uniformly zero.
+  MUTLASS_HOST_DEVICE
+  bool operator!() const {
+    for (int i = 0; i < kRank; ++i) {
+      if (idx[i]) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /// Element-wise addition
+  MUTLASS_HOST_DEVICE
+  Coord operator+(Coord const& b) const {
+    Coord c;
+    for (int i = 0; i < kRank; ++i) {
+      c.idx[i] = idx[i] + b.idx[i];
+    }
+    return c;
+  }
+
+  /// Element-wise subtraction
+  MUTLASS_HOST_DEVICE
+  Coord operator-(Coord const& b) const {
+    Coord c;
+    for (int i = 0; i < kRank; ++i) {
+      c.idx[i] = idx[i] - b.idx[i];
+    }
+    return c;
+  }
+
+  /// Element-wise multiplication
+  MUTLASS_HOST_DEVICE
+  Coord operator*(Coord const& b) const {
+    Coord c;
+    for (int i = 0; i < kRank; ++i) {
+      c.idx[i] = idx[i] * b.idx[i];
+    }
+    return c;
+  }
+
+  /// Element-wise division
+  MUTLASS_HOST_DEVICE
+  Coord operator/(Coord const& b) const {
+    Coord c;
+    for (int i = 0; i < kRank; ++i) {
+      c.idx[i] = idx[i] / b.idx[i];
+    }
+    return c;
+  }
+
+  /// In-place addition
+  MUTLASS_HOST_DEVICE
+  Coord& operator+=(Coord const& b) {
+    for (int i = 0; i < kRank; ++i) {
+      idx[i] += b.idx[i];
+    }
+    return *this;
+  }
+
+  /// In-place subtraction
+  MUTLASS_HOST_DEVICE
+  Coord& operator-=(Coord const& b) {
+    for (int i = 0; i < kRank; ++i) {
+      idx[i] -= b.idx[i];
+    }
+    return *this;
+  }
+
+  /// In-place multiplication
+  MUTLASS_HOST_DEVICE
+  Coord& operator*=(Coord const& b) {
+    for (int i = 0; i < kRank; ++i) {
+      idx[i] *= b.idx[i];
+    }
+    return *this;
+  }
+
+  /// In-place division
+  MUTLASS_HOST_DEVICE
+  Coord& operator/=(Coord const& b) {
+    for (int i = 0; i < kRank; ++i) {
+      idx[i] /= b.idx[i];
+    }
+    return *this;
+  }
+
+  /// Member access operator
+  MUTLASS_HOST_DEVICE Index& operator[](int dim) { return idx[dim]; }
+
+  /// Member access operator
+  MUTLASS_HOST_DEVICE Index const& operator[](int dim) const { return idx[dim]; }
+
+  /// Computes the dot product with anotherCoord object
+  MUTLASS_HOST_DEVICE
+  LongIndex dot(Coord const& b, LongIndex sum = LongIndex(0)) const {
+    for (int i = 0; i < kRank; ++i) {
+      sum += idx[i] * b.idx[i];
+    }
+    return sum;
+  }
+
+  /// Gets the index of a given Coord element
+  template <int Dim>
+  MUTLASS_HOST_DEVICE Index& at() {
+    return idx[Dim];
+  }
+
+  /// Access via index; may limit unrolling potential
+  MUTLASS_HOST_DEVICE
+  Index& at(int dim) { return idx[dim]; }
+
+  /// Gets the index of a given Coord element
+  template <int Dim>
+  MUTLASS_HOST_DEVICE Index const& at() const {
+    return idx[Dim];
+  }
+
+  /// Access via index; may limit unrolling potential
+  MUTLASS_HOST_DEVICE
+  Index const& at(int dim) const { return idx[dim]; }
+
+  /// Determines if two Coord<> objects are equal
+  MUTLASS_HOST_DEVICE
+  bool operator==(Coord const& b) const {
+    bool equal = true;
+    for (int i = 0; equal && i < kRank; ++i) {
+      equal = (idx[i] == b.idx[i]);
+    }
+    return equal;
+  }
+
+  /// Not equal
+  MUTLASS_HOST_DEVICE
+  bool operator!=(Coord const& b) const { return !(*this == b); }
+
+  /// Clamps a coordinate to a range specified by maximum and minimum values
+  MUTLASS_HOST_DEVICE
+  Coord& clamp(Coord const& max, Coord const& min = Coord()) {
+    for (int i = 0; i < kRank; ++i) {
+      idx[i] = __NV_STD_MAX(__NV_STD_MIN(idx[i], max.idx[i]), min.idx[i]);
+    }
+    return *this;
+  }
+
+  /// Returns the sum of all elements
+  MUTLASS_HOST_DEVICE
+  Index sum() const {
+    Index sum_(idx[0]);
+    for (int i = 1; i < kRank; ++i) {
+      sum_ += idx[i];
+    }
+    return sum_;
+  }
+
+  /// Returns the product of all elements
+  MUTLASS_HOST_DEVICE
+  LongIndex product() const {
+    LongIndex product_(idx[0]);
+    for (int i = 1; i < kRank; ++i) {
+      product_ *= idx[i];
+    }
+    return product_;
+  }
+
+  /// Less than operator
+  MUTLASS_HOST_DEVICE
+  bool operator<(Coord const &b) const {
+    for (int i = 0; i < kRank; ++i) {
+      if (!(idx[i] < b[i])) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /// Less than or equals operator
+  MUTLASS_HOST_DEVICE
+  bool operator<=(Coord const &b) const {
+    for (int i = 0; i < kRank; ++i) {
+      if (!(idx[i] <= b[i])) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /// Greater than operator
+  MUTLASS_HOST_DEVICE
+  bool operator>(Coord const &b) const {
+    return !(*this <= b);
+  }
+
+  /// Greater than or equals operator
+  MUTLASS_HOST_DEVICE
+  bool operator>=(Coord const &b) const {
+    return !(*this < b);
+  }
+};
+
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+
+/// Scalar multiplication
+template <int Rank, typename Index>
+MUTLASS_HOST_DEVICE
+Coord<Rank, Index> operator*(Index s, Coord<Rank, Index> coord) {
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank; ++i) {
+    coord[i] *= s;
+  }
+  return coord;
+}
+
+/// Scalar multiplication
+template <int Rank, typename Index>
+MUTLASS_HOST_DEVICE
+Coord<Rank, Index> operator*(Coord<Rank, Index> coord, Index s) {
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank; ++i) {
+    coord[i] *= s;
+  }
+  return coord;
+}
+
+/// Scalar division
+template <int Rank, typename Index>
+MUTLASS_HOST_DEVICE
+Coord<Rank, Index> operator/(Index s, Coord<Rank, Index> coord) {
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank; ++i) {
+    coord[i] = s / coord[i];
+  }
+  return coord;
+}
+
+/// Scalar division
+template <int Rank, typename Index>
+MUTLASS_HOST_DEVICE
+Coord<Rank, Index> operator/(Coord<Rank, Index> coord, Index s) {
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank; ++i) {
+    coord[i] /= s;
+  }
+  return coord;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Integer-valued make_Coord
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper to make a 2-element coordinate
+template <typename T>
+MUTLASS_HOST_DEVICE
+Coord<1, T> make_Coord(T _0) {
+  T values[1] = {_0};
+  return Coord<1, T>(values);
+}
+
+/// Helper to make a 2-element coordinate
+template <typename T>
+MUTLASS_HOST_DEVICE
+Coord<2, T> make_Coord(T _0, T _1) {
+  T values[2] = {_0, _1};
+  return Coord<2, T>(values);
+}
+
+/// Helper to make a 3-element coordinate
+template <typename T>
+MUTLASS_HOST_DEVICE
+Coord<3, T> make_Coord(T _0, T _1, T _2) {
+  T values[3] = {_0, _1, _2};
+  return Coord<3, T>(values);
+}
+
+/// Helper to make a 4-element coordinate
+template <typename T>
+MUTLASS_HOST_DEVICE
+Coord<4, T> make_Coord(T _0, T _1, T _2, T _3) {
+  T values[4] = {_0, _1, _2, _3};
+  return Coord<4, T>(values);
+}
+
+/// Helper to make a 5-element coordinate
+template <typename T>
+MUTLASS_HOST_DEVICE
+Coord<5, T> make_Coord(T _0, T _1, T _2, T _3, T _4) {
+  T values[5] = {_0, _1, _2, _3, _4};
+  return Coord<5, T>(values);
+}
+
+/// Helper to make a 1-element coordinate
+template <int N, typename T>
+MUTLASS_HOST_DEVICE
+Coord<N, T>make_Coord_with_padding(T _0) {
+  Coord<N, T> coord;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = N - 1; i > 0; --i) {
+    coord[i] = 0;
+  }
+
+  coord[0] = _0;
+
+  return coord;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/core_io.h b/vllm_musa/csrc/flash_attn/include/mutlass/core_io.h
new file mode 100644
index 000000000..e286222d9
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/core_io.h
@@ -0,0 +1,287 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Helpers for printing mutlass/core objects
+*/
+#pragma once
+
+#include <iostream>
+#include <typeinfo>
+
+#include "mutlass/array.h"
+#include "mutlass/coord.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/matrix.h"
+#include "mutlass/quaternion.h"
+#include "mutlass/matrix_shape.h"
+#include "mutlass/layout/pitch_linear.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm_enumerated_types.h"
+#include "mutlass/conv/convolution.h"
+#include "mutlass/conv/conv2d_problem_size.h"
+#include "mutlass/conv/conv3d_problem_size.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Output operator for MUSA built-in dim3 type
+inline std::ostream &operator<<(std::ostream &out, dim3 d) {
+  return out << d.x << ", " << d.y << ", " << d.z;
+}
+
+/// Output operator for MUSA built-in error type
+inline std::ostream &operator<<(std::ostream &out, musaError_t error) {
+  return out << musaGetErrorString(error);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//                    stream operators for mutlass namespace                                     //
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Element, int Rank>
+inline
+std::ostream& operator<<(std::ostream& out, Array<Element, Rank> const& v) {
+  for (int i = 0; i < Rank; ++i) {
+    out << (i ? ", " : "") << v[i];
+  }
+  return out;
+}
+
+template <int Rank>
+inline
+std::ostream& operator<<(std::ostream& out, Coord<Rank> const& coord) {
+  for (int i = 0; i < Rank; ++i) {
+    out << (i ? ", " : "") << coord[i];
+  }
+  return out;
+}
+
+inline
+std::istream & operator>>(std::istream &stream, half_t &x) {
+  float tmp;
+  stream >> tmp;
+  x = static_cast<mutlass::half_t>(tmp);
+  return stream;
+}
+
+inline
+std::ostream & operator<<(std::ostream &out, half_t const &x) {
+  return out << float(x);
+}
+
+inline
+std::ostream & operator<<(std::ostream &out, bfloat16_t const &x) {
+  return out << float(x);
+}
+
+inline
+std::ostream & operator<<(std::ostream &out, tfloat32_t const &x) {
+  return out << float(x);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper to enable formatted printing of MUTLASS scalar types to an ostream
+template <typename T>
+struct ScalarIO {
+
+  /// Value to print
+  T value;
+
+  /// Default ctor
+  ScalarIO() { }
+
+  /// Constructs from a value
+  ScalarIO(T value): value(value) {}
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Default printing to ostream
+template <typename T>
+inline std::ostream &operator<<(std::ostream &out, ScalarIO<T> const &scalar) {
+  return out << scalar.value;
+}
+
+/// Printing to ostream of int8_t as integer rather than character
+template <>
+inline std::ostream &operator<<(std::ostream &out, ScalarIO<int8_t> const &scalar) {
+  return out << int(scalar.value);
+}
+
+/// Printing to ostream of uint8_t as integer rather than character
+template <>
+inline std::ostream &operator<<(std::ostream &out, ScalarIO<uint8_t> const &scalar) {
+  return out << unsigned(scalar.value);
+}
+
+
+/// Default printing to ostream for MatrixShape
+template <int Row, int Column>
+inline
+std::ostream & operator<<(std::ostream &out, MatrixShape<Row, Column> const &matrix_shape) {
+  out << "mutlass::MatrixShape::(kRow, kColumn) {"
+    << mutlass::MatrixShape<Row,Column>::kRow <<","
+    << mutlass::MatrixShape<Row,Column>::kColumn <<"}";
+  return out;
+}
+
+
+/// Prints matrix to ostream
+template <typename Element, int Rows, int Columns>
+std::ostream & operator<<(std::ostream &out, Matrix<Element, Rows, Columns> const &rhs) {
+
+  for (int i = 0; i < Rows; ++i) {
+    for (int j = 0; j < Columns; ++j) {
+      ScalarIO<Element> element(rhs.at(i, j));
+      out << (j ? ", " : "") << element;
+    }
+    out << "\\n";
+  }
+
+  return out;
+}
+
+template <typename T>
+std::ostream &operator<<(std::ostream &out, Quaternion<T> const &rhs) {
+
+  out << ScalarIO<T>(rhs.w()) << " ";
+  if (rhs.x() >= 0) {
+    out << "+";
+  }
+
+  out << ScalarIO<T>(rhs.x()) << "*i ";
+  if (rhs.y() >= 0) {
+    out << "+";
+  }
+
+  out << ScalarIO<T>(rhs.y()) << "*j ";
+  if (rhs.z() >= 0) {
+    out << "+";
+  }
+
+  out << ScalarIO<T>(rhs.z()) << "*k";
+
+  return out;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//                         stream operators for mutlass::gemm namespace                          //
+///////////////////////////////////////////////////////////////////////////////////////////////////
+namespace gemm {
+
+/// Default printing to ostream for GemmShape
+template <int M, int N, int K>
+inline
+std::ostream & operator<<(std::ostream &out, GemmShape<M,N,K> const &gemm_shape) {
+  out << "mutlass::gemm::GemmShape::(kM, kN, kK) {"
+    << mutlass::gemm::GemmShape<M,N,K>::kM <<","
+    << mutlass::gemm::GemmShape<M,N,K>::kN <<","
+    << mutlass::gemm::GemmShape<M,N,K>::kK << "}";
+  return out;
+}
+
+/// Default printing to ostream for GemmCoord
+inline
+std::ostream & operator<<(std::ostream &out, GemmCoord const &gemm_coord) {
+  out << "mutlass::gemm::GemmCoord {"
+    << gemm_coord.m() <<","
+    << gemm_coord.n() <<","
+    << gemm_coord.k() << "}";
+  return out;
+}
+
+} //namespace gemm
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//                       stream operators for mutlass namespace                          //
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Default printing to ostream for PitchLinearShape
+template < int Contiguous, int Strided>
+inline
+std::ostream & operator<<(std::ostream &out, PitchLinearShape<Contiguous, Strided> const &pitch_linear_shape) {
+  out << "mutlass::PitchLinearShape:(kContiguous, kStrided) {"
+    << mutlass::layout::PitchLinearShape<Contiguous,Strided>::kContiguous <<","
+    << mutlass::layout::PitchLinearShape<Contiguous,Strided>::kStrided <<"}";
+  return out;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//                         stream operators for mutlass::conv namespace                          //
+///////////////////////////////////////////////////////////////////////////////////////////////////
+namespace conv {
+/// Default printing to ostream for Conv2dProblemSize
+inline
+std::ostream& operator<<(std::ostream& out, Conv2dProblemSize const& problem) {
+  out << "NHWC: (" << problem.N << ", " << problem.H << ", " << problem.W << ", " << problem.C << ")" << std::endl
+      << "KRSC: (" << problem.K << ", " << problem.R << ", " << problem.S << ", " << problem.C / problem.groups << ")" << std::endl
+      << "NPQK: (" << problem.N << ", " << problem.P << ", " << problem.Q << ", " << problem.K << ")" << std::endl
+      << "groups: (" << problem.groups << ")" << std::endl
+      << "Pad_h, Pad_w: (" << problem.pad_h << ", " << problem.pad_w << ")" << std::endl
+      << "Stride_h, Stride_w: (" << problem.stride_h << ", " << problem.stride_w << ")" << std::endl
+      << "Dilation_h, Dilation_w: (" << problem.dilation_h << ", " << problem.dilation_w << ")" << std::endl
+      << "split_k_slices: (" << problem.split_k_slices << ")" << std::endl
+      << "mode: (" << ((problem.mode==conv::Mode::kConvolution) ? "conv" : "xcross") << ")";
+
+  return out;
+}
+
+
+/// Default printing to ostream for Conv3dProblemSize
+inline
+std::ostream& operator<<(std::ostream& out, Conv3dProblemSize const& problem) {
+  out << "NDHWC: (" << problem.N << ", " << problem.D << ", " << problem.H << ", " << problem.W << ", " << problem.C << ")" << std::endl
+      << "KTRSC: (" << problem.K << ", " << problem.T << ", " << problem.R << ", " << problem.S << ", " << problem.C << ")" << std::endl
+      << "NZPQK: (" << problem.N << ", " << problem.Z << ", " << problem.P << ", " << problem.Q << ", " << problem.K << ")" << std::endl
+      << "pad_d, pad_h, pad_w: ("  << problem.pad_d << ", " << problem.pad_h << ", " << problem.pad_w << ")" << std::endl
+      << "stride_d, stride_h, stride_w: ("  << problem.stride_d << ", " << problem.stride_h << ", " << problem.stride_w << ")" << std::endl
+      << "dilation_d, dilation_h, dilation_w: ("  << problem.dilation_d << ", " << problem.dilation_h << ", " << problem.dilation_w << ")" << std::endl
+      << "split_k_slices: (" << problem.split_k_slices << ") " << std::endl
+      << "mode: (" << ((problem.mode==conv::Mode::kConvolution) ? "conv" : "xcross") << ")";
+
+  return out;
+}
+
+} // namespace conv
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/detail/collective.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/detail/collective.hpp
new file mode 100644
index 000000000..f26d084ba
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/detail/collective.hpp
@@ -0,0 +1,66 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mute/container/tuple.hpp"
+#include "mute/layout.hpp" // mute::size(shape)
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <size_t I, class Tuple>
+struct deduce_mixed_width_dtype {
+static_assert(I >= 0u && I <= 2u, "Valid indices are 0, 1, and 2, which represent Operand, Scale, and Bias, respectively.");
+
+private:
+  using underlying_tuple = mute::conditional_t<mute::is_tuple<Tuple>::value, Tuple, mute::tuple<Tuple>>;
+  static constexpr size_t valid_index = mute::min(I, mute::tuple_size_v<underlying_tuple> - 1);
+
+public:
+  using type = mute::conditional_t<(I < mute::tuple_size_v<underlying_tuple>),
+                                    mute::tuple_element_t<valid_index, underlying_tuple>,
+                                    void>;
+};
+
+template <size_t I, class Tuple>
+using deduce_mixed_width_dtype_t = typename deduce_mixed_width_dtype<I, Tuple>::type;
+
+} // namespace detail
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::collective
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/detail/dependent_false.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/detail/dependent_false.hpp
new file mode 100644
index 000000000..8515a47f3
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/detail/dependent_false.hpp
@@ -0,0 +1,87 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::detail {
+
+/// @brief A bool constant that depends on one or more template parameters.
+///
+/// For more detailed documentation and use cases,
+/// please see `dependent_false` below.
+template <bool Value, class... Args>
+inline constexpr bool dependent_bool_value = Value;
+
+/// @brief An always-false value that depends on one or more template parameters.
+///
+/// This exists because `static_assert(false);` always fails,
+/// even if it occurs in the `else` branch of an `if constexpr`.
+/// The following example shows how to use `dependent_false` in that case.
+///
+/// @code
+/// template<class T>
+/// void foo (T t)
+/// {
+///     if constexpr (std::is_integral_v<T>) {
+///         do_integer_stuff(t);
+///     }
+///     else if constexpr (std::is_floating_point_v<T>) {
+///         do_floating_point_stuff(t);
+///     }
+///     else {
+///         static_assert(dependent_false<T>, "T must be "
+///             "an integral or floating-point type.");
+///     }
+/// }
+/// @endcode
+///
+/// This implements the C++ Standard Library proposal P1830R1.
+///
+/// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2019/p1830r1.pdf
+///
+/// That proposal is under review as of 2022/12/05.
+/// The following link shows P1830's current review status.
+///
+/// https://github.com/cplusplus/papers/issues/572
+///
+/// P2593R0 proposes an alternate solution to this problem,
+/// that would change the C++ language itself.
+///
+/// https://www.open-std.org/jtc1/sc22/wg21/docs/papers/2022/p2593r0.html
+///
+/// For headers in this library, however, we only consider library solutions
+/// as work-arounds for future C++ features.
+template <class... Args>
+inline constexpr bool dependent_false = dependent_bool_value<false, Args...>;
+
+}  // end namespace mutlass::detail
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/detail/helper_macros.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/detail/helper_macros.hpp
new file mode 100644
index 000000000..7b25ad562
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/detail/helper_macros.hpp
@@ -0,0 +1,179 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Helper macros for the MUTLASS library
+*/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+#ifdef MUTLASS_NAMESPACE
+#define concat_tok(a, b) a ## b
+#define mkmutlassnamespace(pre, ns) concat_tok(pre, ns)
+#define mutlass mkmutlassnamespace(mutlass_, MUTLASS_NAMESPACE)
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(__MUSACC__)
+#define MUTLASS_HOST_DEVICE __forceinline__ __device__ __host__
+#define MUTLASS_DEVICE __forceinline__ __device__
+#elif defined(__MUSACC_RTC__)
+#define MUTLASS_HOST_DEVICE __forceinline__ __device__
+#define MUTLASS_DEVICE __forceinline__ __device__
+#else
+#define MUTLASS_HOST_DEVICE inline
+#define MUTLASS_DEVICE inline
+#endif
+
+#if ! defined(_MSC_VER)
+#define MUTLASS_LAMBDA_FUNC_INLINE __attribute__((always_inline))
+#else
+#define MUTLASS_LAMBDA_FUNC_INLINE [[msvc::forceinline]]
+#endif
+
+#define MUTLASS_HOST __host__
+#define MUTLASS_GLOBAL __global__ static
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+MUTLASS_HOST_DEVICE void __MUTLASS_UNUSED(T const &)
+{ }
+
+#if defined(__GNUC__)
+  #define MUTLASS_UNUSED(expr) __MUTLASS_UNUSED(expr)
+#else
+  #define MUTLASS_UNUSED(expr) do { ; } while (&expr != &expr)
+#endif
+
+#ifdef _MSC_VER
+// Provides support for alternative operators 'and', 'or', and 'not'
+#include <ciso646>
+#endif // _MSC_VER
+
+#if !defined(__MUSACC_RTC__)
+#include <cassert>
+#endif
+
+#if defined(__MUSA_ARCH__)
+  #if defined(_MSC_VER)
+    #define MUTLASS_NOT_IMPLEMENTED() { printf("%s not implemented\n", __FUNCSIG__); __musa_exit(); }
+  #else
+    #define MUTLASS_NOT_IMPLEMENTED() { printf("%s not implemented\n", __PRETTY_FUNCTION__); __musa_exit(); }
+  #endif
+#else
+  #if defined(_MSC_VER)
+    #define MUTLASS_NOT_IMPLEMENTED() assert(0 && __FUNCSIG__)
+  #else
+    #define MUTLASS_NOT_IMPLEMENTED() assert(0 && __PRETTY_FUNCTION__)
+  #endif
+#endif
+
+#if defined(__MUSA_ARCH__)
+#  define MUTLASS_CMATH_NAMESPACE
+#else
+#  define MUTLASS_CMATH_NAMESPACE std
+#endif
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+
+#ifndef MUTLASS_CONV_UNIT_TEST_RIGOROUS_SIZE_ENABLED
+#define MUTLASS_CONV_UNIT_TEST_RIGOROUS_SIZE_ENABLED 0
+#endif
+
+
+// MUSA 10.1 introduces the mma instruction
+#if !defined(MUTLASS_ENABLE_TENSOR_CORE_MMA)
+#define MUTLASS_ENABLE_TENSOR_CORE_MMA 0
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define MUTLASS_ASSERT(x) assert(x)
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// MUTLASS_PRAGMA_(UNROLL|NO_UNROLL) optimization directives for the MUSA compiler.
+#if defined(__MUSA_ARCH__) && !defined(__INTELLISENSE__)
+  #if defined(__MUSACC_RTC__) || (defined(__clang__) && defined(__MUSA__))
+    #define MUTLASS_PRAGMA_UNROLL _Pragma("unroll")
+    #define MUTLASS_PRAGMA_NO_UNROLL _Pragma("unroll 1")
+  #else
+    #define MUTLASS_PRAGMA_UNROLL #pragma unroll
+    #define MUTLASS_PRAGMA_NO_UNROLL #pragma unroll 1
+  #endif
+
+  #define MUTLASS_GEMM_LOOP MUTLASS_PRAGMA_NO_UNROLL
+
+#else
+
+    #define MUTLASS_PRAGMA_UNROLL
+    #define MUTLASS_PRAGMA_NO_UNROLL
+    #define MUTLASS_GEMM_LOOP
+
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if !defined(__MUSACC_RTC__)
+#define MUTLASS_THREAD_LOCAL thread_local
+#else
+#define MUTLASS_THREAD_LOCAL
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(_MSVC_LANG)
+#  define MUTLASS_CPLUSPLUS _MSVC_LANG
+#else
+#  define MUTLASS_CPLUSPLUS __cplusplus
+#endif
+
+#if (201700L <= MUTLASS_CPLUSPLUS)
+#define MUTLASS_CONSTEXPR_IF_CXX17 constexpr
+#define MUTLASS_CXX17_OR_LATER 1
+#else
+#define MUTLASS_CONSTEXPR_IF_CXX17
+#define MUTLASS_CXX17_OR_LATER 0
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}; // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/detail/layout.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/detail/layout.hpp
new file mode 100644
index 000000000..9d785949c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/detail/layout.hpp
@@ -0,0 +1,374 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mute/layout.hpp"
+#include "mute/pointer_sparse.hpp"       // mute::is_sparse
+#include "mute/swizzle.hpp"              // mute::Swizzle
+#include "mute/swizzle_layout.hpp"       // mute::detail::get_swizzle_portion
+#include "mute/util/type_traits.hpp"
+#include "mute/arch/copy_mp31_tme.hpp"
+#include "mutlass/layout/matrix.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_types.h"
+
+#include "mute/layout.hpp"
+#include "mute/util/type_traits.hpp"
+#include "mute/arch/copy_mp31_tme.hpp"
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::detail {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// For each mutlass::layout, provides its corresponding mute stride types, 64b by default
+
+template <class L>
+struct TagToStrideA {
+  using type = L;
+};
+
+// Maps to modes [M, K, L]
+template <>
+struct TagToStrideA<layout::RowMajor> {
+  using type = mute::Stride<int64_t, mute::Int<1>, int64_t>;
+  using tag = layout::RowMajor;
+};
+
+// Maps to modes [M, K, L]
+template <>
+struct TagToStrideA<layout::ColumnMajor> {
+  using type = mute::Stride<mute::Int<1>, int64_t, int64_t>;
+  using tag = layout::ColumnMajor;
+};
+
+template <class L>
+struct TagToStrideB {
+  using type = L;
+};
+
+// Maps to modes [N, K, L]
+template <>
+struct TagToStrideB<layout::RowMajor> {
+  using type = mute::Stride<mute::Int<1>, int64_t, int64_t>;
+  using tag = layout::RowMajor;
+};
+
+// Maps to modes [N, K, L]
+template <>
+struct TagToStrideB<layout::ColumnMajor> {
+  using type = mute::Stride<int64_t, mute::Int<1>, int64_t>;
+  using tag = layout::ColumnMajor;
+};
+
+// For each mutlass::layout *, provides its corresponding mute stride types, 64b by default
+// Used by pointer array and grouped gemm
+// Maps to modes [M, K, L]
+template <>
+struct TagToStrideA<layout::RowMajor *> {
+  using UnderlyingType = mute::Stride<int64_t, mute::Int<1>, mute::Int<0>>;
+  using type = UnderlyingType*;
+  using tag = layout::RowMajor;
+};
+
+// Maps to modes [M, K, L]
+template <>
+struct TagToStrideA<layout::ColumnMajor *> {
+  using UnderlyingType = mute::Stride<mute::Int<1>, int64_t, mute::Int<0>>;
+  using type = UnderlyingType*;
+  using tag = layout::ColumnMajor;
+};
+
+// Maps to modes [N, K, L]
+template <>
+struct TagToStrideB<layout::RowMajor *> {
+  using UnderlyingType = mute::Stride<mute::Int<1>, int64_t, mute::Int<0>>;
+  using type = UnderlyingType*;
+  using tag = layout::RowMajor;
+};
+
+// Maps to modes [N, K, L]
+template <>
+struct TagToStrideB<layout::ColumnMajor *> {
+  using UnderlyingType = mute::Stride<int64_t, mute::Int<1>, mute::Int<0>>;
+  using type = UnderlyingType*;
+  using tag = layout::ColumnMajor;
+};
+
+// Maps to modes [M, N, L]
+template <class LayoutTag>
+struct TagToStrideC : TagToStrideA<LayoutTag> { };
+
+// Conv: Maps to modes ((P,N), C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorNWC> {
+  using type = mute::Stride<mute::Stride<int64_t, int64_t>, mute::Int<1>, mute::Int<0>>;
+};
+
+// Conv: Maps to modes (PN, C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorLinearizedNWC> {
+  using type = mute::Stride<int64_t, mute::Int<1>, mute::Int<0>>;
+};
+
+// Conv: Maps to modes ((P,Q,N), C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorNHWC> {
+  using type = mute::Stride<mute::Stride<int64_t, int64_t, int64_t>, mute::Int<1>, mute::Int<0>>;
+};
+
+// Conv: Maps to modes (PQN, C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorLinearizedNHWC> {
+  using type = mute::Stride<int64_t, mute::Int<1>, mute::Int<0>>;
+};
+
+// Conv: Maps to modes ((P,Q,Z,N), C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorNDHWC> {
+  using type = mute::Stride<mute::Stride<int64_t, int64_t, int64_t, int64_t>, mute::Int<1>, mute::Int<0>>;
+};
+
+// Conv: Maps to modes (PQZN, C, _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorLinearizedNDHWC> {
+  using type = mute::Stride<int64_t, mute::Int<1>, mute::Int<0>>;
+};
+
+// Conv: Maps to modes (K, (C,S), _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorKCS> {
+  using type = mute::Stride<int64_t, mute::Stride<mute::Int<1>, int64_t>, mute::Int<0>>;
+};
+
+// Conv: Maps to modes (K, (C,S,R), _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorKCSR> {
+  using type = mute::Stride<int64_t, mute::Stride<mute::Int<1>, int64_t, int64_t>, mute::Int<0>>;
+};
+
+// Conv: Maps to modes (K, (C,S,R,T), _0) for compatiblity with GEMM epilogues expecting a batch mode stride
+template <>
+struct TagToStrideC<mutlass::layout::TensorKCSRT> {
+  using type = mute::Stride<int64_t, mute::Stride<mute::Int<1>, int64_t, int64_t, int64_t>, mute::Int<0>>;
+};
+
+// Convenience aliases
+template<class LayoutTag>
+using TagToStrideA_t = typename TagToStrideA<LayoutTag>::type;
+
+template<class LayoutTag>
+using TagToStrideB_t = typename TagToStrideB<LayoutTag>::type;
+
+template<class LayoutTag>
+using TagToStrideC_t = typename TagToStrideC<LayoutTag>::type;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// For 2.x compatibility APIs, provide stride->layout tag mappers
+
+template<int ModeIndex, class Stride>
+constexpr bool
+is_major(Stride = {}) {
+  // Account for stride types with and without batch mode and batch modes with static zero stride
+  return mute::is_constant<1, decltype(mute::front(mute::get<ModeIndex>(mute::remove_pointer_t<Stride>{})))>::value;
+}
+
+// Note : This method can be used for deducing the Layout Tag of A, C, D Matrices
+template<class StrideA>
+constexpr
+auto
+stride_to_layout_tag_A() {
+  using InternalStrideA = mute::remove_pointer_t<StrideA>;
+  if constexpr (is_major<0, StrideA>()) { // M major
+    return layout::ColumnMajor{};
+  }
+  // Specialize for sparse layout
+  else if constexpr (mute::get<0>(InternalStrideA{}) == mute::_2{} &&
+                     mute::rank(mute::get<1>(InternalStrideA{})) == 2 &&
+                     mute::is_same_v<mute::_1, mute::remove_cvref_t<decltype(mute::get<1,0>(InternalStrideA{}))>>) {
+    return layout::ColumnMajor{};
+  }
+  else { // K major
+    return layout::RowMajor{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template<class StrideB>
+constexpr
+auto
+stride_to_layout_tag_B() {
+  if constexpr (is_major<0, StrideB>()) { // N major
+    return layout::RowMajor{};
+  }
+  else { // K major
+    return layout::ColumnMajor{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+template<class StrideC>
+constexpr
+auto
+stride_to_layout_tag_C() {
+  if constexpr (is_major<0, StrideC>()) { // M major
+    return layout::ColumnMajor{};
+  }
+  else { // N major
+    return layout::RowMajor{};
+  }
+
+  MUTE_GCC_UNREACHABLE;
+}
+
+// Utilities to map Stride back on to their corresponding layout tags
+template <class S>
+struct StrideToLayoutTagA {
+  using type = decltype(detail::stride_to_layout_tag_A<S>());
+};
+
+template <class S>
+struct StrideToLayoutTagB {
+  using type = decltype(detail::stride_to_layout_tag_B<S>());
+};
+
+template <class S>
+struct StrideToLayoutTagC {
+  using type = decltype(detail::stride_to_layout_tag_C<S>());
+};
+
+// Convenience aliases
+template<class S>
+using StrideToLayoutTagA_t = typename StrideToLayoutTagA<S>::type;
+
+template<class S>
+using StrideToLayoutTagB_t = typename StrideToLayoutTagB<S>::type;
+
+template<class S>
+using StrideToLayoutTagC_t = typename StrideToLayoutTagC<S>::type;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Inspects a tiled copy and whether its copy engine is TME or not
+template<class GmemTiledCopy>
+constexpr bool is_tme_copy_engine() {
+  if constexpr (mute::is_void_v<GmemTiledCopy>) {
+    return false;
+  }
+  else {
+   if constexpr (    mute::is_same_v<mute::MP31_TME_LOAD,  GmemTiledCopy>
+                  || mute::is_same_v<mute::MP31_TME_STORE, GmemTiledCopy>
+                  ) {
+      return true;
+    }
+  }
+  return false;
+}
+
+template <class X, class = void>
+struct RawDtype { using type = X; };
+
+template <class X>
+struct RawDtype<X,mute::void_t<typename X::raw_type>> { using type = typename X::raw_type; };
+
+
+// Inspects a TiledCopy and returns its alignment in terms of element count
+template <class GmemTiledCopy, class Element, class ElementMma = Element>
+constexpr int
+get_alignment_count_from_gmem_tiled_copy() {
+
+  if constexpr (mute::is_void_v<GmemTiledCopy>) {
+    return 1;
+  }
+
+  // Account for ElementC = void kernels
+  else if constexpr (mute::is_void_v<Element>) {
+    return 0;
+  }
+
+  else {
+    if constexpr (is_tme_copy_engine<GmemTiledCopy>()) {
+      if constexpr (mute::is_sparse_v<ElementMma>) {
+        return 32 / sizeof_bits<Element>::value * ElementMma::sparsity;
+      }
+      return 32 / mute::sizeof_bits_v<Element>;
+    }
+    else {
+      // For non-TME tiled copies, TiledCopy holds the alignment count directly in its TiledShape_MN
+      return GmemTiledCopy::NumValSrc;
+    }
+
+
+  }
+}
+
+// Return the shape that is associated with stride-1 mode, or 1 if not found
+template<typename Shape, typename Stride>
+MUTLASS_HOST_DEVICE constexpr
+auto
+get_contiguous_shape(Shape const & shape, Stride const & stride) {
+  using namespace mute;
+  auto idx = find_if(append(flatten(stride), _1{}), [](auto s){ return is_constant<1,decltype(s)>{}; });
+  return get<decltype(idx)::value>(append(flatten(shape), _1{}));
+}
+
+// Check if tensor shape satisfies a given major alignment
+template<int Alignment, class Shape, class Stride>
+MUTLASS_HOST_DEVICE constexpr
+bool
+check_alignment(Shape const & shape, Stride const & stride) {
+  return is_major<0>(stride)
+    ? get_contiguous_shape(mute::get<0>(shape), mute::get<0>(stride)) % Alignment == 0
+    : get_contiguous_shape(mute::get<1>(shape), mute::get<1>(stride)) % Alignment == 0;
+}
+
+// Check if tensor shape satisfies a given major alignment
+
+template<int B, int M, int S>
+MUTLASS_HOST_DEVICE constexpr
+size_t
+alignment_for_swizzle(mute::Swizzle<B, M, S>) {
+  static_assert(B >= 0 and M >= 0);
+  return size_t(1) << size_t(B + M + mute::abs(S));
+}
+
+template<class Layout>
+MUTLASS_HOST_DEVICE constexpr
+size_t
+alignment_for_swizzle(Layout layout) {
+  return alignment_for_swizzle(mute::detail::get_swizzle_portion(layout));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::detail
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/detail/mainloop_fusion_helper_scale_factor.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/detail/mainloop_fusion_helper_scale_factor.hpp
new file mode 100644
index 000000000..5b44208f7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/detail/mainloop_fusion_helper_scale_factor.hpp
@@ -0,0 +1,76 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Helper macros for the MUTLASS library
+*/
+
+#pragma once
+
+#include <mute/util/type_traits.hpp> // mute::void_t
+
+namespace mutlass::detail {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+template <typename CollectiveMainloop, typename = void>
+struct ElementSFType {
+  using type = void;
+};
+
+template <typename CollectiveMainloop>
+struct ElementSFType<CollectiveMainloop, mute::void_t<typename CollectiveMainloop::ElementSF>> {
+  using type = typename CollectiveMainloop::ElementSF;
+};
+
+template <typename CollectiveMainloop, typename = void>
+struct LayoutSFAType {
+  using type = void;
+};
+
+template <typename CollectiveMainloop>
+struct LayoutSFAType<CollectiveMainloop, mute::void_t<typename CollectiveMainloop::LayoutSFA>> {
+  using type = typename CollectiveMainloop::LayoutSFA;
+};
+
+template <typename CollectiveMainloop, typename = void>
+struct LayoutSFBType {
+  using type = void;
+};
+
+template <typename CollectiveMainloop>
+struct LayoutSFBType<CollectiveMainloop, mute::void_t<typename CollectiveMainloop::LayoutSFB>> {
+  using type = typename CollectiveMainloop::LayoutSFB;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::detail
\ No newline at end of file
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/detail/mma.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/detail/mma.hpp
new file mode 100644
index 000000000..412407136
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/detail/mma.hpp
@@ -0,0 +1,72 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/arch/mma.h"
+#include "mute/layout.hpp"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::detail {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class TiledMma, class = void>
+struct IsSparseTensorOp : mute::false_type { };
+
+// TiledMma for sparse must have ValTypeE
+template <class TiledMma>
+struct IsSparseTensorOp<TiledMma, mute::void_t<typename TiledMma::ValTypeE>>
+    : mute::true_type { };
+
+// The following metafunction is used to extract the OperatorClass from a mutlass 3.x kernel.
+template <class TiledMma>
+struct get_operator_class {
+  static constexpr bool is_sparse_op = IsSparseTensorOp<TiledMma>::value;
+  static constexpr bool is_tensor_op = mute::size<0>(typename TiledMma::AtomShape_MNK{}) >= 8;
+  using type = mute::conditional_t<
+                is_tensor_op,
+                mute::conditional_t<
+                  is_sparse_op,
+                  mutlass::arch::OpClassSparseTensorOp,
+                    mutlass::arch::OpClassTensorOp
+                  >,
+                mutlass::arch::OpClassSimt
+                >;
+};
+
+template <class T>
+using get_operator_class_t = typename get_operator_class<T>::type;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::detail
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/device_kernel.h b/vllm_musa/csrc/flash_attn/include/mutlass/device_kernel.h
new file mode 100644
index 000000000..267b3146b
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/device_kernel.h
@@ -0,0 +1,114 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Template for generic MUTLASS kernel.
+*/
+
+#pragma once
+
+#define MUTLASS_GRID_CONSTANT
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+template <typename T>
+struct Type2Type
+{
+  using type=T;
+};
+// using the simple type to replace the complex type to reduce this symbol size
+template <typename  T>
+struct GetUnderlyingKernel: public Type2Type<T>
+{};
+template <uint64_t shader_guid, unsigned index, template <uint64_t, unsigned> class Wrapper >
+struct GetUnderlyingKernel<Wrapper<shader_guid,index>>: public Wrapper<shader_guid,index>
+{};
+template <typename  T>
+using GetUnderlyingKernel_t = typename GetUnderlyingKernel<T>::type;
+
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Generic MUTLASS kernel template.
+template <typename Operator>
+MUTLASS_GLOBAL
+void Kernel(typename Operator::Params params) {
+  // Dynamic shared memory base pointer
+  extern __shared__ int SharedStorageBase[];
+  // Declare pointer to dynamic shared memory.
+  typename Operator::SharedStorage *shared_storage =
+      reinterpret_cast<typename Operator::SharedStorage *>(SharedStorageBase);
+
+  Operator op;
+
+  op(params, *shared_storage);
+}
+
+
+/// Generic MUTLASS kernel template.
+template <typename Operator>
+MUTLASS_GLOBAL
+void Kernel2(typename Operator::Params params) {
+  // Dynamic shared memory base pointer
+  extern __shared__ int SharedStorageBase[];
+  // Declare pointer to dynamic shared memory.
+  typename Operator::SharedStorage *shared_storage =
+      reinterpret_cast<typename Operator::SharedStorage *>(SharedStorageBase);
+
+  Operator::invoke(params, *shared_storage);
+}
+
+
+////////////////////////////////////////////////////////////////////////////////
+//
+// 3.0 specific launch
+//
+////////////////////////////////////////////////////////////////////////////////
+
+/// Generic MUTLASS kernel template.
+template <typename Operator>
+MUTLASS_GLOBAL
+#ifdef __MUSACC__
+// Enclosing this in __MUSACC__ suppresses MSVC warnings.
+__launch_bounds__(Operator::MaxThreadsPerBlock, Operator::MinBlocksPerMultiprocessor)
+#endif // __MUSACC__
+void device_kernel(MUTLASS_GRID_CONSTANT typename Operator::Params const params)
+{
+  // Dynamic shared memory base pointer
+  extern __shared__ char __attribute__((aligned(Operator::SmemAlignmentBytes))) smem[];
+  Operator op;
+  op(params, smem);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+} /// namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp22_builder.inl b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp22_builder.inl
new file mode 100644
index 000000000..1d295a536
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp22_builder.inl
@@ -0,0 +1,105 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mute/atom/mma_atom.hpp"
+#include "mute/atom/copy_atom.hpp"
+
+#include "mutlass/mutlass.h"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/arch/arch.h"
+#include "mutlass/arch/mma.h"
+#include "mutlass/layout/layout.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+#include "mutlass/epilogue/collective/collective_epilogue.hpp"
+#include "mutlass/epilogue/collective/default_epilogue.hpp"
+#include "mutlass/epilogue/thread/linear_combination.h"
+
+namespace mutlass::epilogue::collective {
+
+template <
+  class OpClass,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class EpilogueTileType,
+  class ElementAccumulator,
+  class ElementCompute,
+  class ElementC_,
+  class GmemLayoutTagC_,
+  int AlignmentC,
+  class ElementD,
+  class GmemLayoutTagD,
+  int AlignmentD,
+  class EpilogueScheduleType,
+  class MainloopCollective,
+  FloatRoundStyle RoundStyle
+>
+struct CollectiveBuilder<
+  mutlass::arch::Mp22,
+  OpClass,
+  TileShape_MNK,
+  ClusterShape_MNK,
+  EpilogueTileType,
+  ElementAccumulator,
+  ElementCompute,
+  ElementC_,
+  GmemLayoutTagC_,
+  AlignmentC,
+  ElementD,
+  GmemLayoutTagD,
+  AlignmentD,
+  EpilogueScheduleType,
+  fusion::LinearCombination<ElementD,ElementCompute,ElementC_,ElementCompute,RoundStyle>,
+  MainloopCollective,
+  void
+> {
+  // Passing void C disables source load
+  using ElementC = mute::conditional_t<mute::is_void_v<ElementC_>,
+    ElementD, ElementC_>; // prevents mute breakages
+  using GmemLayoutTagC = mute::conditional_t<mute::is_void_v<ElementC_>,
+    GmemLayoutTagD, GmemLayoutTagC_>;
+  static constexpr thread::ScaleType::Kind ScaleType = mute::is_void_v<ElementC_> ?
+    thread::ScaleType::OnlyAlphaScaling : thread::ScaleType::Default;
+
+  static constexpr int FragmentSize = 1;
+  using ThreadOp = thread::LinearCombination<
+    ElementD, FragmentSize, ElementAccumulator, ElementCompute,
+    ScaleType, FloatRoundStyle::round_to_nearest, ElementC>;
+
+  using CollectiveOp = mutlass::epilogue::collective::DefaultEpilogue<
+                        mutlass::detail::TagToStrideC_t<GmemLayoutTagC>,
+                        mutlass::detail::TagToStrideC_t<GmemLayoutTagD>,
+                        ThreadOp,
+                        mutlass::gemm::EpilogueDefault
+                      >;
+};
+
+} // namespace mutlass::epilogue::collective
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp31_builder.inl b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp31_builder.inl
new file mode 100644
index 000000000..5cc2656db
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp31_builder.inl
@@ -0,0 +1,373 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Functor performing elementwise operations used by epilogues.
+*/
+
+#pragma once
+
+#include "mute/atom/mma_traits_mp31.hpp"
+#include "mute/atom/mma_traits_mp31_sqmma.hpp"
+#include "mute/atom/copy_traits_mp31.hpp"
+
+#include "mutlass/detail/dependent_false.hpp"
+#include "mutlass/detail/layout.hpp"
+#include "mutlass/epilogue/dispatch_policy.hpp"
+#include "mutlass/epilogue/collective/collective_epilogue.hpp"
+#include "mutlass/epilogue/thread/linear_combination.h"
+#include "mutlass/epilogue/fusion/callbacks.hpp"
+#include "mutlass/epilogue/collective/builders/mp31_builder_common.inl"
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/type_traits>
+#else
+#include <type_traits>
+#endif
+using namespace mute;
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::collective {
+
+///////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+// Returns the parameterized dispatch policy for the TME epilogue
+template<class TileShapeMNK, class EpilogueTileMN,  class ElementC, class ElementD>
+constexpr auto
+mp31_get_tme_dispatch_policy() {
+  using AccAtomM = _16;
+  using AccAtomN =  _8;
+  auto EpilogueShapeMN = make_shape(size<0>(EpilogueTileMN{}), size<1>(EpilogueTileMN{}));
+  constexpr int EpiTiles = size(shape_div(take<0,2>(TileShapeMNK{}), EpilogueShapeMN));
+  constexpr int AtomM = get<0>(shape<0>(EpilogueTileMN{}));
+  constexpr int AtomN = get<0>(shape<1>(EpilogueTileMN{}));
+  constexpr int FragmentSize = (AtomN / AccAtomN{}) * (AtomM / AccAtomM{});
+  constexpr bool ReuseSmem = (sizeof_bits_v<ElementC> == sizeof_bits_v<ElementD>) && (sizeof_bits_v<ElementD> > 8);
+  constexpr int StagesD = mute::min(EpiTiles, 3);
+  constexpr int StagesC = ReuseSmem ? mute::max(mute::min(EpiTiles, 3), StagesD)
+                                    : mute::min(EpiTiles, 2);
+  constexpr bool DelayTmeStore = (mute::is_void_v<ElementC> || not ReuseSmem) && StagesD > 1;
+  return Mp31CollectiveEpilogue<StagesC, StagesD, FragmentSize, ReuseSmem, DelayTmeStore>{};
+}
+
+// Returns the smem layout atom to be used for C or D matrix
+template<class GmemStrideType, class Element, class EpilogueTile_MN>
+constexpr auto
+mp31_get_epilogue_smem_swizzle_layout_atom() {
+  // ColMajor C/D (M-major)
+  if constexpr (mutlass::gemm::detail::is_major<0>(GmemStrideType{})) {
+    return mutlass::epilogue::collective::detail::ss_smem_selector_C<
+      mute::TCE::Major::MN, Element, EpilogueTile_MN
+    >();
+  }
+  // RowMajor C/D (N-major)
+  else if constexpr (mutlass::gemm::detail::is_major<1>(GmemStrideType{})) {
+    return mutlass::epilogue::collective::detail::ss_smem_selector_C<
+      mute::TCE::Major::K , Element, EpilogueTile_MN
+    >();
+  }
+  else {
+    static_assert(mutlass::detail::dependent_false<GmemStrideType>, "Unsupported gmem layout.");
+  }
+}
+
+// Attempts to compute a reasonable epilogue tile based on block tile shape or allows the user to provide one.
+template <class TileShape_MNK, class AtomShape_MNK>
+constexpr auto
+mp31_compute_tile_shape() {
+  constexpr int tile_m = size<0>(TileShape_MNK{});
+  constexpr int tile_n = size<1>(TileShape_MNK{});
+  constexpr int atom_tile_m = size<0>(AtomShape_MNK{});
+  constexpr int atom_tile_n = size<1>(AtomShape_MNK{});
+  constexpr int atom_m = tile_m / atom_tile_m;
+  constexpr int atom_n = tile_n / atom_tile_n;
+
+  if constexpr (atom_m > 1 && atom_n == 1) {
+    using tiler_m = Layout<Shape <Int<atom_tile_m>, Int<atom_m>>,
+                           Stride<              _1, Int<atom_tile_m>>>;
+    using tiler_n = Layout<Shape <Int<8>, Int<atom_n>>,
+                           Stride<    _1, Int<atom_tile_n>>>;
+    return make_tile(tiler_m{}, tiler_n{});
+  }
+  else {
+    using tiler_m = Layout<Shape <_16, Int<atom_m>>,
+                           Stride< _1, Int<atom_tile_m>>>;
+    using tiler_n = Layout<Shape <Int<atom_tile_n>, Int<atom_n>>,
+                           Stride<              _1, Int<atom_tile_n>>>;
+    return make_tile(tiler_m{}, tiler_n{});
+  }
+
+}
+
+// Helper for building TME collective epilogues, specialized by
+// the fusion operation performed and the dispatch policy to use.
+template <
+  class TileShape_MNK,
+  class EpilogueTile_MN,
+  class ElementAccumulator,
+  class ElementCompute,
+  class ElementC_,
+  class GmemLayoutTagC_,
+  int AlignmentC,
+  class ElementD_,
+  class GmemLayoutTagD,
+  int AlignmentD,
+  class ThreadOp,
+  class DispatchPolicy
+>
+struct Mp31TmeBuilderImpl {
+
+  using ElementD = ElementD_;
+  // Passing void C disables source load + smem allocation
+  using ElementC = mute::conditional_t<mute::is_void_v<ElementC_>,ElementD,ElementC_>; // prevents void ref breakages
+  using GmemLayoutTagC = mute::conditional_t<mute::is_void_v<ElementC_>,GmemLayoutTagD,GmemLayoutTagC_>;
+
+  using GmemStrideTypeC = mutlass::detail::TagToStrideC_t<GmemLayoutTagC>;
+  using GmemStrideTypeD = mutlass::detail::TagToStrideC_t<GmemLayoutTagD>;
+
+  using UnderlyingGmemStrideTypeC = mute::remove_pointer_t<GmemStrideTypeC>;
+  using UnderlyingGmemStrideTypeD = mute::remove_pointer_t<GmemStrideTypeD>;
+
+  using CopyOpS2G = MP31_TME_STORE;
+  using CopyOpG2S = MP31_TME_LOAD;
+
+  using CopyOpS2R = Copy_Atom<DefaultCopy,ElementC>;
+  using CopyOpR2S = Copy_Atom<UniversalCopy<ElementD>, ElementD>;
+
+  using CollectiveOp = mutlass::epilogue::collective::CollectiveEpilogue<
+      DispatchPolicy,
+      ThreadOp,
+      TileShape_MNK,
+      EpilogueTile_MN,
+      ElementC_, // Need to pass void through to expose via GemmUniversal
+      GmemStrideTypeC,
+      ElementD,
+      GmemStrideTypeD,
+      CopyOpG2S,
+      CopyOpS2G,
+      decltype(detail::mp31_get_epilogue_smem_swizzle_layout_atom<UnderlyingGmemStrideTypeC, ElementC, EpilogueTile_MN>()),
+      decltype(detail::mp31_get_epilogue_smem_swizzle_layout_atom<UnderlyingGmemStrideTypeD, ElementD, EpilogueTile_MN>()),
+      CopyOpS2R,
+      CopyOpR2S
+    >;
+};
+} // namespace detail
+
+// No-smem builder
+template <
+  class OpClass,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class EpilogueTileType,
+  class ElementAccumulator,
+  class ElementCompute,
+  class ElementC_,
+  class GmemLayoutTagC_,
+  int AlignmentC,
+  class ElementD,
+  class GmemLayoutTagD,
+  int AlignmentD,
+  class Schedule,
+  class MainloopCollective,
+  FloatRoundStyle RoundStyle
+>
+struct CollectiveBuilder<
+    arch::Mp31,
+    OpClass,
+    TileShape_MNK,
+    ClusterShape_MNK,
+    EpilogueTileType,
+    ElementAccumulator,
+    ElementCompute,
+    ElementC_,
+    GmemLayoutTagC_,
+    AlignmentC,
+    ElementD,
+    GmemLayoutTagD,
+    AlignmentD,
+    Schedule,
+    fusion::LinearCombination<ElementD,ElementCompute,ElementC_,ElementCompute,RoundStyle>,
+    MainloopCollective,
+    mute::enable_if_t<mute::is_same_v<Schedule, NoSmem>>> {
+
+  // Passing void C disables source load
+  using ElementC = mute::conditional_t<mute::is_void_v<ElementC_>,
+      ElementD, ElementC_>; // prevents mute breakages
+  using GmemLayoutTagC = mute::conditional_t<mute::is_void_v<ElementC_>,
+      GmemLayoutTagD, GmemLayoutTagC_>;
+  static constexpr thread::ScaleType::Kind ScaleType = mute::is_void_v<ElementC_> ?
+      thread::ScaleType::OnlyAlphaScaling : thread::ScaleType::Default;
+  static constexpr int FragmentSize = 1;
+  using ThreadOp = thread::LinearCombination<
+    ElementD, FragmentSize, ElementAccumulator, ElementCompute,
+    ScaleType, RoundStyle, ElementC>;
+
+  using CollectiveOp = mutlass::epilogue::collective::DefaultEpilogue<
+        mutlass::detail::TagToStrideC_t<GmemLayoutTagC>,
+        mutlass::detail::TagToStrideC_t<GmemLayoutTagD>,
+        ThreadOp,
+        mutlass::gemm::EpilogueDefault>;
+};
+
+// Tme builder
+template <
+  class OpClass,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class EpilogueTileType,
+  class ElementAccumulator,
+  class ElementCompute,
+  class ElementC,
+  class GmemLayoutTagC,
+  int AlignmentC,
+  class ElementD,
+  class GmemLayoutTagD,
+  int AlignmentD,
+  class Schedule,
+  class MainloopCollective
+>
+struct CollectiveBuilder<
+    arch::Mp31,
+    OpClass,
+    TileShape_MNK,
+    ClusterShape_MNK,
+    EpilogueTileType,
+    ElementAccumulator,
+    ElementCompute,
+    ElementC,
+    GmemLayoutTagC,
+    AlignmentC,
+    ElementD,
+    GmemLayoutTagD,
+    AlignmentD,
+    Schedule,
+    fusion::LinearCombination<ElementD,ElementCompute,ElementC,ElementCompute>,
+    MainloopCollective,
+    mute::enable_if_t<mute::is_same_v<Schedule, WithTme> &&
+    mute::is_void_v<mute::void_t<typename MainloopCollective::TiledMma>>>
+    > {
+private:
+
+  using TiledMma = typename MainloopCollective::TiledMma;
+  using AtomShape_MNK = typename TiledMma::AtomShape_MNK;
+  using ThreadOp = thread::LinearCombination<ElementD, 1, ElementAccumulator, ElementCompute>;
+  using EpilogueTile_MN = decltype(detail::mp31_compute_tile_shape<TileShape_MNK, AtomShape_MNK>());
+  using DispatchPolicy =
+    decltype(detail::mp31_get_tme_dispatch_policy<TileShape_MNK, EpilogueTile_MN, ElementC, ElementD>());
+
+public:
+  using CollectiveOp =
+    typename detail::Mp31TmeBuilderImpl<
+      TileShape_MNK,
+      EpilogueTile_MN,
+      ElementAccumulator,
+      ElementCompute,
+      ElementC,
+      GmemLayoutTagC,
+      AlignmentC,
+      ElementD,
+      GmemLayoutTagD,
+      AlignmentD,
+      ThreadOp,
+      DispatchPolicy
+    >::CollectiveOp;
+};
+
+// Auto builder
+template <
+  class OpClass,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class EpilogueTileType,
+  class ElementAccumulator,
+  class ElementCompute,
+  class ElementC,
+  class GmemLayoutTagC,
+  int AlignmentC,
+  class ElementD,
+  class GmemLayoutTagD,
+  int AlignmentD,
+  class FusionOperation,
+  class MainloopCollective
+>
+struct CollectiveBuilder<
+    arch::Mp31,
+    OpClass,
+    TileShape_MNK,
+    ClusterShape_MNK,
+    EpilogueTileType,
+    ElementAccumulator,
+    ElementCompute,
+    ElementC,
+    GmemLayoutTagC,
+    AlignmentC,
+    ElementD,
+    GmemLayoutTagD,
+    AlignmentD,
+    EpilogueScheduleAuto,
+    FusionOperation,
+    MainloopCollective,
+    void> {
+private:
+  static_assert(mute::is_same_v<FusionOperation, fusion::LinearCombination<ElementD,ElementCompute,ElementC,ElementCompute>>,
+                "Auto schedule doesn't support fusion. Use one of the TME schedules instead.");
+
+  // Pick No-Smem epilogue as the Auto Epilogue Schedule (Auto schedules do not guarantee best performance)
+  // since TME epilogues are not compatible with non-TME non-WS mainloops
+  using EpilogueSchedule = NoSmem;
+  using _CollectiveBuilder = CollectiveBuilder<
+    arch::Mp31,
+    OpClass,
+    TileShape_MNK,
+    ClusterShape_MNK,
+    EpilogueTileType,
+    ElementAccumulator,
+    ElementCompute,
+    ElementC,
+    GmemLayoutTagC,
+    AlignmentC,
+    ElementD,
+    GmemLayoutTagD,
+    AlignmentD,
+    EpilogueSchedule,
+    FusionOperation,
+    MainloopCollective
+  >;
+
+public:
+  using CollectiveOp = typename _CollectiveBuilder::CollectiveOp;
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::collective
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp31_builder_common.inl b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp31_builder_common.inl
new file mode 100644
index 000000000..d670ebf81
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/builders/mp31_builder_common.inl
@@ -0,0 +1,167 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mute/tensor.hpp"
+#include "mute/atom/copy_traits_mp31_tme_swizzle.hpp"
+#include "mute/atom/mma_atom.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+using namespace mute;
+
+template <
+  int MN,
+  int K,
+  class Type,
+  TCE::Major tnsp,
+  class SmemAtomLayout,
+  TME::SmemSwizzleGranularity sg = TME::SmemSwizzleGranularity::NONE,
+  TME::SmemSwizzleStride      ss = TME::SmemSwizzleStride::B128,
+  TME::SmemSwizzleLine        sl = TME::SmemSwizzleLine::B128
+>
+MUTE_HOST_DEVICE constexpr
+auto
+get_layout_swizzle_Atom() {
+  constexpr int BITS = sizeof_bits_v<Type>;
+  using swizzle = decltype(mute::detail::get_tme_swizzle<sl, ss, sg>());
+  using layout_atom_bits = mute::conditional_t<tnsp == TCE::Major::MN,
+                            decltype(make_layout(Shape<Int<MN * sizeof_bits_v<Type>>, Int<K>>{}, LayoutLeft{})),
+                            decltype(make_layout(Shape<Int<MN>, Int<K * sizeof_bits_v<Type>>>{}, LayoutRight{}))>;
+  return upcast<BITS>(ComposedLayout<
+              swizzle,
+              smem_ptr_flag,
+              decltype(blocked_product(
+              layout_atom_bits{},
+              SmemAtomLayout{}))>{});
+};
+
+template <class Type, TCE::Major tnsp, class AtomM, class AtomN, class SmemAtomLayout>
+MUTE_HOST_DEVICE constexpr
+auto
+make_tme_smem_atom_layout_C() {
+  constexpr int BITS = sizeof_bits_v<Type>;
+  constexpr int ATOM_M = size(AtomM{});
+  constexpr int ATOM_N = size(AtomN{});
+  constexpr int BYTE = sizeof(Type);
+  static_assert(BITS <= 32, "Unsupported Tme Type");
+  if constexpr (tnsp == TCE::Major::MN) {
+    constexpr auto SL = ATOM_M * BYTE <= 128 ? TME::SmemSwizzleLine::B128 : TME::SmemSwizzleLine::B256;
+    if constexpr (shape<0>(SmemAtomLayout{}) > 1 && shape<1>(SmemAtomLayout{}) == 1) {
+      return get_layout_swizzle_Atom<ATOM_M, ATOM_N, Type, tnsp, SmemAtomLayout,
+        TME::SmemSwizzleGranularity::B16, TME::SmemSwizzleStride::B128, SL>();
+    } else {
+      // S4/S8/U8/FP8
+      if constexpr (BITS <= 8) {
+        return get_layout_swizzle_Atom<ATOM_M, ATOM_N, Type, tnsp, SmemAtomLayout,
+          TME::SmemSwizzleGranularity::NONE, TME::SmemSwizzleStride::B128, TME::SmemSwizzleLine::B256>();
+      }
+      // FP16/BF16
+      else if constexpr (BITS == 16) {
+        return get_layout_swizzle_Atom<ATOM_M, ATOM_N, Type, tnsp, SmemAtomLayout,
+            TME::SmemSwizzleGranularity::B16, TME::SmemSwizzleStride::B32, TME::SmemSwizzleLine::B128>();
+      }
+      // TF32
+      else {
+        return get_layout_swizzle_Atom<ATOM_M, ATOM_N, Type, tnsp, SmemAtomLayout,
+            TME::SmemSwizzleGranularity::B16, TME::SmemSwizzleStride::B64, TME::SmemSwizzleLine::B128>();
+      }
+    }
+  }
+  else if constexpr (tnsp == TCE::Major::K) {
+    constexpr auto SL = ATOM_N * BYTE <= 128 ? TME::SmemSwizzleLine::B128 : TME::SmemSwizzleLine::B256;
+    if constexpr (shape<0>(SmemAtomLayout{}) > 1 && shape<1>(SmemAtomLayout{}) == 1) {
+      return get_layout_swizzle_Atom<ATOM_M, ATOM_N, Type, tnsp, SmemAtomLayout,
+        TME::SmemSwizzleGranularity::NONE, TME::SmemSwizzleStride::B128, TME::SmemSwizzleLine::B256>();
+    } else {
+      return get_layout_swizzle_Atom<ATOM_M, ATOM_N, Type, tnsp, SmemAtomLayout,
+        TME::SmemSwizzleGranularity::B32, TME::SmemSwizzleStride::B128, SL>();
+    }
+  }
+  else {
+    static_assert(tnsp != TCE::Major::MN && tnsp != TCE::Major::K, "Unrecognized MajorMode!");
+  }
+}
+
+namespace mutlass::epilogue::collective::detail {
+template <TCE::Major major, class ElementType, class EpilogueTileMN>
+MUTE_HOST_DEVICE constexpr
+auto
+ss_smem_selector_C()
+{
+  using AtomN = decltype(get<0>(shape<1>(EpilogueTileMN{})));
+  using AtomM = decltype(get<0>(shape<0>(EpilogueTileMN{})));
+
+  constexpr int smem_atom_m = size<1>(shape<0>(EpilogueTileMN{}));
+  constexpr int smem_atom_n = size<1>(shape<1>(EpilogueTileMN{}));
+
+  using SmemAtomLayout = Layout<Shape <Int<smem_atom_m>, Int<smem_atom_n>>,
+                                Stride<Int<smem_atom_n>,               _1>>;
+
+  return make_tme_smem_atom_layout_C<ElementType, major, AtomM, AtomN, SmemAtomLayout>();
+}
+
+template<class CopyOp, class CopyV, class EpilogueTile, class... MArgs>
+MUTLASS_HOST_DEVICE constexpr
+auto make_tiled_copy_C_sqmma(gemm::TiledMMA<MArgs...> const& mma){
+  auto layoutC_TV = mma.get_layoutC_TV();
+  MUTE_STATIC_ASSERT_V(CopyV{} <= size<1>(layoutC_TV));
+  constexpr int AtomM = get<1>(shape<0>(EpilogueTile{}));
+  constexpr int AtomN = get<1>(shape<1>(EpilogueTile{}));
+  constexpr int stride_v = (AtomM > 1 && AtomN == 1) ? get<2>(get<0>(stride<1>(layoutC_TV))) : get<0>(get<0>(stride<1>(layoutC_TV)));
+  auto layout_TV = make_layout(make_shape(shape<0>(layoutC_TV), CopyV{}),  make_stride(stride<0>(layoutC_TV), Int<stride_v>{}));
+
+  auto mma_tiler = make_shape(tile_size<0>(mma), tile_size<1>(mma));
+  auto mma_zeros = repeat_like(mma_tiler, Int<0>{});
+
+  auto orig_tiler = mute::transform(make_seq<rank(mma_tiler)>{}, [&](auto i) {
+    return filter(composition(make_layout(mma_tiler, replace<i>(mma_zeros, Int<1>{})), layout_TV));
+  });
+
+  auto orig_tiler_m = get<0>(EpilogueTile{});
+  auto orig_tiler_n = get<1>(EpilogueTile{});
+
+  auto tiler_n = append<3>(orig_tiler_n, Layout<_1, _0>{});
+  auto permute_tiler_n = make_layout(get<0>(tiler_n), get<2>(tiler_n), get<1>(tiler_n));
+
+  auto tiler = make_tile(orig_tiler_m, permute_tiler_n);
+  auto tile2mma = composition(make_layout(mma_tiler), tiler);
+
+  auto tiler_mn = make_tile(make_layout(size<0>(tiler)),
+                            make_layout(size<1>(tiler)));
+
+  auto layout_tv = composition(left_inverse(tile2mma), layout_TV);
+  return make_tiled_copy_impl(CopyOp{}, layout_tv, tiler_mn);
+}
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::collective::detail
\ No newline at end of file
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/collective_builder.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/collective_builder.hpp
new file mode 100644
index 000000000..8dfe6164c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/collective_builder.hpp
@@ -0,0 +1,124 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mute/arch/copy.hpp>         // mute::DefaultCopy
+#include <mute/util/type_traits.hpp>  // mute::is_base_of_v
+
+#include "mutlass/detail/dependent_false.hpp"
+#include "mutlass/epilogue/fusion/callbacks.hpp"
+#include "mutlass/epilogue/dispatch_policy.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Used to specify epilogue subtile shape or dispatch to automatic computation of subtile shape
+struct EpilogueTileAuto {};
+
+// Used to let the builder pick the epilogue schedule automatically.
+// Can be overridden with kernel schedule tags in mutlass/epilogue/dispatch_policy.hpp
+struct EpilogueScheduleAuto {};
+struct EpilogueIm2ColScheduleAuto {};
+
+template <
+  class ArchTag,
+  class OpClass,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class EpilogueTileType,
+  class ElementAccumulator,
+  class ElementCompute,
+  class ElementC,
+  class GmemLayoutTagC,
+  int AlignmentC,
+  class ElementD,
+  class GmemLayoutTagD,
+  int AlignmentD,
+  class EpilogueScheduleType,
+  class FusionOpOrCallbacks = fusion::LinearCombination<ElementD,ElementCompute>,
+  class MainloopCollective = void,
+  class Enable = void
+>
+struct CollectiveBuilder {
+  static_assert(mutlass::detail::dependent_false<ArchTag>,
+      "Could not build a collective epilogue for given parameters.");
+};
+
+// helper sub-builder for epilogue fusion callbacks (for internal use by CollectiveBuilder only)
+namespace detail {
+
+// callbacks builder with operation tag
+template<
+  class DispatchPolicy,
+  class FusionOp,
+  class TileShape_MNK,
+  class EpilogueTile_MN,
+  class ElementAccumulator,
+  class = void
+>
+struct CallbacksBuilder {
+  using Callbacks = fusion::FusionCallbacks<DispatchPolicy, FusionOp, TileShape_MNK, EpilogueTile_MN>;
+};
+
+// callbacks builder with callbacks passthrough
+template <
+  class DispatchPolicy,
+  class FusionCallbacks,
+  class TileShape_MNK,
+  class EpilogueTile_MN,
+  class ElementAccumulator
+>
+struct CallbacksBuilder<
+  DispatchPolicy,
+  FusionCallbacks,
+  TileShape_MNK,
+  EpilogueTile_MN,
+  ElementAccumulator,
+  mute::enable_if_t<not mute::is_base_of_v<fusion::FusionOperation, FusionCallbacks>>
+> {
+  using Callbacks = FusionCallbacks;
+};
+
+} // namespace detail
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "builders/mp22_builder.inl"
+#include "builders/mp31_builder.inl"
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/collective_epilogue.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/collective_epilogue.hpp
new file mode 100644
index 000000000..edc80675a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/collective_epilogue.hpp
@@ -0,0 +1,61 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <mutlass/detail/dependent_false.hpp>
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class DispatchPolicy,
+  class... Args
+>
+class CollectiveEpilogue {
+  static_assert(mutlass::detail::dependent_false<DispatchPolicy>, "Could not find an epilogue specialization.");
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "detail.hpp"
+#include "default_epilogue.hpp"
+#include "epilogue_tensor_broadcast.hpp"
+#include "mp31_epilogue_tme.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/default_epilogue.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/default_epilogue.hpp
new file mode 100644
index 000000000..267b0ffd1
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/default_epilogue.hpp
@@ -0,0 +1,244 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Functor performing elementwise operations used by epilogues.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+#include "mutlass/epilogue/collective/detail.hpp"
+
+#include "mute/tensor.hpp"
+#include "mute/numeric/numeric_types.hpp"
+#include "mutlass/musa_host_adapter.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Applies an element wise operation to all elements within the fragment
+/// and writes them out to destination storage.
+template <
+  class StrideC_,
+  class StrideD_,
+  class ThreadEpilogueOp_,
+  class EpilogueSchedule_
+>
+class DefaultEpilogue {
+public:
+  //
+  // Type Aliases
+  //
+  using EpilogueSchedule = EpilogueSchedule_;
+  using DispatchPolicy = EpilogueSchedule_;
+
+  // derived types of output thread level operator
+  using ThreadEpilogueOp = ThreadEpilogueOp_;
+  using ElementOutput = typename ThreadEpilogueOp::ElementOutput;
+  using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator;
+  using ElementCompute = typename ThreadEpilogueOp::ElementCompute;
+  using ElementScalar = ElementCompute;
+  using ElementC = typename ThreadEpilogueOp::ElementC;
+  using StrideC = StrideC_;
+  using ElementD = typename ThreadEpilogueOp::ElementD;
+  using StrideD = StrideD_;
+
+  using GmemTiledCopyC = void;
+  using GmemTiledCopyD = void;
+
+  static const int kOutputAlignment = ThreadEpilogueOp::kCount;
+  using AlignmentType = typename mute::uint_bit<sizeof_bits<ElementOutput>::value * kOutputAlignment>::type;
+  static constexpr int NumBarriers = 0;
+
+  static_assert(mute::rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]");
+  static_assert(mute::rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]");
+
+  struct SharedStorage { };
+
+  using TensorStorage = SharedStorage;
+
+  // Host side epilogue arguments
+  struct Arguments {
+    typename ThreadEpilogueOp::Params thread{};
+    ElementC const* ptr_C = nullptr;
+    StrideC dC{};
+    ElementD* ptr_D = nullptr;
+    StrideD dD{};
+  };
+
+  // Device side epilogue params
+  using Params = Arguments;
+
+  //
+  // Methods
+  //
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(
+      [[maybe_unused]] ProblemShape const& _,
+      Arguments const& args,
+      [[maybe_unused]] void* workspace) {
+    return args;
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return 0;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream,
+    MusaHostAdapter* musa_adapter = nullptr) {
+    return mutlass::Status::kSuccess;
+  }
+
+  template<class ProblemShape>
+  MUTLASS_HOST_DEVICE static bool
+  can_implement(
+      [[maybe_unused]] ProblemShape const& problem_shape,
+      [[maybe_unused]] Arguments const& args) {
+    return true;
+  }
+
+  // Note: SharedStorage is unused for DefaultEpilogue
+  MUTLASS_HOST_DEVICE
+  DefaultEpilogue(Params const& params_, SharedStorage const& shared_storage = SharedStorage())
+      : params(params_), epilogue_op(params_.thread) { }
+
+  MUTLASS_DEVICE
+  bool
+  is_source_needed() {
+    return epilogue_op.is_source_needed();
+  }
+
+  template<
+    class ProblemShapeMNKL,
+    class BlockShapeMNK,
+    class BlockCoordMNKL,
+    class FrgEngine, class FrgLayout,
+    class TiledMma,
+    class ResidueMNK
+  >
+  MUTLASS_HOST_DEVICE void
+  operator()(
+      ProblemShapeMNKL problem_shape_mnkl,
+      BlockShapeMNK blk_shape_MNK,
+      BlockCoordMNKL blk_coord_mnkl,
+      mute::Tensor<FrgEngine, FrgLayout> const& accumulators,
+      TiledMma tiled_mma,
+      ResidueMNK residue_mnk,
+      int thread_idx,
+      [[maybe_unused]] char* smem_buf)
+  {
+    using namespace mute;
+    using X = Underscore;
+
+    static_assert(mute::rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4");
+    static_assert(is_static<BlockShapeMNK>::value, "ThreadBlock tile shape must be static");
+    static_assert(mute::rank(BlockShapeMNK{}) == 3, "BlockShapeMNK must be rank 3");
+    static_assert(mute::rank(BlockCoordMNKL{}) == 4, "BlockCoordMNKL must be rank 3");
+
+    // Separate out problem shape for convenience
+    auto M = get<0>(problem_shape_mnkl);
+    auto N = get<1>(problem_shape_mnkl);
+    auto L = get<3>(problem_shape_mnkl);
+
+    auto stride_c = detail::get_epilogue_stride<EpilogueSchedule>(params.dC);
+    auto stride_d = detail::get_epilogue_stride<EpilogueSchedule>(params.dD);
+
+    // Represent the full output tensor
+    Tensor mC_mnl = make_tensor(make_gmem_ptr(params.ptr_C), make_shape(M,N,L), stride_c);                 // (m,n,l)
+    Tensor mD_mnl = make_tensor(make_gmem_ptr(params.ptr_D), make_shape(M,N,L), stride_d);                 // (m,n,l)
+    Tensor gC_mnl = local_tile(mC_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});    // (BLK_M,BLK_N,m,n,l)
+    Tensor gD_mnl = local_tile(mD_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});    // (BLK_M,BLK_N,m,n,l)
+
+    // Slice to get the tile this CTA is responsible for
+    auto [m_coord, n_coord, k_coord, l_coord] = blk_coord_mnkl;
+    Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord);                                                 // (BLK_M,BLK_N)
+    Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord);                                                 // (BLK_M,BLK_N)
+
+    // Partition source and destination tiles to match the accumulator partitioning
+    auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
+    Tensor tCgD = thr_mma.partition_C(gD);                                       // (VEC,THR_M,THR_N)
+    Tensor tCgC = thr_mma.partition_C(gC);                                       // (VEC,THR_M,THR_N)
+
+    static_assert(is_static<FrgLayout>::value, "Accumulator layout must be static");
+    MUTE_STATIC_ASSERT_V(size(tCgC) == size(tCgD),
+        "Source and destination must have the same number of elements.");
+    MUTE_STATIC_ASSERT_V(size(tCgD) == size(accumulators),
+        "Accumulator count must have the same destination element count.");
+
+    // Make an identity coordinate tensor for predicating our output MN tile
+    auto cD = make_identity_tensor(make_shape(unwrap(shape<0>(gD)), unwrap(shape<1>(gD))));
+    Tensor tCcD = thr_mma.partition_C(cD);
+
+    // source is needed
+    if (epilogue_op.is_source_needed()) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < size(accumulators); ++i) {
+        if (elem_less(tCcD(i), make_coord(get<0>(residue_mnk), get<1>(residue_mnk)))) {
+          tCgD(i) = epilogue_op(accumulators(i), tCgC(i));
+        }
+      }
+    }
+    // source is not needed, avoid load
+    else {
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < size(accumulators); ++i) {
+        if (elem_less(tCcD(i), make_coord(get<0>(residue_mnk), get<1>(residue_mnk)))) {
+          tCgD(i) = epilogue_op(accumulators(i));
+        }
+      }
+    }
+  }
+
+private:
+  Params params;
+  ThreadEpilogueOp epilogue_op;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace collective
+} // namespace epilogue
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/detail.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/detail.hpp
new file mode 100644
index 000000000..a17249b92
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/detail.hpp
@@ -0,0 +1,147 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+#include "mutlass/epilogue/dispatch_policy.hpp"
+
+#include "mute/tensor.hpp"
+#include "mute/numeric/numeric_types.hpp"
+#include "mute/util/type_traits.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace collective {
+
+namespace detail {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class Stride>
+constexpr bool
+is_m_major() {
+  return mutlass::gemm::detail::is_major<0,Stride>();
+}
+
+template <class Stride>
+constexpr bool
+is_n_major() {
+  return mutlass::gemm::detail::is_major<1,Stride>();
+}
+
+template <class Stride>
+constexpr bool
+is_im2col() {
+  return mute::is_same_v<Stride, mutlass::detail::TagToStrideC_t<mutlass::layout::TensorNWC>>
+      || mute::is_same_v<Stride, mutlass::detail::TagToStrideC_t<mutlass::layout::TensorNHWC>>
+      || mute::is_same_v<Stride, mutlass::detail::TagToStrideC_t<mutlass::layout::TensorNDHWC>>;
+}
+
+using mutlass::atomic_maximum;
+
+template <class T>
+static constexpr int elements_per_access_v = mutlass::sizeof_bits<uint32_t>::value / mutlass::sizeof_bits<T>::value;
+
+template <class GmemLayoutTag>
+static constexpr bool is_im2col_mode =
+  mute::is_same_v<GmemLayoutTag, mutlass::layout::TensorNWC> ||
+  mute::is_same_v<GmemLayoutTag, mutlass::layout::TensorNHWC> ||
+  mute::is_same_v<GmemLayoutTag, mutlass::layout::TensorNDHWC>;
+
+template <class T>
+struct EmptyStorage {
+  MUTLASS_HOST_DEVICE
+  T* data() { return nullptr; }
+};
+
+template<class EpilogueSchedule, class Stride>
+MUTLASS_HOST_DEVICE
+auto get_epilogue_stride(Stride stride){
+  if constexpr (mute::is_base_of_v<mutlass::gemm::EpilogueTransposed, EpilogueSchedule>) {
+    return mute::make_stride(mute::get<1>(stride), mute::get<0>(stride), mute::get<2>(stride));
+  }
+  else {
+    return stride;
+  }
+}
+
+template <typename ThreadEpilogueOp, typename = void>
+struct IsThreadEpilogueOpWithBias {
+  static constexpr bool value = false;
+  using type = typename ThreadEpilogueOp::ElementCompute;
+};
+
+template <typename ThreadEpilogueOp>
+struct IsThreadEpilogueOpWithBias <ThreadEpilogueOp, mute::void_t<typename ThreadEpilogueOp::ElementBias>> {
+  static constexpr bool value = true;
+  using type = typename ThreadEpilogueOp::ElementBias;
+};
+template <typename ThreadEpilogueOp, typename = void>
+struct IsThreadEpilogueOpWithPerChannelScaling {
+  static constexpr bool value = false;
+};
+
+template <typename ThreadEpilogueOp>
+struct IsThreadEpilogueOpWithPerChannelScaling <ThreadEpilogueOp, mute::enable_if_t<ThreadEpilogueOp::IsPerChannelScalingSupported>> {
+  static constexpr bool value = true;
+};
+
+template <typename ThreadEpilogueOp, typename = void>
+struct IsThreadEpilogueOpWithActivation {
+  static constexpr bool value = false;
+  using type = void;
+};
+
+template <typename ThreadEpilogueOp>
+struct IsThreadEpilogueOpWithActivation <ThreadEpilogueOp, mute::enable_if_t<ThreadEpilogueOp::IsEltActSupported>> {
+  static constexpr bool value = true;
+  using type = typename ThreadEpilogueOp::ActivationFn;
+};
+
+template <typename ThreadEpilogueOp, typename = void>
+struct IsThreadEpilogueOpWithElementwiseArguments : mute::false_type {};
+
+template <typename ThreadEpilogueOp>
+struct IsThreadEpilogueOpWithElementwiseArguments<
+        ThreadEpilogueOp,
+        mute::void_t<typename ThreadEpilogueOp::ElementwiseOp::Arguments>> : mute::true_type {};
+
+
+} // namespace detail
+} // namespace collective
+} // namespace epilogue
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/epilogue_tensor_broadcast.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/epilogue_tensor_broadcast.hpp
new file mode 100644
index 000000000..d2556a586
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/epilogue_tensor_broadcast.hpp
@@ -0,0 +1,273 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Functor for performing tensor-tensor broadacasts atop existing epilogues.
+
+  Concretely, the opeartion performed is the following:
+    UnaryOp(
+        BinaryOp1(
+            BinaryOp0(
+                Activation((alpha * A @ B) + bias),
+                beta * C0
+            ),
+            beta * C1
+        )
+    )
+
+    where:
+        - C0 and C1 have the same extents as the output
+        - BinaryOp0 and BinaryOp1 perform elementwise binary operations
+        - UnaryOp is an elementwise operation
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/epilogue/collective/detail.hpp"
+
+#include "mute/tensor.hpp"
+#include "mutlass/musa_host_adapter.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace collective {
+/////////////////////////////////////////////////////////////////////////////////////////////////
+using namespace mute;
+/// Collective epilogue that applies elementwise tensor-tensor operations atop other epilogues
+///
+template <
+  class StrideC_,
+  class StrideD_,
+  class ThreadEpilogueOp_,
+  class EpilogueSchedule_,
+  bool PerColumnBias_ = false
+>
+class EpilogueTensorBroadcast {
+public:
+  //
+  // Type Aliases
+  //
+  using EpilogueSchedule = EpilogueSchedule_;
+
+  // derived types of output thread level operator
+  using ThreadEpilogueOp = ThreadEpilogueOp_;
+  using ElementOutput = typename ThreadEpilogueOp::ElementOutput;
+  using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator;
+  using ElementCompute = typename ThreadEpilogueOp::ElementCompute;
+  using ElementScalar = ElementCompute;
+  using ElementBias = typename ThreadEpilogueOp::ElementBias;
+  using ElementC = typename ThreadEpilogueOp::ElementC;
+  using StrideC = StrideC_;
+  using ElementD = typename ThreadEpilogueOp::ElementD;
+  using StrideD = StrideD_;
+  using ActivationFunctor = typename ThreadEpilogueOp::ActivationFunctor;
+
+  static_assert(mute::rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]");
+  static_assert(mute::rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]");
+
+  static constexpr int kOutputAlignment = ThreadEpilogueOp::kCount;
+  using AlignmentType = typename mute::uint_bit<sizeof_bits<ElementOutput>::value * kOutputAlignment>::type;
+
+  static constexpr bool IsBinaryOp0Enabled = ThreadEpilogueOp::IsBinaryOp0Enabled;
+  static constexpr bool IsBinaryOp1Enabled = ThreadEpilogueOp::IsBinaryOp1Enabled;
+  static constexpr bool IsUnaryOpEnabled = ThreadEpilogueOp::IsUnaryOpEnabled;
+
+  static constexpr bool PerColumnBias = PerColumnBias_;
+  using BiasStride = typename mute::conditional_t<PerColumnBias, Stride<_0, _1, _0>, Stride<_1, _0, _0>>;
+  static constexpr int NumBarriers = 0;
+
+  struct SharedStorage { };
+
+  // Host side epilogue arguments
+  struct Arguments {
+    typename ThreadEpilogueOp::Params thread{};
+    StrideC dC{};
+    ElementD* ptr_D = nullptr;
+    StrideD dD{};
+    ElementBias* ptr_Bias = nullptr;
+    ElementC* ptr_C0 = nullptr;
+    ElementC* ptr_C1 = nullptr;
+  };
+
+  // Device side epilogue params
+  using Params = Arguments;
+
+  //
+  // Methods
+  //
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(
+      [[maybe_unused]] ProblemShape const& _,
+      Arguments const& args,
+      [[maybe_unused]] void* workspace) {
+    return args;
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return 0;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream,
+    MusaHostAdapter* musa_adapter = nullptr) {
+    return mutlass::Status::kSuccess;
+  }
+
+  template <class ProblemShape>
+  MUTLASS_HOST_DEVICE static bool
+  can_implement(
+      [[maybe_unused]] ProblemShape const& problem_shape,
+      [[maybe_unused]] Arguments const& args) {
+    return true;
+  }
+
+  MUTLASS_HOST_DEVICE
+  EpilogueTensorBroadcast(Params const& params_)
+      : params(params_), epilogue_op(params_.thread) { }
+
+  MUTLASS_DEVICE
+  bool
+  is_source_needed() {
+    return epilogue_op.is_source0_needed() || epilogue_op.is_source1_needed();
+  }
+
+  template<
+    class ProblemShapeMNKL,
+    class BlockShapeMNK,
+    class BlockCoordMNKL,
+    class FrgEngine, class FrgLayout,
+    class TiledMma,
+    class ResidueMNK
+  >
+  MUTLASS_HOST_DEVICE void
+  operator()(
+      ProblemShapeMNKL problem_shape_mnkl,
+      BlockShapeMNK blk_shape_MNK,
+      BlockCoordMNKL blk_coord_mnkl,
+      mute::Tensor<FrgEngine, FrgLayout> const& accumulators,
+      TiledMma tiled_mma,
+      ResidueMNK residue_mnk,
+      int thread_idx,
+      [[maybe_unused]] char* smem_buf)
+  {
+    using namespace mute;
+    using X = Underscore;
+
+    static_assert(mute::rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4");
+    static_assert(is_static<BlockShapeMNK>::value, "ThreadBlock tile shape must be static");
+    static_assert(mute::rank(BlockShapeMNK{}) == 3, "BlockShapeMNK must be rank 3");
+    static_assert(mute::rank(BlockCoordMNKL{}) == 4, "BlockCoordMNKL must be rank 4");
+
+    // Separate out problem shape for convenience
+    auto M = get<0>(problem_shape_mnkl);
+    auto N = get<1>(problem_shape_mnkl);
+    auto L = get<3>(problem_shape_mnkl);
+
+    auto stride_c    = detail::get_epilogue_stride<EpilogueSchedule>(params.dC);
+    auto stride_d    = detail::get_epilogue_stride<EpilogueSchedule>(params.dD);
+    auto stride_bias = detail::get_epilogue_stride<EpilogueSchedule>(BiasStride{});
+
+    // Represent the full output tensor
+    Tensor mC0_mnl = make_tensor(make_gmem_ptr(params.ptr_C0), make_shape(M,N,L), stride_c);                   // (m,n,l)
+    Tensor mC1_mnl = make_tensor(make_gmem_ptr(params.ptr_C1), make_shape(M,N,L), stride_c);                   // (m,n,l)
+    Tensor mD_mnl = make_tensor(make_gmem_ptr(params.ptr_D), make_shape(M,N,L), stride_d);                     // (m,n,l)
+    Tensor mBias_mnl = make_tensor(make_gmem_ptr(params.ptr_Bias), make_shape(M,N,L), stride_bias);            // (m,n,l)
+
+    Tensor gC0_mnl = local_tile(mC0_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});      // (BLK_M,BLK_N,m,n,l)
+    Tensor gC1_mnl = local_tile(mC1_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});      // (BLK_M,BLK_N,m,n,l)
+
+    Tensor gD_mnl = local_tile(mD_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});        // (BLK_M,BLK_N,m,n,l)
+    Tensor gBias_mnl = local_tile(mBias_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});  // (BLK_M,BLK_N,m,n,l)
+
+    // Slice to get the tile this thread block is responsible for
+    auto [m_coord, n_coord, k_coord, l_coord] = blk_coord_mnkl;
+    Tensor gC0 = gC0_mnl(_,_,m_coord,n_coord,l_coord);                                                   // (BLK_M,BLK_N)
+    Tensor gC1 = gC1_mnl(_,_,m_coord,n_coord,l_coord);                                                   // (BLK_M,BLK_N)
+    Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord);                                                     // (BLK_M,BLK_N)
+    Tensor gBias = gBias_mnl(_,_,m_coord,n_coord,l_coord);                                               // (BLK_M,BLK_N)
+
+    // Partition source and destination tiles to match the accumulator partitioning
+    auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
+    Tensor tCgD = thr_mma.partition_C(gD);                                                           // (VEC,THR_M,THR_N)
+    Tensor tCgC0 = thr_mma.partition_C(gC0);                                                         // (VEC,THR_M,THR_N)
+    Tensor tCgC1 = thr_mma.partition_C(gC1);                                                         // (VEC,THR_M,THR_N)
+    Tensor tCgBias = thr_mma.partition_C(gBias);                                                     // (VEC,THR_M,THR_N)
+
+    static_assert(is_static<FrgLayout>::value,
+        "Accumulator layout must be static");
+    MUTE_STATIC_ASSERT_V(size(tCgC0) == size(tCgD),
+        "Source and destination must have the same number of elements.");
+    MUTE_STATIC_ASSERT_V(size(tCgC1) == size(tCgD),
+        "Source and destination must have the same number of elements.");
+    MUTE_STATIC_ASSERT_V(size(tCgD) == size(accumulators),
+        "Accumulator count must have the same destination element count.");
+    MUTE_STATIC_ASSERT_V(size(tCgBias) == size(accumulators),
+        "Accumulator count must have the same destination element count.");
+
+    auto cD = make_identity_tensor(make_shape(unwrap(shape<0>(gD)), unwrap(shape<1>(gD))));
+    Tensor tCcD = thr_mma.partition_C(cD);
+
+    bool bias_needed = params.ptr_Bias != nullptr;
+    bool c0_needed = (params.ptr_C0 != nullptr) && epilogue_op.is_source0_needed();
+    bool c1_needed = (params.ptr_C1 != nullptr) && epilogue_op.is_source1_needed();
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accumulators); ++i) {
+      if (elem_less(tCcD(i), make_coord(get<0>(residue_mnk), get<1>(residue_mnk)))) {
+        ElementBias bias = bias_needed ? tCgBias(i) : ElementBias(0);
+        ElementC c0 = c0_needed ? tCgC0(i) : ElementC(0);
+        ElementC c1 = c1_needed ? tCgC1(i) : ElementC(0);
+
+        tCgD(i) = epilogue_op(accumulators(i), c0, c1, bias);
+      }
+    }
+  }
+
+private:
+  Params params;
+  ThreadEpilogueOp epilogue_op;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace collective
+} // namespace epilogue
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp22_epilogue_evt.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp22_epilogue_evt.hpp
new file mode 100644
index 000000000..a3d2e43c3
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp22_epilogue_evt.hpp
@@ -0,0 +1,336 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/epilogue/dispatch_policy.hpp"
+#include "mutlass/epilogue/collective/detail.hpp"
+#include "mutlass/epilogue/fusion/callbacks.hpp"
+#include "mutlass/epilogue/fusion/mp22_callbacks.hpp"
+#include "mutlass/detail/layout.hpp"
+
+#include "mute/tensor.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::collective {
+using namespace mute;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  int Stages_,
+  int FragmentSize_,
+  class SmemLayout_,
+  class TiledCopyS2R_,
+  class ElementC_,
+  class StrideC_,
+  class ElementD_,
+  class StrideD_,
+  class ElementAccumulator_,
+  class FusionCallbacks_
+>
+class CollectiveEpilogue <
+  Mp22CollectiveEpilogue<Stages_, FragmentSize_>,
+  SmemLayout_,
+  TiledCopyS2R_,
+  ElementC_,
+  StrideC_,
+  ElementD_,
+  StrideD_,
+  ElementAccumulator_,
+  FusionCallbacks_
+> {
+public:
+  //
+  // Type Aliases
+  //
+  using DispatchPolicy = Mp22CollectiveEpilogue<Stages_, FragmentSize_>;
+  using SmemLayout = SmemLayout_;
+  using TiledCopyS2R = TiledCopyS2R_;
+  using ElementC = ElementC_;
+  using StrideC = StrideC_;
+  using ElementD = ElementD_;
+  using StrideD = StrideD_;
+  using ElementAccumulator = ElementAccumulator_;
+  using FusionCallbacks = FusionCallbacks_;
+
+  using ThreadEpilogueOp = typename epilogue::fusion::FusionCallbacksTraits<FusionCallbacks>::Operation;
+
+  static_assert(mute::rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]");
+  static_assert(mute::rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]");
+
+private:
+  static constexpr int Stages = Stages_;
+  static constexpr int FragmentSize = FragmentSize_;
+
+  static constexpr bool is_source_supported = not mute::is_void_v<ElementC>;
+  static constexpr bool is_destination_supported = not mute::is_void_v<ElementD>;
+  static constexpr int NumBarriers = 0;
+
+public:
+  struct SharedStorage {
+    mute::array_aligned<ElementAccumulator, mute::cosize_v<SmemLayout>> smem_epi;
+
+    using FusionStorage = typename FusionCallbacks::SharedStorage;
+    FusionStorage thread;
+  };
+
+  struct Arguments {
+    typename FusionCallbacks::Arguments thread{};
+    ElementC const* ptr_C;
+    StrideC dC;
+    ElementD* ptr_D;
+    StrideD dD;
+  };
+
+  using Params = Arguments;
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(
+      [[maybe_unused]] ProblemShape const& _,
+      Arguments const& args,
+      [[maybe_unused]] void* workspace) {
+    return args;
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return FusionCallbacks::get_workspace_size(problem_shape, args.thread);
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream) {
+    return FusionCallbacks::initialize_workspace(problem_shape, args.thread, workspace, stream);
+  }
+
+  MUTLASS_HOST_DEVICE
+  CollectiveEpilogue(Params const& params_, SharedStorage& shared_storage)
+    : params(params_), fusion_callbacks(params_.thread, shared_storage.thread) {}
+
+
+  template <
+    class ProblemShapeMNKL,
+    class BlockShapeMNK,
+    class BlockCoordMNKL,
+    class FrgEngine, class FrgLayout,
+    class TiledMma,
+    class ResidueMNK
+  >
+  MUTLASS_HOST_DEVICE void
+  operator() (
+      ProblemShapeMNKL problem_shape_mnkl,
+      BlockShapeMNK blk_shape_MNK,
+      BlockCoordMNKL blk_coord_mnkl,
+      mute::Tensor<FrgEngine, FrgLayout> const& accumulators,
+      TiledMma tiled_mma,
+      ResidueMNK residue_mnk,
+      int thread_idx,
+      char* smem_buf)
+  {
+    using namespace mute;
+    using X = Underscore;
+    using ElementCompute_ = typename epilogue::fusion::FusionCallbacksTraits<FusionCallbacks>::ElementCompute;
+    using ElementCompute  = mute::conditional_t<mute::is_void_v<ElementCompute_>, ElementAccumulator, ElementCompute_>;
+
+    static_assert(mute::rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4");
+    static_assert(is_static<BlockShapeMNK>::value, "ThreadBlock tile shape must be static");
+    static_assert(mute::rank(BlockShapeMNK{}) == 3, "BlockShapeMNK must be rank 3");
+    static_assert(mute::rank(BlockCoordMNKL{}) == 4, "BlockCoordMNKL must be rank 4");
+    static_assert(rank(FrgLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA,MMA_M,MMA_N)");
+    static_assert(is_rmem<FrgEngine>::value, "Accumulator must be RF resident.");
+    static_assert(is_same_v<typename FrgEngine::value_type, ElementAccumulator>, "Accumulator type doesn't match");
+
+
+    // Separate out problem shape for convenience
+    auto M = get<0>(problem_shape_mnkl);
+    auto N = get<1>(problem_shape_mnkl);
+    auto L = get<3>(problem_shape_mnkl);
+
+    auto residue_mn = make_tuple(get<0>(residue_mnk), get<1>(residue_mnk));
+
+    // Represent the full output tensor
+    Tensor mC_mnl = make_tensor(make_gmem_ptr(params.ptr_C), make_shape(M,N,L), params.dC);      //             (m,n,l)
+    Tensor mD_mnl = make_tensor(make_gmem_ptr(params.ptr_D), make_shape(M,N,L), params.dD);      //             (m,n,l)
+    Tensor gC_mnl = local_tile(mC_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});      // (BLK_M,BLK_N,m,n,l)
+    Tensor gD_mnl = local_tile(mD_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});      // (BLK_M,BLK_N,m,n,l)
+
+    // Slice to get the tile this CTA is responsible for
+    auto [m_coord, n_coord, k_coord, l_coord] = blk_coord_mnkl;
+    Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord);                                                   // (BLK_M,BLK_N)
+    Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord);                                                   // (BLK_M,BLK_N)
+
+
+    auto mma_tile_m = tile_size<0>(tiled_mma);
+    auto mma_tile_n = tile_size<1>(tiled_mma);
+    auto epi_tile_m = size<0>(SmemLayout{});
+    auto epi_tile_n = size<1>(SmemLayout{});
+
+    static_assert(epi_tile_m % mma_tile_m == 0, "MMA_TILE_M must divide EPI_TILE_M");
+    static_assert(epi_tile_n % mma_tile_n == 0, "MMA_TILE_N must divide EPI_TILE_N");
+    static_assert(typename TiledCopyS2R::TiledNumThr{} == size(tiled_mma));
+
+    // Construct a tensor in SMEM that we can partition for rearranging data
+    SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+    Tensor sC = make_tensor(make_smem_ptr(storage.smem_epi.data()), SmemLayout{});
+
+    using CopyAtomR2S = Copy_Atom<DefaultCopy, ElementAccumulator>;
+    auto tiled_copy_r2s = make_tiled_copy_C(CopyAtomR2S{}, tiled_mma);
+    auto thread_r2s = tiled_copy_r2s.get_thread_slice(thread_idx);
+    Tensor tCaC = thread_r2s.retile_S(accumulators);                                          // ((Atom,AtomNum), MMA_M, MMA_N)
+    Tensor tCsC = thread_r2s.partition_D(sC);                                                 // ((Atom,AtomNum),PIPE_M,PIPE_N)
+
+    // Tile gD and gC by the shape of SmemLayout first
+    auto tile  = make_shape(size<0>(sC), size<1>(sC));
+    Tensor gCt = flat_divide(gC, tile);                                                // (SMEM_M,SMEM_N,TILE_M,TILE_N)
+    Tensor gDt = flat_divide(gD, tile);                                                // (SMEM_M,SMEM_N,TILE_M,TILE_N)
+
+    auto tiled_copy_s2r = TiledCopyS2R{};
+    auto thread_s2r = tiled_copy_s2r.get_thread_slice(thread_idx);
+    Tensor tDsC = thread_s2r.partition_S(sC);                                   //               ((Atom,AtomNum),ATOM_M,ATOM_N)
+    Tensor tDgC = thread_s2r.partition_D(gCt);                                  // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N)
+    Tensor tDgD = thread_s2r.partition_D(gDt);                                  // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N)
+
+    // Allocate intermediate registers on the dst tensors
+    Tensor tDrC = make_tensor<ElementAccumulator>(take<0,3>(shape(tDgC)));       // ((Atom,AtomNum),ATOM_M,ATOM_N)
+    Tensor tDrD = make_tensor<ElementD>(shape(tDrC));                            // ((Atom,AtomNum),ATOM_M,ATOM_N)
+
+    // Allocate intermediate registers on the src tensors
+    Tensor tCrC = make_tensor<ElementC>(shape(tDrC));                            // ((Atom,AtomNum),ATOM_M,ATOM_N)
+
+    // Construct proper tiled copy
+    constexpr int vec_elem = decltype(max_common_vector(tDgC, tDgD))::value;
+    constexpr int vec_bits = vec_elem * sizeof_bits_v<ElementC>;
+    using CopyAtomG2R = Copy_Atom<
+            UniversalCopy<uint_bit_t<mute::min(1024, vec_bits)>>, ElementC>;
+
+    using TiledCopyG2R = TiledCopy<CopyAtomG2R,
+                              typename TiledCopyS2R::TiledLayout_TV,
+                              typename TiledCopyS2R::Tiler_MN>;
+    TiledCopyG2R tiled_copy_g2r{};
+    auto thr_copy_g2r = tiled_copy_g2r.get_slice(thread_idx);
+    Tensor tCgC       = thr_copy_g2r.partition_S(gCt);
+    Tensor tCrC_view  = thr_copy_g2r.retile_D(tCrC);
+
+    // Repeat the D-partitioning for coordinates and predication
+    Tensor cD   = make_identity_tensor(make_shape(size<0>(gD),size<1>(gD)));          // (BLK_M,BLK_N) -> (blk_m,blk_n)
+    Tensor cDt  = flat_divide(cD, tile);                                //                (SMEM_M,SMEM_N,TILE_M,TILE_N)
+    Tensor tDcD = thread_s2r.partition_D(cDt);                                  // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N)
+
+    bool is_C_load_needed = is_source_supported and fusion_callbacks.is_C_load_needed();
+
+    auto args = mutlass::epilogue::fusion::detail::VisitorArgs {
+                  problem_shape_mnkl,
+                  blk_shape_MNK,
+                  blk_coord_mnkl,
+                  residue_mn,
+                  tile,
+                  tiled_copy_s2r,
+                  thread_idx,
+                  cD,
+                  tDcD,
+                  tCrC,
+                  tDgD
+                };
+
+    auto callbacks = fusion_callbacks.get_callbacks(args);
+
+    callbacks.begin_epilogue();
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int epi_m = 0; epi_m < size<2>(cDt); ++epi_m) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int epi_n = 0; epi_n < size<3>(cDt); ++epi_n) {
+        int step_idx = epi_m * size<3>(cDt) + epi_n;
+        callbacks.begin_step(epi_m, epi_n, step_idx);
+
+        MUTLASS_PRAGMA_UNROLL
+        for (int pipe_m = 0; pipe_m < size<1>(tCsC); ++pipe_m) {
+          MUTLASS_PRAGMA_UNROLL
+          for (int pipe_n = 0; pipe_n < size<2>(tCsC); ++pipe_n) {
+            int mma_m = epi_m * size<1>(tCsC) + pipe_m;
+            int mma_n = epi_n * size<2>(tCsC) + pipe_n;
+            copy(tiled_copy_r2s, tCaC(_, mma_m, mma_n), tCsC(_, pipe_m, pipe_n));
+          } // for pipe_n
+        } // for pipe_m
+
+        __syncthreads();
+
+        copy(tiled_copy_s2r, tDsC, tDrC);
+
+        __syncthreads();
+
+        Tensor tDgDmn = tDgD(_,_,_,epi_m,epi_n);
+        Tensor tDcDmn = tDcD(_,_,_,epi_m,epi_n);
+
+        if (is_C_load_needed) {
+          tCrC_view = tCgC(_, _, _, epi_m, epi_n);
+        }
+
+        Tensor tSR_rAcc_frg = recast<Array<ElementAccumulator, FragmentSize>>(tDrC);
+        Tensor tSR_rD_frg   = recast<Array<ElementD          , FragmentSize>>(tDrD);
+
+        MUTLASS_PRAGMA_UNROLL
+        for (int epi_v = 0; epi_v < size(tSR_rAcc_frg); ++epi_v) {
+          tSR_rD_frg(epi_v) = callbacks.visit(tSR_rAcc_frg(epi_v), epi_v, epi_m, epi_n);
+        }
+
+        callbacks.end_step(epi_m, epi_n, step_idx);
+
+        if constexpr (is_destination_supported) {
+          MUTLASS_PRAGMA_UNROLL
+          for (int m = 0; m < size<1>(tDgDmn); ++m) {
+            MUTLASS_PRAGMA_UNROLL
+            for (int n = 0; n < size<2>(tDgDmn); ++n) {
+              if (elem_less(tDcDmn(0, m, n), residue_mn)) {
+                copy(tDrD(_, m, n), tDgDmn(_, m, n));
+              }
+            } // for n
+          } // for m
+        }
+      } // for epi_n
+    } // for epi_m
+    callbacks.end_epilogue();
+  }
+private:
+  Params const& params;
+  FusionCallbacks fusion_callbacks;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp22_epilogue_vectorized.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp22_epilogue_vectorized.hpp
new file mode 100644
index 000000000..ea5a4468c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp22_epilogue_vectorized.hpp
@@ -0,0 +1,597 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Functor performing elementwise operations used by epilogues.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+
+#include "mute/tensor.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class StrideC,
+  class StrideD,
+  class ThreadEpilogueOp,
+  class SmemLayout,
+  class CopyAtomR2S,
+  class TiledCopyS2R,
+  class CopyAtomR2G,
+  class EpilogueScheduleType = EpilogueSimtVectorized,
+  class Enable = void
+>
+class Epilogue {
+  static_assert(mute::is_same_v<EpilogueScheduleType, EpilogueSimtVectorized>,
+                "Could not find an epilogue specialization.");
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// Epilogue Vectorized
+/// Applies an element wise operation to all elements within the fragment
+/// and writes it out to destination storage.
+///
+/// Ways to generalize this:
+/// - CTA tile shape
+/// - vectorization requirements (GMEM)
+/// - vectoriz(able) transform()
+///
+template <
+  class StrideC_,
+  class StrideD_,
+  class ThreadEpilogueOp_,
+  class SmemLayout_,
+  class CopyAtomR2S_,
+  class TiledCopyS2R_,
+  class CopyAtomR2G_,
+  class EpilogueScheduleType_
+>
+class Epilogue<
+        StrideC_,
+        StrideD_,
+        ThreadEpilogueOp_,
+        SmemLayout_,
+        CopyAtomR2S_,
+        TiledCopyS2R_,
+        CopyAtomR2G_,
+        EpilogueScheduleType_,
+        mute::enable_if_t<
+          mute::is_same_v<EpilogueScheduleType_, EpilogueSimtVectorized>
+        >
+      > {
+public:
+  //
+  // Type Aliases
+  //
+  // derived types of output thread level operator
+  using ThreadEpilogueOp = ThreadEpilogueOp_;
+  using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator;
+  using ElementCompute = typename ThreadEpilogueOp::ElementCompute;
+  using ElementScalar = ElementCompute;
+  using ElementOutput = typename ThreadEpilogueOp::ElementOutput;
+  using ElementC = typename ThreadEpilogueOp::ElementC;
+  using StrideC = StrideC_;
+  using ElementD = typename ThreadEpilogueOp::ElementD;
+  using StrideD = StrideD_;
+  using DispatchPolicy = void;
+
+  using ElementBias = typename detail::IsThreadEpilogueOpWithBias<ThreadEpilogueOp>::type;
+  using SmemLayout   = SmemLayout_;
+  using CopyAtomR2S  = CopyAtomR2S_;
+  using TiledCopyS2R = TiledCopyS2R_;
+  using CopyAtomR2G  = CopyAtomR2G_;
+
+  static const int kOutputAlignment = ThreadEpilogueOp::kCount;
+  using GmemTiledCopyC = void;
+  using GmemTiledCopyD = CopyAtomR2G;
+  using AlignmentType = typename mute::uint_bit<sizeof_bits<ElementOutput>::value * kOutputAlignment>::type;
+  static constexpr int NumBarriers = 0;
+  static constexpr bool IsEpilogueBiasSupported = detail::IsThreadEpilogueOpWithBias<ThreadEpilogueOp>::value;
+  using StrideBias = mute::conditional_t<detail::is_m_major<StrideD>(), Stride<_1,_0,int64_t>, Stride<_0,_1,int64_t>>;
+
+  static_assert(rank(StrideC{}) == 3, "StrideCD must be rank-3: [M, N, L]");
+  static_assert(rank(StrideD{}) == 3, "StrideCD must be rank-3: [M, N, L]");
+
+  struct SharedStorage
+  {
+    mute::array_aligned<ElementAccumulator, mute::cosize_v<SmemLayout>> smem_epilogue;
+  };
+
+  static constexpr bool IsActHasArgs = detail::IsThreadEpilogueOpWithElementwiseArguments<ThreadEpilogueOp>::value;
+
+  // Host side epilogue arguments
+  template<class ThreadEpiOp, class = void>
+  struct ThreadEpilogueOpArguments {
+    ElementScalar alpha{0};
+    ElementScalar beta{0};
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+    ElementBias const* bias_ptr = nullptr;
+    StrideBias dBias{};
+  };
+
+  template<class ThreadEpiOp>
+  struct ThreadEpilogueOpArguments<
+          ThreadEpiOp,
+          mute::enable_if_t<detail::IsThreadEpilogueOpWithElementwiseArguments<ThreadEpiOp>::value>> {
+    ElementScalar alpha{0};
+    ElementScalar beta{0};
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+    ElementBias const* bias_ptr = nullptr;
+    StrideBias dBias{};
+    typename ThreadEpiOp::ElementwiseArguments activation{};
+  };
+
+  struct Arguments {
+    ThreadEpilogueOpArguments<ThreadEpilogueOp> thread{};
+    using StrideBias = decltype(thread.dBias);
+    ElementC const* ptr_C = nullptr;
+    StrideC dC{};
+    ElementD* ptr_D = nullptr;
+    StrideD dD{};
+  };
+
+  // Device side epilogue params
+  // using Params = Arguments;
+  template<class ThreadEpiOp, class = void>
+  struct ParamsType {
+    typename ThreadEpiOp::Params thread{};
+    ElementC const* ptr_C = nullptr;
+    StrideC dC{};
+    ElementD* ptr_D = nullptr;
+    StrideD dD{};
+    ElementBias const* ptr_Bias = nullptr;
+    StrideBias dBias{};
+  };
+
+  template<class ThreadEpiOp>
+  struct ParamsType<
+          ThreadEpiOp,
+          mute::enable_if_t<detail::IsThreadEpilogueOpWithElementwiseArguments<ThreadEpiOp>::value>> {
+    typename ThreadEpiOp::Params thread{};
+    typename ThreadEpiOp::ElementwiseArguments activation{};
+    ElementC const* ptr_C = nullptr;
+    StrideC dC{};
+    ElementD* ptr_D = nullptr;
+    StrideD dD{};
+    ElementBias const* ptr_Bias = nullptr;
+    StrideBias dBias{};
+  };
+
+  using Params = ParamsType<ThreadEpilogueOp>;
+
+  //
+  // Methods
+  //
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(
+      [[maybe_unused]] ProblemShape const& _,
+      Arguments const& args,
+      [[maybe_unused]] void* workspace) {
+    typename ThreadEpilogueOp::Params thread_op_args;
+    thread_op_args.alpha = args.thread.alpha;
+    thread_op_args.beta = args.thread.beta;
+    thread_op_args.alpha_ptr = args.thread.alpha_ptr;
+    thread_op_args.beta_ptr = args.thread.beta_ptr;
+
+    if constexpr (IsActHasArgs) {
+      return {
+        thread_op_args,
+        args.thread.activation,
+        args.ptr_C,
+        args.dC,
+        args.ptr_D,
+        args.dD,
+        args.thread.bias_ptr,
+        args.thread.dBias
+      };
+    }
+    else {
+      return {
+        thread_op_args,
+        args.ptr_C,
+        args.dC,
+        args.ptr_D,
+        args.dD,
+        args.thread.bias_ptr,
+        args.thread.dBias
+      };
+    }
+  }
+
+  template<class ProblemShape>
+  MUTLASS_HOST_DEVICE static bool
+  can_implement(
+      [[maybe_unused]] ProblemShape const& problem_shape,
+      [[maybe_unused]] Arguments const& args) {
+    return true;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Epilogue(Params const& params_)
+      : params(params_), epilogue_op(params_.thread) { }
+
+  MUTLASS_DEVICE
+  bool
+  is_source_needed() {
+    return epilogue_op.is_source_needed();
+  }
+
+  template<
+    class ProblemShapeMNKL,
+    class BlockShapeMNK,
+    class BlockCoordMNKL,
+    class FrgEngine, class FrgLayout,
+    class TiledMma,
+    class ResidueMNK
+  >
+  MUTLASS_DEVICE void
+  operator()(
+      ProblemShapeMNKL problem_shape_mnkl,
+      BlockShapeMNK blk_shape_MNK,
+      BlockCoordMNKL blk_coord_mnkl,
+      mute::Tensor<FrgEngine,FrgLayout> const& accumulators,                   // (MMA,MMA_M,MMA_N)
+      TiledMma tiled_mma,
+      ResidueMNK residue_mnk,
+      int thread_idx,
+      char* smem_buf)
+  {
+    using namespace mute;
+    using X = Underscore;
+
+    static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4");
+    static_assert(is_static<BlockShapeMNK>::value, "ThreadBlock tile shape must be static");
+    static_assert(rank(BlockShapeMNK{}) == 3, "BlockShapeMNK must be rank 3");
+    static_assert(rank(BlockCoordMNKL{}) == 4, "BlockCoordMNKL must be rank 3");
+
+    // synchronizing function for smem reads/writes
+    auto synchronize = [] () { __syncthreads(); };
+
+    // Separate out problem shape for convenience
+    auto M = get<0>(problem_shape_mnkl);
+    auto N = get<1>(problem_shape_mnkl);
+    auto L = get<3>(problem_shape_mnkl);
+
+    constexpr int major_index = detail::is_m_major<StrideD>() ? 0 : 1;
+
+    // Represent the full output tensor
+    Tensor mC_mnl = make_tensor(make_gmem_ptr(params.ptr_C), make_shape(M,N,L), params.dC);             //             (m,n,l)
+    Tensor mD_mnl = make_tensor(make_gmem_ptr(params.ptr_D), make_shape(M,N,L), params.dD);             //             (m,n,l)
+    Tensor mBias_mnl = make_tensor(make_gmem_ptr(params.ptr_Bias), make_shape(M,N,L), params.dBias);    //             (m,n,l)
+
+    Tensor gC_mnl = local_tile(mC_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});             // (BLK_M,BLK_N,m,n,l)
+    Tensor gD_mnl = local_tile(mD_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});             // (BLK_M,BLK_N,m,n,l)
+    Tensor gBias_mnl = local_tile(mBias_mnl, blk_shape_MNK, make_coord(_,_,_), Step<_1,_1, X>{});       // (BLK_M,BLK_N,m,n,l)
+
+    // Slice to get the tile this CTA is responsible for
+    auto [m_coord, n_coord, k_coord, l_coord] = blk_coord_mnkl;
+    Tensor gC = gC_mnl(_,_,m_coord,n_coord,l_coord);                                                   // (BLK_M,BLK_N)
+    Tensor gD = gD_mnl(_,_,m_coord,n_coord,l_coord);                                                   // (BLK_M,BLK_N)
+    Tensor gBias = gBias_mnl(_,_,m_coord,n_coord,l_coord);                                             // (BLK_M,BLK_N)
+
+    // Construct a tensor in SMEM that we can partition for rearranging data
+    SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+    // Tensor sC = make_tensor(make_smem_ptr(storage.smem_epilogue.data()), SmemLayout{});              // (SMEM_M,SMEM_N)
+    Tensor sAcc = make_tensor(make_smem_ptr(storage.smem_epilogue.data()), SmemLayout{});            // (SMEM_M,SMEM_N)
+
+    // Partition sAcc to match the accumulator partitioning
+    auto tiled_r2s = make_tiled_copy_C(CopyAtomR2S{}, tiled_mma);
+    auto thread_r2s     = tiled_r2s.get_thread_slice(thread_idx);
+    Tensor tRS_rAcc = thread_r2s.retile_S(accumulators);                              // ((Atom,AtomNum), MMA_M, MMA_N)
+    Tensor tRS_sAcc = thread_r2s.partition_D(sAcc);                                   // ((Atom,AtomNum),PIPE_M,PIPE_N)
+
+    // Tile gD and gC by the shape of SmemLayout first
+    auto tile  = make_shape(size<0>(sAcc), size<1>(sAcc));
+    Tensor gCt = flat_divide(gC, tile);                                              // (SMEM_M,SMEM_N,TILE_M,TILE_N)
+    Tensor gDt = flat_divide(gD, tile);                                              // (SMEM_M,SMEM_N,TILE_M,TILE_N)
+    Tensor gBiast = flat_divide(gBias, tile);                                        // (SMEM_M,SMEM_N,TILE_M,TILE_N)
+
+    // Partition sAcc, gC, and gD for the output
+    auto tiled_s2r = TiledCopyS2R{};
+    auto thread_s2r     = tiled_s2r.get_thread_slice(thread_idx);
+    Tensor tSR_sAcc = thread_s2r.partition_S(sAcc);                      //               ((Atom,AtomNum),ATOM_M,ATOM_N)
+    Tensor tSR_gC = thread_s2r.partition_D(gCt);                         // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N)
+    Tensor tSR_gD = thread_s2r.partition_D(gDt);                         // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N)
+    Tensor tSR_gBias = thread_s2r.partition_D(gBiast);                   // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N)
+
+    // Allocate intermediate registers on the dst tensors
+    Tensor tSR_rAcc = make_tensor<ElementAccumulator>(take<0,3>(shape(tSR_gC)));       // ((Atom,AtomNum),ATOM_M,ATOM_N)
+    Tensor tSR_rC = make_tensor<ElementC>(shape(tSR_rAcc));                            // ((Atom,AtomNum),ATOM_M,ATOM_N)
+    Tensor tSR_rD = make_tensor<ElementD>(shape(tSR_rAcc));                            // ((Atom,AtomNum),ATOM_M,ATOM_N)
+    Tensor tSR_rBias = make_tensor_like(tSR_gBias);                      // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N)
+
+    // Repeat the D-partitioning for coordinates and predication
+    Tensor cD   = make_identity_tensor(make_shape(size<0>(gD),size<1>(gD)));          // (BLK_M,BLK_N) -> (blk_m,blk_n)
+    Tensor cDt  = flat_divide(cD, tile);                                 //                (SMEM_M,SMEM_N,TILE_M,TILE_N)
+    Tensor tSR_cD = thread_s2r.partition_D(cDt);                         // ((Atom,AtomNum),ATOM_M,ATOM_N,TILE_M,TILE_N)
+
+
+    MUTE_STATIC_ASSERT(size<1>(tRS_rAcc) % size<3>(tSR_gC) == 0);  // TILE_M divides MMA_M
+    MUTE_STATIC_ASSERT(size<2>(tRS_rAcc) % size<4>(tSR_gC) == 0);  // TILE_N divides MMA_N
+
+#if 0
+    if (thread_idx == 0 && m_coord == 0 && n_coord == 0) {
+      print("aC   : "); print(accumulators.layout()); print("\n");
+      print("gC   : "); print(gC.layout()); print("\n");
+      print("gD   : "); print(gD.layout()); print("\n");
+      print("gBias   : "); print(gBias.layout()); print("\n");
+      print("sAcc   : "); print(sAcc.layout()); print("\n");
+      print("\n");
+      print("tRS_sAcc : "); print(tRS_sAcc.layout()); print("\n");
+      print("tRS_rAcc : "); print(tRS_rAcc.layout()); print("\n");
+      print("\n");
+      print("gDt  : "); print(gDt.layout()); print("\n");
+      print("tDsC : "); print(tDsC.layout()); print("\n");
+      print("tDrC : "); print(tDrC.layout()); print("\n");
+      print("\n");
+      print("tSR_rC : "); print(tSR_rC.layout()); print("\n");
+      print("tSR_rD : "); print(tSR_rD.layout()); print("\n");
+      print("tSR_gC : "); print(tSR_gC.layout()); print("\n");
+      print("tSR_gD : "); print(tSR_gD.layout()); print("\n");
+      print("\n");
+      print("gBiast  : "); print(gBiast.layout()); print("\n");
+      print("tSR_gBias  : "); print(tSR_gBias.layout()); print("\n");
+      print("tSR_rBias  : "); print(tSR_rBias.layout()); print("\n");
+      print("\n");
+    }
+#endif
+
+    if constexpr (IsEpilogueBiasSupported) {
+      if (params.ptr_Bias) {
+        // Filter so we don't issue redundant copies over stride-0 modes
+        // (only works if 0-strides are in same location, which is by construction)
+        Tensor tSR_gBias_flt = filter_zeros(tSR_gBias);
+        Tensor tSR_rBias_flt = filter_zeros(tSR_rBias);
+        Tensor tSR_cD_flt = filter_zeros(tSR_cD, tSR_gBias.stride());
+
+        // Step 0. Copy Bias from GMEM to fragment
+        auto pred_fn = [&] (auto const&... coords) { return elem_less(tSR_cD_flt(coords...), take<0, 2>(residue_mnk)); };
+        copy_if(pred_fn, tSR_gBias_flt, tSR_rBias_flt);
+      }
+    }
+
+    // For each tiling needed for SmemLayout to cover shape(gD)
+    MUTLASS_PRAGMA_UNROLL
+    for (int step_m = 0; step_m < size<2>(cDt); ++step_m) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int step_n = 0; step_n < size<3>(cDt); ++step_n) {
+        // Step 1. Copy to SMEM
+        MUTLASS_PRAGMA_UNROLL
+        for (int pipe_m = 0; pipe_m < size<1>(tRS_sAcc); ++pipe_m) {
+          MUTLASS_PRAGMA_UNROLL
+          for (int pipe_n = 0; pipe_n < size<2>(tRS_sAcc); ++pipe_n) {
+            int mma_m = step_m * size<1>(tRS_sAcc) + pipe_m;
+            int mma_n = step_n * size<2>(tRS_sAcc) + pipe_n;
+            copy(tiled_r2s, tRS_rAcc(_,mma_m,mma_n), tRS_sAcc(_,pipe_m,pipe_n));
+          }
+        }
+
+        // Step 2. Wait for SMEM writes to complete
+        synchronize();
+
+        // Step 3. Copy from SMEM into a fragment
+        copy(tiled_s2r, tSR_sAcc, tSR_rAcc);
+
+        // Step 4. Wait for SMEM reads to complete
+        synchronize();
+
+        Tensor tSR_gDmn = tSR_gD(_,_,_,step_m,step_n);
+        Tensor tSR_cDmn = tSR_cD(_,_,_,step_m,step_n);
+
+        if constexpr (IsEpilogueBiasSupported) {
+          Tensor tSR_rBiasmn = tSR_rBias(_,_,_,step_m,step_n);
+          if (epilogue_op.is_source_needed()) {
+            // source is needed
+            Tensor tSR_gCmn  = tSR_gC(_,_,_,step_m,step_n);
+            if (get<0>(residue_mnk) >= get<0>(BlockShapeMNK{}) &&
+                get<1>(residue_mnk) >= get<1>(BlockShapeMNK{})) {
+              MUTLASS_PRAGMA_UNROLL
+              for (int m = 0; m < size<1>(tSR_gDmn); ++m)
+              {
+                MUTLASS_PRAGMA_UNROLL
+                for (int n = 0; n < size<2>(tSR_gDmn); ++n)
+                {
+                  MUTLASS_PRAGMA_UNROLL
+                  for (int i = 0; i < size<0>(tSR_rAcc); ++i) {
+                    tSR_rC(i,m,n) = tSR_gCmn(i,m,n);
+                  }
+                }
+              }
+            } else {
+              for (int m = 0; m < size<1>(tSR_gDmn); ++m)
+              {
+                MUTLASS_PRAGMA_UNROLL
+                for (int n = 0; n < size<2>(tSR_gDmn); ++n)
+                {
+                  if (elem_less(tSR_cDmn(0,m,n), take<0,2>(residue_mnk))) {
+                    MUTLASS_PRAGMA_UNROLL
+                    for (int i = 0; i < size<0>(tSR_rAcc); ++i) {
+                      tSR_rC(i,m,n) = tSR_gCmn(i,m,n);
+                    }
+                  }
+                }
+              }
+            }
+            // Step 6. Elementwise operation with conversion
+            MUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < size(tSR_rAcc); ++i) {
+              if constexpr (IsActHasArgs) {
+                epilogue_op(tSR_rD(i), tSR_rD(i), tSR_rAcc(i), tSR_rC(i), tSR_rBiasmn(i), params.activation);
+              } else {
+                epilogue_op(tSR_rD(i), tSR_rD(i), tSR_rAcc(i), tSR_rC(i), tSR_rBiasmn(i));
+              }
+            }
+          }
+          else {
+            // source is not needed, avoid load and lift compute
+            MUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < size(tSR_rAcc); ++i) {
+              if constexpr (IsActHasArgs) {
+                epilogue_op(tSR_rD(i), tSR_rD(i), tSR_rAcc(i), tSR_rBiasmn(i), params.activation);
+              } else {
+                epilogue_op(tSR_rD(i), tSR_rD(i), tSR_rAcc(i), tSR_rBiasmn(i));
+              }
+            }
+          }
+          if (get<0>(residue_mnk) >= get<0>(BlockShapeMNK{}) &&
+              get<1>(residue_mnk) >= get<1>(BlockShapeMNK{})) {
+            MUTLASS_PRAGMA_UNROLL
+            for (int m = 0; m < size<1>(tSR_gDmn); ++m)
+            {
+              MUTLASS_PRAGMA_UNROLL
+              for (int n = 0; n < size<2>(tSR_gDmn); ++n)
+              {
+                copy(CopyAtomR2G{}, tSR_rD(_,m,n), tSR_gDmn(_,m,n));
+              }
+            }
+          } else {
+            MUTLASS_PRAGMA_UNROLL
+            for (int m = 0; m < size<1>(tSR_gDmn); ++m)
+            {
+              MUTLASS_PRAGMA_UNROLL
+              for (int n = 0; n < size<2>(tSR_gDmn); ++n)
+              {
+                if (elem_less(tSR_cDmn(0,m,n), take<0,2>(residue_mnk))) {
+                  Tensor tDpD = make_tensor<bool>(make_shape(size(tSR_rD(_,m,n))));
+                  MUTLASS_PRAGMA_UNROLL
+                  for (int i = 0; i < size<0>(tDpD); i++) {
+                    tDpD(i) = get<major_index>(tSR_cDmn(0,m,n)) + i < get<major_index>(residue_mnk);
+                  }
+                  copy_if(tDpD, tSR_rD(_,m,n), tSR_gDmn(_,m,n));
+                }
+              }
+            }
+          }
+        } else {
+          if (epilogue_op.is_source_needed()) {
+            // source is needed
+            Tensor tSR_gCmn = tSR_gC(_,_,_,step_m,step_n);
+            if (get<0>(residue_mnk) >= get<0>(BlockShapeMNK{}) &&
+                get<1>(residue_mnk) >= get<1>(BlockShapeMNK{})) {
+              MUTLASS_PRAGMA_UNROLL
+              for (int m = 0; m < size<1>(tSR_gDmn); ++m)
+              {
+                MUTLASS_PRAGMA_UNROLL
+                for (int n = 0; n < size<2>(tSR_gDmn); ++n)
+                {
+                  MUTLASS_PRAGMA_UNROLL
+                  for (int i = 0; i < size<0>(tSR_rAcc); ++i) {
+                    tSR_rC(i,m,n) = tSR_gCmn(i,m,n);
+                  }
+                }
+              }
+            } else {
+              for (int m = 0; m < size<1>(tSR_gDmn); ++m)
+              {
+                MUTLASS_PRAGMA_UNROLL
+                for (int n = 0; n < size<2>(tSR_gDmn); ++n)
+                {
+                  if (elem_less(tSR_cDmn(0,m,n), take<0,2>(residue_mnk))) {
+                    MUTLASS_PRAGMA_UNROLL
+                    for (int i = 0; i < size<0>(tSR_rAcc); ++i) {
+                      tSR_rC(i,m,n) = tSR_gCmn(i,m,n);
+                    }
+                  }
+                }
+              }
+            }
+            // Step 6. Elementwise operation with conversion
+            MUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < size(tSR_rAcc); ++i) {
+              tSR_rD(i) = epilogue_op(tSR_rAcc(i), tSR_rC(i));
+            }
+          }
+          else {
+            // source is not needed, avoid load and lift compute
+            MUTLASS_PRAGMA_UNROLL
+            for (int i = 0; i < size(tSR_rAcc); ++i) {
+              tSR_rD(i) = epilogue_op(tSR_rAcc(i));
+            }
+          }
+          if (get<0>(residue_mnk) >= get<0>(BlockShapeMNK{}) &&
+              get<1>(residue_mnk) >= get<1>(BlockShapeMNK{})) {
+            MUTLASS_PRAGMA_UNROLL
+            for (int m = 0; m < size<1>(tSR_gDmn); ++m)
+            {
+              MUTLASS_PRAGMA_UNROLL
+              for (int n = 0; n < size<2>(tSR_gDmn); ++n)
+              {
+                copy(CopyAtomR2G{}, tSR_rD(_,m,n), tSR_gDmn(_,m,n));
+              }
+            }
+          } else {
+            MUTLASS_PRAGMA_UNROLL
+            for (int m = 0; m < size<1>(tSR_gDmn); ++m)
+            {
+              MUTLASS_PRAGMA_UNROLL
+              for (int n = 0; n < size<2>(tSR_gDmn); ++n)
+              {
+                if (elem_less(tSR_cDmn(0,m,n), take<0,2>(residue_mnk))) {
+                  Tensor tDpD = make_tensor<bool>(make_shape(size(tSR_rD(_,m,n))));
+                  MUTLASS_PRAGMA_UNROLL
+                  for (int i = 0; i < size<0>(tDpD); i++) {
+                    tDpD(i) = get<major_index>(tSR_cDmn(0,m,n)) + i < get<major_index>(residue_mnk);
+                  }
+                  copy_if(tDpD, tSR_rD(_,m,n), tSR_gDmn(_,m,n));
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
+private:
+  Params params;
+  ThreadEpilogueOp epilogue_op;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace collective
+} // namespace epilogue
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp31_epilogue_tme.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp31_epilogue_tme.hpp
new file mode 100644
index 000000000..88c9d595c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/collective/mp31_epilogue_tme.hpp
@@ -0,0 +1,508 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mute/tensor.hpp"
+#include "mutlass/epilogue/collective/builders/mp31_builder_common.inl"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  int StagesC_,
+  int StagesD_,
+  int FragmentSize_,
+  bool ReuseSmemC_,
+  bool DelayTmeStore_,
+  class ThreadEpilogueOp_,
+  class CtaTileMNK_,
+  class EpilogueTile_,
+  class ElementC_,
+  class StrideC_,
+  class ElementD_,
+  class StrideD_,
+  class CopyOpG2S,
+  class CopyOpS2G,
+  class SmemLayoutAtomC_,
+  class SmemLayoutAtomD_,
+  class CopyOpS2R_,
+  class CopyOpR2S_
+>
+class CollectiveEpilogue<
+    Mp31CollectiveEpilogue<StagesC_, StagesD_, FragmentSize_, ReuseSmemC_, DelayTmeStore_>,
+    ThreadEpilogueOp_,
+    CtaTileMNK_,
+    EpilogueTile_,
+    ElementC_,
+    StrideC_,
+    ElementD_,
+    StrideD_,
+    CopyOpG2S,
+    CopyOpS2G,
+    SmemLayoutAtomC_,
+    SmemLayoutAtomD_,
+    CopyOpS2R_,
+    CopyOpR2S_
+> {
+public:
+  //
+  // Type Aliases
+  //
+  using DispatchPolicy = Mp31CollectiveEpilogue<StagesC_, StagesD_, FragmentSize_,ReuseSmemC_, DelayTmeStore_>;
+  using CtaTileMNK = CtaTileMNK_;
+  using ThreadEpilogueOp = ThreadEpilogueOp_;
+  using EpilogueTile = EpilogueTile_;
+  using ElementC = ElementC_;
+  using StrideC = StrideC_;
+  using ElementD = ElementD_;
+  using StrideD = StrideD_;
+  using SmemLayoutAtomC = SmemLayoutAtomC_;
+  using SmemLayoutAtomD = SmemLayoutAtomD_;
+  using CopyOpR2S = CopyOpR2S_;
+  using CopyOpS2R = CopyOpS2R_;
+  using ElementCompute = typename ThreadEpilogueOp::ElementCompute;
+  using ElementAccumulator = typename ThreadEpilogueOp::ElementAccumulator;
+  using GmemTiledCopyC = CopyOpG2S;
+  using GmemTiledCopyD = CopyOpS2G;
+
+  static_assert(!is_layout<EpilogueTile>::value && is_tuple<EpilogueTile>::value, "EpilogueTile must be a mute::Tile or mute::Shape");
+  static_assert(mute::rank(CtaTileMNK{}) == 3, "CtaTileMNK must be rank-3: [CTA_M, CTA_N, CTA_K]");
+  static_assert(mute::rank(EpilogueTile{}) == 2, "EpilogueTile must be rank-2: [EPI_TILE_M, EPI_TILE_N]");
+  static_assert(size<0>(CtaTileMNK{}) % size<0>(shape(EpilogueTile{})) == 0, "EPI_TILE_M must divide CTA_M");
+  static_assert(size<1>(CtaTileMNK{}) % size<1>(shape(EpilogueTile{})) == 0, "EPI_TILE_N must divide CTA_N");
+  static_assert(mute::rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]");
+  static_assert(mute::rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]");
+
+private:
+  constexpr static bool is_source_supported = not mute::is_void_v<ElementC>;
+  constexpr static bool is_destination_supported = not mute::is_void_v<ElementD>;
+  using NonVoidElementD = ElementD;
+  static_assert(not mute::is_void_v<NonVoidElementD>, "SmemElementD is void");
+  using NonVoidElementC = mute::conditional_t<not is_source_supported,NonVoidElementD,ElementC>; // prevents void ref breakages
+
+  using SmemElementC = typename mutlass::detail::get_unpacked_element_type<NonVoidElementC>::type;
+  using SmemElementD = typename mutlass::detail::get_unpacked_element_type<NonVoidElementD>::type;
+  constexpr static int StagesC = StagesC_;
+  constexpr static int StagesD = StagesD_;
+  constexpr static bool is_m_major_C = detail::is_m_major<StrideC>();
+  constexpr static bool is_m_major_D = detail::is_m_major<StrideD>();
+  constexpr static bool ReuseSmemC = ReuseSmemC_ and is_destination_supported;
+  constexpr static bool DelayTmeStore = DelayTmeStore_;
+
+  using SmemLayoutC = decltype(tile_to_shape(
+      SmemLayoutAtomC{},
+      make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<StagesC>{}),
+      mute::conditional_t<is_m_major_C, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
+  using SmemLayoutD = decltype(tile_to_shape(
+      SmemLayoutAtomD{},
+      make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<ReuseSmemC ? StagesC : StagesD>{}),
+      mute::conditional_t<is_m_major_D, Step<_2,_1,_3>, Step<_1,_2,_3>>{} ));
+  constexpr static bool support_smem_reuse = is_source_supported && is_destination_supported && StagesD <= StagesC
+                                            && cosize(take<0,2>(SmemLayoutC{})) == cosize(take<0,2>(SmemLayoutD{}));
+  static_assert(not (ReuseSmemC && not support_smem_reuse), "Smem reuse requirements not met");
+  constexpr static size_t SmemAlignmentD = mutlass::detail::alignment_for_swizzle(SmemLayoutD{});
+  constexpr static size_t SmemAlignmentC = mutlass::detail::alignment_for_swizzle(SmemLayoutC{});
+  constexpr static size_t MaxSmemAlignment = mute::max(SmemAlignmentC, SmemAlignmentD);
+
+  struct CollectiveStorageWithC {
+    mute::array_aligned<SmemElementC, mute::cosize_v<SmemLayoutC>, SmemAlignmentC> smem_C;
+    mute::array_aligned<SmemElementD, mute::cosize_v<SmemLayoutD>, SmemAlignmentD> smem_D;
+  };
+
+  union CollectiveStorageWithoutC {
+    mute::array<SmemElementC, 0> smem_C;
+    mute::array_aligned<SmemElementD, mute::cosize_v<SmemLayoutD>, SmemAlignmentD> smem_D;
+  };
+
+  union CollectiveStorageReuseC {
+    mute::array_aligned<SmemElementC, mute::cosize_v<SmemLayoutC>, MaxSmemAlignment> smem_C;
+    mute::array_aligned<SmemElementD, mute::cosize_v<SmemLayoutD>, MaxSmemAlignment> smem_D;
+  };
+
+public:
+  // TME pipeline for loading C
+  using LoadPipeline       = mutlass::Mp31PipelineTmeAsync<DispatchPolicy::StagesC>;
+  using LoadPipelineParams = typename LoadPipeline::Params;
+  using LoadPipelineState  = mutlass::PipelineState<DispatchPolicy::StagesC>;
+  using StorePipelineState = mute::conditional_t<ReuseSmemC,
+                             mutlass::PipelineState<DispatchPolicy::StagesC>,
+                             mutlass::PipelineState<DispatchPolicy::StagesD>>;
+  static constexpr int NumBarriers = is_source_supported ? LoadPipeline::NumBarriers: 0;
+  struct SharedStorage
+  {
+    using CollectiveStorage = mute::conditional_t<not is_source_supported, CollectiveStorageWithoutC,
+                              mute::conditional_t<ReuseSmemC, CollectiveStorageReuseC, CollectiveStorageWithC>>;
+    CollectiveStorage collective;
+  };
+
+
+  // Host side epilogue arguments
+  struct Arguments {
+    typename ThreadEpilogueOp::Params thread{};
+    ElementC const* ptr_C;
+    StrideC dC;
+    ElementD const* ptr_D;
+    StrideD dD;
+  };
+
+  // Device side epilogue params
+  struct Params {
+    using TME_C = decltype(make_tme_copy(
+        CopyOpG2S{},
+        make_tensor(make_gmem_ptr(static_cast<NonVoidElementC const*>(nullptr)),
+            repeat_like(StrideC{}, int32_t(0)), StrideC{}),
+        take<0, 2>(SmemLayoutC{}),
+        make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}))));
+    using TME_D = decltype(make_tme_copy(
+        CopyOpS2G{},
+        make_tensor(make_gmem_ptr(static_cast<NonVoidElementD const*>(nullptr)),
+            repeat_like(StrideD{}, int32_t(0)), StrideD{}),
+        take<0, 2>(SmemLayoutD{}),
+        make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}))));
+    typename ThreadEpilogueOp::Params thread{};
+    TME_C tme_load_c;
+    TME_D tme_store_d;
+  };
+
+  //
+  // Methods
+  //
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    typename Params::TME_C tme_load_c = {};
+    if constexpr (is_source_supported) {
+      Tensor tensor_c = make_tensor(args.ptr_C, make_layout(make_shape(M,N,L), args.dC));
+      tme_load_c = make_tme_copy(
+        MP31_TME_LOAD{},
+        tensor_c,
+        take<0, 2>(SmemLayoutC{}),
+        make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{})));
+    }
+
+    typename Params::TME_D tme_store_d = {};
+    if constexpr (is_destination_supported) {
+      Tensor tensor_d = make_tensor(args.ptr_D, make_layout(make_shape(M,N,L), args.dD));
+      tme_store_d = make_tme_copy(
+          MP31_TME_STORE{},
+          tensor_d,
+          take<0, 2>(SmemLayoutD{}),
+          make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{})));
+    }
+
+    typename ThreadEpilogueOp::Params thread{args.thread};
+
+    return {thread, tme_load_c, tme_store_d};
+  }
+
+  template <class ProblemShape>
+  MUTLASS_HOST_DEVICE static bool
+  can_implement(
+      ProblemShape const& problem_shape,
+      [[maybe_unused]] Arguments const& args) {
+    constexpr int tme_alignment_bits = 32;
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    bool implementable = true;
+    if constexpr (not mute::is_void_v<ElementC>) {
+      constexpr int min_tme_aligned_elements_C = tme_alignment_bits / mutlass::sizeof_bits<ElementC>::value;
+      implementable = implementable && mutlass::detail::check_alignment<min_tme_aligned_elements_C>(mute::make_shape(M,N), StrideC{});
+    }
+
+    if constexpr (is_destination_supported) {
+      constexpr int min_tme_aligned_elements_D = tme_alignment_bits / mutlass::sizeof_bits<ElementD>::value;
+      implementable = implementable && mutlass::detail::check_alignment<min_tme_aligned_elements_D>(mute::make_shape(M,N), StrideD{});
+    }
+
+    if (!implementable) {
+      MUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TME.\n");
+    }
+    return implementable;
+  }
+
+  MUTLASS_HOST_DEVICE
+  CollectiveEpilogue(Params const& params_)
+      : params(params_), epilogue_op(params_.thread) {}
+
+  MUTLASS_DEVICE
+  bool
+  is_source_needed() {
+    return is_source_supported && epilogue_op.is_source_needed();
+  }
+
+  template<
+    class ProblemShapeMNKL,
+    class TileShapeMNK,
+    class TileCoordMNKL,
+    class AccEngine, class AccLayout,
+    class TiledMma,
+    class ResidueMNK
+  >
+  MUTLASS_DEVICE auto
+  operator()(
+      ProblemShapeMNKL problem_shape_mnkl,
+      TileShapeMNK tile_shape_MNK,
+      TileCoordMNKL tile_coord_mnkl,
+      mute::Tensor<AccEngine,AccLayout> accumulators,
+      TiledMma tiled_mma,
+      ResidueMNK residue_mnk,
+      int thread_idx,
+      char* smem_buf) {
+    using namespace mute;
+    using ElementAccumulator = typename AccEngine::value_type;
+
+    static_assert(is_rmem<AccEngine>::value, "Accumulator must be RF resident.");
+    static_assert(rank(AccLayout{}) == 3, "Accumulator must be MMA-partitioned: (MMA,MMA_M,MMA_N)");
+    static_assert(rank(ProblemShapeMNKL{}) == 4, "ProblemShapeMNKL must be rank 4");
+    static_assert(is_static<TileShapeMNK>::value, "TileShapeMNK must be static");
+    static_assert(rank(TileShapeMNK{}) == 3, "TileShapeMNK must be rank 3");
+    static_assert(rank(TileCoordMNKL{}) == 4, "TileCoordMNKL must be rank 4");
+
+    // Indexing variables
+    auto [M, N, K, L] = problem_shape_mnkl;
+    auto [m_coord, n_coord, k_coord, l_coord] = tile_coord_mnkl;
+
+    Tensor mD_mn = params.tme_store_d.get_tme_tensor(make_shape(M,N,L));                               // (M,N,L)
+    Tensor mD = coalesce(mD_mn, take<0,2>(CtaTileMNK{}));
+    Tensor gD = local_tile(mD, take<0,2>(CtaTileMNK{}), make_coord(m_coord, n_coord, l_coord));     // (CTA_M,CTA_N)
+
+    Tensor mC_mn = params.tme_load_c.get_tme_tensor(make_shape(M,N,L));                             //  (M,N,L)
+    Tensor mC = coalesce(mC_mn, take<0,2>(CtaTileMNK{}));
+    Tensor gC = local_tile(mC, take<0,2>(CtaTileMNK{}), make_coord(m_coord, n_coord, l_coord));
+
+    Tensor gC_epi = flat_divide(gC, EpilogueTile{});                             // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
+    Tensor gD_epi = flat_divide(gD, EpilogueTile{});                             // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
+
+    // Apply epilogue subtiling
+    SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+    Tensor sC_epi = make_tensor(make_smem_ptr(storage.collective.smem_C.begin()), SmemLayoutC{});             // (EPI_TILE_M,EPI_TILE_N,PIPE_C)
+    Tensor sD_epi = make_tensor(make_smem_ptr(storage.collective.smem_D.begin()), SmemLayoutD{});             // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
+
+    constexpr int FragmentSize = DispatchPolicy::FragmentSize;
+    using CopyV = Int<FragmentSize>;
+    TiledCopy r2s_tiled_copy = detail::make_tiled_copy_C_sqmma<CopyOpR2S, CopyV, EpilogueTile>(tiled_mma);
+    ThrCopy thread_r2s = r2s_tiled_copy.get_thread_slice(thread_idx);
+    Tensor tRS_rAcc = thread_r2s.retile_S(accumulators);                                   // ((R2S,R2S_V),MMA_M,MMA_N)
+    Tensor tRS_sD   = thread_r2s.partition_D(sD_epi);                                      // (R2S,R2S_M,R2S_N,PIPE_D)
+
+    auto mma_tile_m = size<0>(TileShapeMNK{}) / size<1>(tRS_rAcc);
+    auto mma_tile_n = size<1>(TileShapeMNK{}) / size<2>(tRS_rAcc);
+    auto epi_tile_m = size<0>(EpilogueTile{});
+    auto epi_tile_n = size<1>(EpilogueTile{});
+    // Allocate D registers
+    Layout tRS_rD_layout = make_layout(take<0,3>(shape(thread_r2s.partition_D(sD_epi))));
+    Tensor tRS_rD = make_tensor<SmemElementD>(tRS_rD_layout);                                   // (R2S,R2S_M,R2S_N)
+
+    TiledCopy tiled_s2r  = detail::make_tiled_copy_C_sqmma<CopyOpS2R, CopyV, EpilogueTile>(tiled_mma);
+    ThrCopy thread_s2r   = tiled_s2r.get_slice(thread_idx);
+    Tensor tSR_sC        = thread_s2r.partition_S(sC_epi);                                  // (S2R,S2R_M,S2R_N,PIPE_C)
+    Layout tSR_rC_layout = thread_s2r.retile_D(tRS_rD).layout();                            // (S2R,S2R_M,S2R_N)
+
+    Tensor tRS_rC = make_tensor<SmemElementC>(tRS_rD_layout);                             // (R2S,R2S_M,R2S_N)
+    Tensor tSR_rC = thread_s2r.retile_D(tRS_rC);                                            // (S2R,S2R_M,S2R_N)
+
+    ThrCopy thrblk_g2s = params.tme_load_c.get_slice(Int<0>{});
+    Tensor bGS_gC      = thrblk_g2s.partition_S(gC_epi);                                    // (G2S,G2S_M,G2S_N,EPI_M,EPI_N)
+    Tensor bGS_sC      = thrblk_g2s.partition_D(sC_epi);                                    // (G2S,G2S_M,G2S_N,PIPE_C)
+    // thread(b)lock-partition for (s)mem to (g)mem copy (bSG_)
+    ThrCopy thrblk_s2g = params.tme_store_d.get_slice(Int<0>{});
+    Tensor bSG_sD      = thrblk_s2g.partition_S(sD_epi);                                    // (S2G,S2G_M,S2G_N,PIPE_D)
+    Tensor bSG_gD      = thrblk_s2g.partition_D(gD_epi);                                    // (S2G,S2G_M,S2G_N,EPI_M,EPI_N)
+
+    MUTE_STATIC_ASSERT(mma_tile_n % epi_tile_n == 0, "MMA_TILE_M must divide EPI_TILE_M");
+
+    MUTE_STATIC_ASSERT(mma_tile_m % epi_tile_m == 0, "EPI_TILE_N must divide MMA_TILE_N");
+
+    // Predication for TMA store (one warp issues TMA store)
+    bool issue_tme_store = canonical_warp_idx() == 0;
+    bool issue_tme_load  = canonical_warp_idx() == 0;
+    constexpr uint32_t start_store = DelayTmeStore ? StagesD - 1: 0;
+    constexpr uint32_t TmeTransactionBytes = static_cast<uint32_t>(
+        size<0>(EpilogueTile{}) * size<1>(EpilogueTile{}) * sizeof(SmemElementC));
+    LoadPipelineParams pipe_params;
+    pipe_params.transaction_bytes = TmeTransactionBytes;
+    constexpr int NumWarps = size(TiledMma{}) / NumThreadsPerWarp;
+    pipe_params.num_consumers = NumWarps;
+    LoadPipeline load_pipeline(pipe_params);
+
+    LoadPipelineState smem_pipe_load_read;
+    LoadPipelineState smem_pipe_load_write = mutlass::make_producer_start_state<LoadPipelineState>();
+    LoadPipelineState smem_pipe_release;
+    StorePipelineState smem_pipe_store_read;
+    StorePipelineState smem_pipe_store_write;
+
+
+    uint32_t tme_load_tile = 0;
+    uint32_t tme_store_tile = 0;
+    __syncthreads();
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int pipe_c = 0; pipe_c < StagesC-1; ++pipe_c) {
+      if (issue_tme_load &&
+          is_source_needed() &&
+          tme_load_tile < size<2>(gD_epi) * size<3>(gD_epi)) {
+        load_pipeline.producer_acquire(smem_pipe_load_write);
+        uint32_t bar_id = load_pipeline.producer_get_barrier_id(smem_pipe_load_write);
+        uint32_t cur_epi_m = tme_load_tile / size<3>(gD_epi);
+        uint32_t cur_epi_n = tme_load_tile % size<3>(gD_epi);
+        copy(params.tme_load_c.with(bar_id),bGS_gC(_,_,_,cur_epi_m,cur_epi_n), bGS_sC(_,_,_,pipe_c));
+      }
+      ++smem_pipe_load_write;
+      ++tme_load_tile;
+    }
+    MUTLASS_PRAGMA_UNROLL
+    for (int epi_m = 0; epi_m < size<2>(gD_epi); ++epi_m) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int epi_n = 0; epi_n < size<3>(gD_epi); ++epi_n) {
+        if (is_source_needed()) {
+          load_pipeline.consumer_wait(smem_pipe_load_read);
+          copy(tiled_s2r, tSR_sC(_,_,_,smem_pipe_load_read.index()), tSR_rC);
+        }
+        int mma_m = (epi_m * size<0>(EpilogueTile{})) / mma_tile_m;
+        int mma_n = (epi_n * size<1>(EpilogueTile{})) / mma_tile_n;
+        Tensor tRS_rAcc_mn = tRS_rAcc(_,mma_m,mma_n);
+        if constexpr (size<3>(gD_epi) > 1 && size<2>(gD_epi) == 1) {
+          MUTLASS_PRAGMA_UNROLL
+          for (int epi_v = 0; epi_v < size(tRS_rD); ++epi_v) {
+            if (is_source_needed()) {
+              tRS_rD(epi_v) = epilogue_op(tRS_rAcc_mn(epi_n + epi_v * size<3>(gD_epi)), tSR_rC(epi_v));
+            } else {
+              tRS_rD(epi_v) = ElementD(tRS_rAcc_mn(epi_n + epi_v * size<3>(gD_epi)));
+            }
+          }
+        } else {
+          int epi_m_in_mma = epi_m % (mma_tile_m / epi_tile_m);
+          int epi_n_in_mma = epi_n % (mma_tile_n / epi_tile_n);
+          int r2s_v = epi_m_in_mma * size<3>(gD_epi) * size(tRS_rD) + epi_n_in_mma * size(tRS_rD);
+          MUTLASS_PRAGMA_UNROLL
+          for (int epi_v = 0; epi_v < size(tRS_rD); ++epi_v) {
+            if (is_source_needed()) {
+              tRS_rD(epi_v) = epilogue_op(tRS_rAcc_mn(r2s_v + epi_v), tSR_rC(epi_v));
+            } else {
+              tRS_rD(epi_v) = ElementD(tRS_rAcc_mn(r2s_v + epi_v));
+            }
+          }
+        }
+
+        // copy tile from register to smem
+        if (tme_store_tile > start_store) {
+          if (issue_tme_store) {
+            mute::tme_store_wait();
+          }
+          if (not is_source_needed()) {
+            __syncthreads();
+          }
+        }
+        copy(r2s_tiled_copy, tRS_rD, tRS_sD(_,_,_,smem_pipe_store_write.index()));
+        // if reuse smemC or StagesD == 1, supposed to do s2g immediately
+        if constexpr (not DelayTmeStore) {
+           __syncthreads();
+           if (issue_tme_store) {
+            copy(params.tme_store_d, bSG_sD(_,_,_,smem_pipe_store_read.index()), bSG_gD(_,_,_,epi_m, epi_n));
+            mute::tme_store_arrive();
+           }
+          ++smem_pipe_store_read;
+        } else if (tme_store_tile >= start_store){
+            uint32_t cur_epi_m = (tme_store_tile - StagesD + 1) / size<3>(gD_epi);
+            uint32_t cur_epi_n = (tme_store_tile - StagesD + 1) % size<3>(gD_epi);
+            if (issue_tme_store) {
+              copy(params.tme_store_d, bSG_sD(_,_,_,smem_pipe_store_read.index()), bSG_gD(_,_,_,cur_epi_m, cur_epi_n));
+              mute::tme_store_arrive();
+            }
+            ++smem_pipe_store_read;
+        }
+
+        if (issue_tme_load &&
+            is_source_needed() &&
+            tme_load_tile < size<2>(gD_epi) * size<3>(gD_epi)) {
+          load_pipeline.producer_acquire(smem_pipe_load_write);
+          uint32_t bar_id = load_pipeline.producer_get_barrier_id(smem_pipe_load_write);
+          uint32_t cur_epi_m = tme_load_tile / size<3>(gD_epi);
+          uint32_t cur_epi_n = tme_load_tile % size<3>(gD_epi);
+          copy(params.tme_load_c.with(bar_id), bGS_gC(_,_,_,cur_epi_m, cur_epi_n), bGS_sC(_,_,_,smem_pipe_load_write.index()));
+        }
+
+        if (is_source_needed()) {
+          load_pipeline.consumer_release(smem_pipe_release);
+        }
+
+        ++tme_load_tile;
+        ++tme_store_tile;
+        ++smem_pipe_load_write;
+        ++smem_pipe_load_read;
+        ++smem_pipe_store_write;
+        ++smem_pipe_release;
+
+      } // for epi_n
+    } // for epi_m
+    // store residual smem_D
+    if constexpr (DelayTmeStore) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int pip_d = 0; pip_d < StagesD -1; ++pip_d) {
+        if (issue_tme_store) {
+          mute::tme_store_wait();
+          uint32_t cur_epi_m = (tme_store_tile - StagesD + 1) / size<3>(gD_epi);
+          uint32_t cur_epi_n = (tme_store_tile - StagesD + 1) % size<3>(gD_epi);
+          copy(params.tme_store_d, bSG_sD(_,_,_,smem_pipe_store_read.index()), bSG_gD(_,_,_,cur_epi_m, cur_epi_n));
+          ++smem_pipe_store_read;
+          ++tme_store_tile;
+          mute::tme_store_arrive();
+        }
+      }
+    }
+  }
+
+private:
+  Params const& params;
+  ThreadEpilogueOp epilogue_op;
+};
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace collective
+} // namespace epilogue
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/dispatch_policy.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/dispatch_policy.hpp
new file mode 100644
index 000000000..1565488d7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/dispatch_policy.hpp
@@ -0,0 +1,78 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/epilogue/thread/scale_type.h"
+
+//////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue {
+
+//////////////////////////////////////////////////////////////////////////////
+//
+// Builder Epilogue Schedules
+//
+//////////////////////////////////////////////////////////////////////////////
+
+struct NoSmem  {};
+struct WithTme {};
+struct EpilogueSimtVectorized {};
+//////////////////////////////////////////////////////////////////////////////
+
+template <
+  int Stages_,
+  int FragmentSize_
+>
+struct Mp22CollectiveEpilogue {
+  constexpr static int Stages = Stages_;
+  constexpr static int FragmentSize = FragmentSize_;
+};
+
+template <
+  int StagesC_,
+  int StagesD_,
+  int FragmentSize_,
+  bool ReuseSmemC_,
+  bool DelayTmeStore_ = false
+>
+struct Mp31CollectiveEpilogue {
+  constexpr static int StagesC = StagesC_;
+  constexpr static int StagesD = StagesD_;
+  constexpr static int FragmentSize = FragmentSize_;
+  constexpr static bool ReuseSmemC = ReuseSmemC_;
+  constexpr static bool DelayTmeStore = DelayTmeStore_;
+};
+
+//////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/callbacks.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/callbacks.hpp
new file mode 100644
index 000000000..d61a4030f
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/callbacks.hpp
@@ -0,0 +1,90 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mutlass/detail/dependent_false.hpp"
+#include "mutlass/epilogue/fusion/operations.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::fusion {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Dispatch interface for epilogue fusion callbacks
+// For visitor fusions, this is just a convenience wrapper to provide metadata and non-nested args.
+// It is also valid to just pass visitor callbacks directly to the collective,
+// provided the collective supports a visitor callbacks interface. This is useful for implementing custom fusions.
+template <
+  class DispatchPolicy,  // specialize on collective's dispatch policy since callbacks API will depend on collective's algorithm
+  class Operation,       // the fusion operation being performed, e.g. fusion::LinearCombination
+  class CtaTile_MNK,     // computed tile per CTA
+  class EpilogueTile_MN, // epilogue subtile size
+  class... Args          // callbacks implementation dependent args (e.g. copy atoms, smem layouts)
+>
+struct FusionCallbacks {
+  static_assert(mutlass::detail::dependent_false<DispatchPolicy, Operation>, "Could not find a callbacks specialization.");
+};
+
+// Metadata helper to handle custom EVTs or other non-FusionCallbacks types
+template <class T>
+struct FusionCallbacksTraits {
+  using DispatchPolicy = void;
+  using Operation = T;
+  using CtaTile_MNK = void;
+  using EpilogueTile_MN = void;
+  using ElementCompute = void;
+};
+
+template <
+  class DispatchPolicy_,
+  class Operation_,
+  class CtaTile_MNK_,
+  class EpilogueTile_MN_,
+  class... Args
+>
+struct FusionCallbacksTraits<
+  FusionCallbacks<DispatchPolicy_, Operation_, CtaTile_MNK_, EpilogueTile_MN_, Args...>
+> {
+  using DispatchPolicy = DispatchPolicy_;
+  using Operation = Operation_;
+  using CtaTile_MNK = CtaTile_MNK_;
+  using EpilogueTile_MN = EpilogueTile_MN_;
+  using ElementCompute = typename Operation::ElementCompute;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::fusion
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_callbacks.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_callbacks.hpp
new file mode 100644
index 000000000..6d6ee3f72
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_callbacks.hpp
@@ -0,0 +1,168 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/mutlass.h"
+
+#include "mute/tensor.hpp"
+
+#include "mutlass/epilogue/dispatch_policy.hpp"
+#include "mutlass/epilogue/fusion/callbacks.hpp"
+#include "mutlass/epilogue/fusion/mp22_visitor.hpp"
+#include "mutlass/epilogue/fusion/mp22_visitor_load.hpp"
+#include "mutlass/epilogue/fusion/mp22_visitor_store.hpp"
+#include "mutlass/epilogue/fusion/mp22_visitor_compute.hpp"
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::fusion {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// D = alpha * acc
+template <
+  int Stages,
+  int FragmentSize,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementScalar,
+  FloatRoundStyle RoundStyle,
+  class CtaTileShapeMNK,
+  class EpilogueTile
+>
+struct FusionCallbacks<
+  epilogue::Mp22CollectiveEpilogue<Stages, FragmentSize>,
+  fusion::ScaledAcc<ElementOutput, ElementCompute, ElementScalar, RoundStyle>,
+  CtaTileShapeMNK,
+  EpilogueTile
+> : Mp22EVT<Mp22Compute<multiplies, ElementOutput, ElementCompute, RoundStyle>,
+      Mp22ScalarBroadcast<ElementScalar>,
+      Mp22AccFetch
+    > {
+  using Impl =
+    Mp22EVT<Mp22Compute<multiplies, ElementOutput, ElementCompute, RoundStyle>,
+      Mp22ScalarBroadcast<ElementScalar>,
+      Mp22AccFetch
+    >;
+  using Operation = fusion::ScaledAcc<ElementOutput, ElementCompute, ElementScalar, RoundStyle>;
+
+  struct Arguments {
+    // Give a name and flat ordering to the fusion callback args
+    ElementScalar alpha = ElementScalar(1);
+    ElementScalar beta = ElementScalar(0);
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+
+    // Conversion to the args expected by the visitor implementation
+    // to_underlying_arguments will implicitly call this
+    operator typename Impl::Arguments() const {
+      return
+        {    // binary op : alpha * acc
+          {{alpha}, {alpha_ptr}}, // leaf args : alpha
+          {},                     // leaf args : acc
+          {} // binary args : multiplies
+        };   // end binary op
+    }
+  };
+
+  // Ctor inheritance
+  using Impl::Impl;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// D = alpha * acc + beta * C
+template<
+  class ElementOutput,
+  class ElementCompute,
+  class ElementSource = ElementOutput,
+  class ElementScalar = ElementCompute,
+  FloatRoundStyle RoundStyle = FloatRoundStyle::round_to_nearest
+>
+using Mp22LinearCombination =
+  Mp22EVT<Mp22Compute<homogeneous_multiply_add, ElementOutput, ElementCompute, RoundStyle>, // beta * C + (alpha * acc)
+    Mp22ScalarBroadcast<ElementScalar>, // beta
+    Mp22SrcFetch<ElementSource>, // C
+    Mp22EVT<Mp22Compute<multiplies, ElementCompute, ElementCompute, RoundStyle>, // alpha * acc
+      Mp22ScalarBroadcast<ElementScalar>, // alpha
+      Mp22AccFetch // acc
+    >
+  >;
+
+template <
+  int Stages,
+  int FragmentSize,
+  class ElementOutput,
+  class ElementCompute,
+  class ElementSource,
+  class ElementScalar,
+  FloatRoundStyle RoundStyle,
+  class CtaTileShapeMNK,
+  class EpilogueTile
+>
+struct FusionCallbacks<
+    epilogue::Mp22CollectiveEpilogue<Stages, FragmentSize>,
+    fusion::LinearCombination<ElementOutput, ElementCompute, ElementSource, ElementScalar, RoundStyle>,
+    CtaTileShapeMNK,
+    EpilogueTile
+> : Mp22LinearCombination<typename mutlass::detail::get_unpacked_element_type<ElementOutput>::type, ElementCompute, ElementSource, ElementScalar, RoundStyle> {
+
+  using Impl = Mp22LinearCombination<typename mutlass::detail::get_unpacked_element_type<ElementOutput>::type, ElementCompute, ElementSource, ElementScalar, RoundStyle>;
+  using Operation = fusion::LinearCombination<ElementOutput, ElementCompute, ElementSource, ElementScalar, RoundStyle>;
+
+  struct Arguments {
+    ElementScalar alpha = ElementScalar(1);
+    ElementScalar beta = ElementScalar(0);
+    ElementScalar const* alpha_ptr = nullptr;
+    ElementScalar const* beta_ptr = nullptr;
+
+    operator typename Impl::Arguments() const {
+      return
+        {    // ternary op : beta * C + (alpha * acc)
+          {{beta}, {beta_ptr}}, // leaf args : beta
+          {},                   // leaf args : C
+          {                     // binary op : alpha * acc
+            {{alpha}, {alpha_ptr}}, // leaf args : alpha
+            {},                     // leaf args : acc
+            {}                  // binary args : multiplies
+          },                    // end binary op
+          {} // ternary args : multiply_add
+        };   // end ternary op
+    }
+  };
+
+  // Ctor inheritance
+  using Impl::Impl;
+};
+
+} // namespace mutlass::epilogue::fusion
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor.hpp
new file mode 100644
index 000000000..400769193
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor.hpp
@@ -0,0 +1,259 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "visitor_base.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::fusion {
+
+using namespace mute;
+using mute::tuple;
+
+namespace detail {
+
+template <
+  class ProblemShapeMNKL,
+  class TileShapeMNK,
+  class TileCoordMNKL,
+  class ResidueMN,
+  class EpilogueTile,
+  class TiledCopy,
+  class CoordTensor,
+  class ThrCoordTensor,
+  class SrcTensor,
+  class DstTensor
+>
+struct VisitorArgs {
+  ProblemShapeMNKL problem_shape_mnkl;
+  TileShapeMNK tile_shape_mnk;
+  TileCoordMNKL tile_coord_mnkl;
+  ResidueMN residue_mn;
+  EpilogueTile epi_tile;
+  TiledCopy tiled_copy;
+  int thread_idx;
+  CoordTensor cD;
+  ThrCoordTensor tCcD;
+  SrcTensor const& tCrC;
+  DstTensor const& tDgD;
+
+  MUTLASS_DEVICE
+  VisitorArgs(
+     ProblemShapeMNKL problem_shape_mnkl,
+     TileShapeMNK tile_shape_mnk,
+     TileCoordMNKL tile_coord_mnkl,
+     ResidueMN residue_mn,
+     EpilogueTile epi_tile,
+     TiledCopy tiled_copy,
+     int thread_idx,
+     CoordTensor cD,
+     ThrCoordTensor tCcD,
+     SrcTensor const& tCrC,
+     DstTensor const& tDgD)
+  : problem_shape_mnkl(problem_shape_mnkl),
+    tile_shape_mnk(tile_shape_mnk),
+    tile_coord_mnkl(tile_coord_mnkl),
+    residue_mn(residue_mn),
+    epi_tile(epi_tile),
+    tiled_copy(tiled_copy),
+    thread_idx(thread_idx),
+    cD(cD),
+    tCcD(tCcD),
+    tCrC(tCrC),
+    tDgD(tDgD) {}
+};
+
+template <class... Ops>
+struct Mp22VisitorImpl : VisitorImplBase<Ops...> {
+  using Impl = VisitorImplBase<Ops...>;
+  using Params = typename Impl::Params;
+  using SharedStorage = typename Impl::SharedStorage;
+
+  MUTLASS_HOST_DEVICE
+  Mp22VisitorImpl() {}
+
+  MUTLASS_HOST_DEVICE
+  Mp22VisitorImpl(Params const& params, SharedStorage const& shared_storage)
+    : Impl(params, shared_storage) {}
+
+  using Impl::ops;
+
+  MUTLASS_DEVICE bool
+  is_C_load_needed() const {
+    return apply(ops,
+      [] (auto const&... op) {
+        return (false || ... || op.is_C_load_needed());
+      }
+    );
+  }
+
+  template <class CallbacksTuple>
+  struct Callbacks {
+    // Callbacks can store non-persistent variables (e.g. tensors) or copies of persistent variables
+    CallbacksTuple callbacks_tuple;
+
+    MUTLASS_DEVICE void
+    begin_epilogue() {
+      for_each(callbacks_tuple,
+        [] (auto& callbacks) {
+          callbacks.begin_epilogue();
+        }
+      );
+    }
+
+    MUTLASS_DEVICE void
+    begin_step(int epi_m, int epi_n, int step_idx) {
+      for_each(callbacks_tuple,
+        [&] (auto& callbacks) {
+          callbacks.begin_step(epi_m, epi_n, step_idx);
+        }
+      );
+    }
+
+    // Perform the fused elementwise computation
+    template <typename ElementAccumulator, typename... ElementInputs, int FragmentSize>
+    MUTLASS_DEVICE auto // returns an Array
+    visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n,
+          Array<ElementInputs, FragmentSize> const&... frg_inputs) // depends on the N-naryness of the op
+      = delete; // Must be implemented for each operation
+
+    MUTLASS_DEVICE void
+    end_step(int epi_m, int epi_n, int step_idx) {
+      for_each(callbacks_tuple,
+        [&] (auto& callbacks) {
+          callbacks.end_step(epi_m, epi_n, step_idx);
+        }
+      );
+    }
+
+    MUTLASS_DEVICE void
+    end_epilogue() {
+      for_each(callbacks_tuple,
+        [&] (auto& callbacks) {
+          callbacks.end_epilogue();
+        }
+      );
+    }
+  };
+
+
+  // Callbacks factory
+  // All operations must redefine this
+  template <class... Args>
+  MUTLASS_DEVICE auto
+  get_callbacks(VisitorArgs<Args...> const& args) {
+    return transform_apply(ops,
+        [&] (auto& op) {
+          return op.get_callbacks(args);
+        },
+        [] (auto&&... callbacks) {
+          auto callbacks_tuple = mute::make_tuple(callbacks...);
+          return Callbacks<decltype(callbacks_tuple)>{callbacks_tuple};
+        }
+    );
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+// Convenience aliases
+using EmptyCallbacks = Mp22VisitorImpl<>::Callbacks<mute::tuple<>>;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace detail
+
+using namespace detail;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Tree visitor
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class NodeOp, class... ChildOps>
+struct Mp22TreeVisitor : Mp22VisitorImpl<ChildOps..., NodeOp> {
+
+  using Impl = Mp22VisitorImpl<ChildOps..., NodeOp>;
+  using Params = typename Impl::Params;
+  using SharedStorage = typename Impl::SharedStorage;
+
+  MUTLASS_HOST_DEVICE
+  Mp22TreeVisitor() {}
+
+  MUTLASS_HOST_DEVICE
+  Mp22TreeVisitor(
+      Params const& params,
+      SharedStorage const& shared_storage)
+    : Impl(params, shared_storage) {}
+
+
+  template <class CallbacksImpl>
+  struct Callbacks : CallbacksImpl {
+    MUTLASS_DEVICE
+    Callbacks(CallbacksImpl&& impl)
+      : CallbacksImpl(mute::forward<CallbacksImpl>(impl)) {}
+
+    using CallbacksImpl::callbacks_tuple;
+
+    template <typename ElementAccumulator, int FragmentSize>
+    MUTLASS_DEVICE auto
+    visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
+      constexpr int Rm1 = sizeof...(ChildOps);
+      return mute::detail::tapply(callbacks_tuple,
+        [&] (auto& child_callbacks) {
+          return child_callbacks.visit(frg_acc, epi_v, epi_m, epi_n);
+        },
+        [&] (auto&&... frg_inputs) {
+          return get<Rm1>(callbacks_tuple).visit(frg_acc, epi_v, epi_m, epi_n, frg_inputs...);
+        },
+        make_seq<Rm1>{}
+      );
+    }
+  };
+
+  // Callbacks factory
+  template <class... Args>
+  MUTLASS_DEVICE auto
+  get_callbacks(VisitorArgs<Args...> const& args) {
+    return Callbacks<
+      decltype(Mp22VisitorImpl<ChildOps..., NodeOp>::get_callbacks(args))>(
+        Mp22VisitorImpl<ChildOps..., NodeOp>::get_callbacks(args));
+  }
+};
+
+template <class NodeOp, class... ChildOps>
+using Mp22EVT = Mp22TreeVisitor<NodeOp, ChildOps...>;
+
+} // namespace mutlass::epilogue::fusion
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_compute.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_compute.hpp
new file mode 100644
index 000000000..728aeb7df
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_compute.hpp
@@ -0,0 +1,183 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/epilogue/fusion/mp22_visitor.hpp"
+
+#include "mute/tensor.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::fusion {
+
+using namespace mute;
+using namespace detail;
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// N-nary Elementwise Compute Operation
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// The template argument provided for ComputeFn must be able to accept
+// exactly one template parameter.  In Standard C++, it's OK for
+// ComputeFn to have other template parameters, as long as those have
+// defaults.  For example, the following struct Foo would work.
+//
+// template<class A, class B = A>
+// struct Foo {
+//   MUTLASS_HOST_DEVICE auto operator() (A a, B b);
+// };
+//
+// However, some compilers, such as Clang, require that the argument
+// take _exactly_ one template parameter.  This is nonstandard C++
+// behavior.  One work-around for this case is to create a subclass
+// with exactly one template parameter, and then use that subclass as
+// the template argument.
+//
+// template<class A>
+// struct FooHomogeneous : public Foo<A, B> {};
+//
+template <
+  template <class> class ComputeFn,
+  class ElementOutput,
+  class ElementCompute,
+  FloatRoundStyle RoundStyle,
+  class = void
+>
+struct Mp22Compute {
+private:
+  using EmptyArguments = typename Mp22VisitorImpl<>::Arguments;
+
+  template <class Fn, class = void>
+  struct ComputeArguments {
+    using type = EmptyArguments;
+  };
+
+
+  // partial specialization for compute fns that define an Arguments member, e.g. activation hyperparameters
+  template <class Fn>
+  struct ComputeArguments<Fn, platform::void_t<typename Fn::Arguments>> {
+    using type = typename Fn::Arguments;
+  };
+
+public:
+  struct SharedStorage { };
+
+  using Arguments = typename ComputeArguments<ComputeFn<ElementCompute>>::type;
+
+  using Params = Arguments;
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    return args;
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return 0;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream) {
+    return mutlass::Status::kSuccess;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Mp22Compute() { }
+
+  MUTLASS_HOST_DEVICE
+  Mp22Compute(Params const& params, SharedStorage const& shared_storage)
+      : params(params) {}
+
+  Params const params;
+
+  MUTLASS_DEVICE bool
+  is_C_load_needed() const {
+    return false;
+  }
+
+  struct Callbacks : EmptyCallbacks {
+    MUTLASS_DEVICE
+    Callbacks(Params const& params)
+      : params(params) {}
+
+    Params const& params;
+
+    template <typename ElementAccumulator, typename... ElementInputs, int FragmentSize>
+    MUTLASS_DEVICE Array<ElementOutput, FragmentSize>
+    visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n,
+          Array<ElementInputs, FragmentSize> const&... frg_inputs) {
+      return transform_apply(mute::make_tuple(frg_inputs...),
+        [&] (auto&& frg_input) {
+          using ElementInput = typename mute::remove_cvref_t<decltype(frg_input)>::Element;
+          using ConvertInput = NumericArrayConverter<ElementCompute, ElementInput, FragmentSize, RoundStyle>;
+          ConvertInput convert_input{};
+
+          return convert_input(frg_input);
+        },
+        [&] (auto&&... cvt_frg_inputs) {
+          using ComputeOutput = ComputeFn<Array<ElementCompute, FragmentSize>>;
+          using ConvertOutput = NumericArrayConverter<ElementOutput, ElementCompute, FragmentSize, RoundStyle>;
+          ComputeOutput compute_output{};
+          ConvertOutput convert_output{};
+
+          if constexpr (is_same_v<Arguments, EmptyArguments>) {
+            return convert_output(compute_output(cvt_frg_inputs...));
+          }
+          else {
+            return convert_output(compute_output(cvt_frg_inputs..., params));
+          }
+        }
+      );
+    }
+  };
+
+  template <class... Args>
+  MUTLASS_DEVICE auto
+  get_callbacks(VisitorArgs<Args...> const& args) {
+    return Callbacks(params);
+  }
+
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::fusion
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_load.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_load.hpp
new file mode 100644
index 000000000..73f708bbf
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_load.hpp
@@ -0,0 +1,393 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/epilogue/fusion/mp22_visitor.hpp"
+
+#include "mute/tensor.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::fusion {
+
+using namespace mute;
+using namespace detail;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Elementwise Fetch Operations
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// returns accumulator
+struct Mp22AccFetch : Mp22VisitorImpl<> {
+
+  using Mp22VisitorImpl<>::Mp22VisitorImpl;
+
+  struct Callbacks : EmptyCallbacks {
+    template <typename ElementAccumulator, int FragmentSize>
+    MUTLASS_DEVICE Array<ElementAccumulator, FragmentSize>
+    visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
+      return frg_acc;
+    }
+  };
+
+  template <class... Args>
+  MUTLASS_DEVICE auto
+  get_callbacks(VisitorArgs<Args...> const& args) {
+    return Callbacks{};
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// return C
+template <class Element>
+struct Mp22SrcFetch : Mp22VisitorImpl<> {
+  MUTLASS_DEVICE bool
+  is_C_load_needed() const {
+    return not is_void_v<Element>;
+  }
+
+  using Mp22VisitorImpl<>::Mp22VisitorImpl;
+
+  template <class SrcTensor>
+  struct Callbacks : EmptyCallbacks {
+    MUTLASS_DEVICE
+    Callbacks(SrcTensor const& tCrC)
+      : tCrC(tCrC) {}
+
+    SrcTensor const& tCrC;
+
+    template <typename ElementAccumulator, int FragmentSize>
+    MUTLASS_DEVICE Array<typename SrcTensor::value_type, FragmentSize>
+    visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
+      return recast<Array<typename SrcTensor::value_type, FragmentSize>>(tCrC)(epi_v);
+    }
+  };
+
+  template <class... Args>
+  MUTLASS_DEVICE auto
+  get_callbacks(VisitorArgs<Args...> const& args) {
+    return Callbacks<mute::remove_cvref_t<decltype(args.tCrC)>>(args.tCrC);
+  }
+};
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Broadcast Load Operations
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Scalar broadcast
+// Supports reduction over multiple broadcasts to support fusions such as fp8 scaling factors
+template <
+  class Element,
+  class StrideMNL = Stride<_0, _0, _0>,
+  int BroadcastCount = 1,
+  template <class> class ReductionFn = multiplies
+>
+struct Mp22ScalarBroadcast {
+  static_assert(
+    (mute::is_same_v<StrideMNL, Stride<_0, _0, _0>>) || // scalar broadcast, e.g. alpha
+    (mute::is_same_v<StrideMNL, Stride<_0, _0, _1>>) || // batched scalar broadcast, e.g. per-batch alpha
+    (mute::is_same_v<StrideMNL, Stride<_0, _0,int>>));
+
+  struct SharedStorage {};
+
+  struct Arguments {
+    Element scalars[BroadcastCount] = {};
+    Element const* scalar_ptrs[BroadcastCount] = {};
+    StrideMNL dScalar = {};
+  };
+
+  using Params = Arguments;
+
+  template <typename ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    return args;
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return 0;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args,
+                       void* workspace, musaStream_t stream) {
+    return mutlass::Status::kSuccess;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Mp22ScalarBroadcast() {}
+
+  MUTLASS_HOST_DEVICE
+  Mp22ScalarBroadcast(Params const& params, SharedStorage const& shared_storage)
+    : params_ptr(&params) {
+    // Get the scalar for non-batched broadcast
+    if constexpr (mute::is_same_v<StrideMNL, Stride<_0, _0, _0>>) {
+      update_scalar();
+    }
+  }
+
+  Element scalar;
+  Params const* params_ptr;
+
+  MUTLASS_DEVICE bool
+  is_C_load_needed() const {
+    return false;
+  }
+
+  struct Callbacks : EmptyCallbacks {
+    MUTLASS_DEVICE
+    Callbacks(Element scalar)
+      : scalar(scalar) {}
+
+    Element scalar;
+
+    template <typename ElementAccumulator, int FragmentSize>
+    MUTLASS_DEVICE Array<Element, FragmentSize>
+    visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
+      Array<Element, FragmentSize> frg_scalar;
+      frg_scalar.fill(scalar);
+
+      return frg_scalar;
+    }
+  };
+
+  template <class... Args>
+  MUTLASS_DEVICE auto
+  get_callbacks(VisitorArgs<Args...> const& args) {
+    if constexpr (
+      mute::is_same_v<StrideMNL, Stride<_0, _0, _0>> ||
+      mute::is_same_v<StrideMNL, Stride<_0, _0,int>>) {
+      auto [m_coord, n_coord, k_coord, l_coord] = args.tile_coord_mnkl;
+      update_scalar(l_coord);
+    }
+
+    return Callbacks(scalar);
+  }
+
+private:
+  MUTLASS_DEVICE void
+  update_scalar(int l_coord = 0) {
+    int l_offset = l_coord * size<2>(params_ptr->dScalar);
+
+    if (params_ptr->scalar_ptrs[0] != nullptr) {
+      scalar = params_ptr->scalar_ptrs[0][l_offset];
+    } else {
+      // batch stride is ignored for nullptr fallback
+      scalar = params_ptr->scalars[0];
+    }
+
+    // Do reduction over multiple broadcasts if necessary
+    ReductionFn<Element> reduction_fn;
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < BroadcastCount; ++i) {
+      if (params_ptr->scalar_ptrs[i] != nullptr) {
+        scalar = reduction_fn(scalar, params_ptr->scalar_ptrs[i][l_offset]);
+      } else {
+        // batch stride is ignored for nullptr fallback
+        scalar = reduction_fn(scalar, params_ptr->scalars[i]);
+      }
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Elementwise Load Operations
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class Element,
+  class StrideMNL
+>
+struct Mp22AuxLoad {
+
+  constexpr static bool is_m_major = epilogue::collective::detail::is_m_major<StrideMNL>();
+  constexpr static bool is_n_major = epilogue::collective::detail::is_n_major<StrideMNL>();
+
+  struct SharedStorage {};
+
+  struct Arguments {
+    Element* ptr_aux = nullptr;
+    Element null_default = Element(0);
+    StrideMNL dAux = {};
+  };
+
+  struct Params {
+    Element* ptr_aux;
+    Element null_default = Element(0);
+    StrideMNL dAux = {};
+    RobustDescriptor desc_aux;
+  };
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    auto problem_shape_MNKL = mute::append<4>(problem_shape, _1{});
+    auto stride_aux         = mute::append<3>(args.dAux, _0{});
+
+    auto [M, N, K, L] = problem_shape_MNKL;
+    auto [dM, dN, dL] = stride_aux;
+
+    size_t elem = 0;
+    if constexpr (is_n_major) {
+      elem = N * dM;
+    } else if constexpr (is_m_major) {
+      elem = M * dN;
+    }
+
+    if (dL != 0) {
+      elem *= L;
+    }
+
+    auto desc_aux = make_robust_desc(args.ptr_aux, elem);
+
+    return {
+      args.ptr_aux,
+      args.null_default,
+      args.dAux,
+      desc_aux
+    };
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return 0;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args,
+                       void* workspace, musaStream_t stream) {
+    return mutlass::Status::kSuccess;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Mp22AuxLoad() {}
+
+  MUTLASS_HOST_DEVICE
+  Mp22AuxLoad(Params const& params, SharedStorage const& shared_storage)
+    : params_ptr(&params) {}
+
+  Params const* params_ptr;
+
+  MUTLASS_DEVICE bool
+  is_C_load_needed() const {
+    return false;
+  }
+
+  template <class GTensor,
+            class RTensor,
+            class TiledCopyG2R>
+  struct Callbacks : EmptyCallbacks {
+    MUTLASS_DEVICE
+    Callbacks(
+      GTensor&& tGRgAux,
+      RTensor&& tGRrAux,
+      TiledCopyG2R tiled_g2r,
+      Params const* params_ptr
+    ):
+      tGRgAux(mute::forward<GTensor>(tGRgAux)),
+      tGRrAux(mute::forward<RTensor>(tGRrAux)),
+      tiled_g2r(tiled_g2r),
+      params_ptr(params_ptr) { }
+
+    GTensor tGRgAux;
+    RTensor tGRrAux;
+    TiledCopyG2R tiled_g2r;
+    Params const* params_ptr;
+
+    MUTLASS_DEVICE void
+    begin_step(int epi_m, int epi_n, int step_idx) {
+      auto src_v = tGRgAux(_, _, _, epi_m, epi_n);
+      auto dst_v = tGRrAux;
+      copy(tiled_g2r.with(params_ptr->desc_aux), src_v, dst_v);
+    }
+
+    template <class ElementAccumulator, int FragmentSize>
+    MUTLASS_DEVICE auto
+    visit(Array<ElementAccumulator, FragmentSize> const& frag_acc, int epi_v, int epi_m, int epi_n) {
+      Tensor tGR_rAux_frg = recast<Array<Element, FragmentSize>>(coalesce(tGRrAux));                          // (EPI_V)
+      return tGR_rAux_frg(epi_v);
+    }
+  };
+
+  template <class... Args>
+  MUTLASS_DEVICE auto
+  get_callbacks(VisitorArgs<Args...> const& args) {
+    auto [M, N, K, L] = args.problem_shape_mnkl;
+    auto [m, n, k, l] = args.tile_coord_mnkl;
+
+    Tensor mAux = make_tensor(
+      make_gmem_ptr(params_ptr->ptr_aux),
+      make_shape(M, N, L),
+      params_ptr->dAux); // (M, N, L)
+
+    Tensor gAux = local_tile(mAux, take<0, 2>(args.tile_shape_mnk), make_coord(m, n, l)); // (BLK_M, BLK_N)
+    Tensor gAux_epi = flat_divide(gAux, args.epi_tile); // (EPI_M, EPI_N, TILE_M, TILE_N)
+    auto thread_copy = args.tiled_copy.get_thread_slice(args.thread_idx);
+    Tensor tDgAux    = thread_copy.partition_S(gAux_epi);
+    constexpr int vec_elem = decltype(max_common_vector(tDgAux, args.tDgD))::value;
+    constexpr int vec_bits = vec_elem * sizeof_bits_v<Element>;
+    using CopyAtomG2R      = Copy_Atom<
+            UniversalCopy<uint_bit_t<mute::min(1024, vec_bits)>>, Element>;
+    using TiledCopyG2R = TiledCopy<CopyAtomG2R,
+                            typename decltype(args.tiled_copy)::TiledLayout_TV,
+                            typename decltype(args.tiled_copy)::Tiler_MN>;
+    TiledCopyG2R tiled_copy_g2r{};
+    auto thr_copy_g2r = tiled_copy_g2r.get_slice(args.thread_idx);
+    auto tGRgAux      = thr_copy_g2r.partition_S(gAux_epi);
+    auto tGRrAux      = thr_copy_g2r.retile_D(make_tensor<Element>(take<0, 3>(shape(tGRgAux))));
+
+    return Callbacks<
+      decltype(tGRgAux), decltype(tGRrAux), TiledCopyG2R>(
+          mute::move(tGRgAux),
+          mute::move(tGRrAux),
+          tiled_copy_g2r,
+          params_ptr);
+  }
+};
+} // namespace mutlass::epilogue::fusion
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_store.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_store.hpp
new file mode 100644
index 000000000..84aeabc3f
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/mp22_visitor_store.hpp
@@ -0,0 +1,240 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/epilogue/fusion/mp22_visitor.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::fusion {
+
+using namespace mute;
+using namespace detail;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Elementwise Store Operations
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class Element,
+  FloatRoundStyle RoundStyle,
+  class StrideMNL
+>
+struct Mp22AuxStore {
+  constexpr static bool is_m_major = epilogue::collective::detail::is_m_major<StrideMNL>();
+  constexpr static bool is_n_major = epilogue::collective::detail::is_n_major<StrideMNL>();
+
+  struct SharedStorage {};
+
+  struct Arguments {
+    Element* ptr_aux = nullptr;
+    StrideMNL dAux = {};
+  };
+
+  struct Params {
+    Element* ptr_aux;;
+    StrideMNL dAux;
+    RobustDescriptor desc_aux;
+  };
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    auto problem_shape_MNKL = mute::append<4>(problem_shape, _1{});
+    auto stride_aux         = mute::append<3>(args.dAux, _0{});
+
+    auto [M, N, K, L] = problem_shape_MNKL;
+    auto [dM, dN, dL] = stride_aux;
+
+    size_t elem = 0;
+    if constexpr (is_n_major) {
+      elem = N * dM;
+    } else if constexpr (is_m_major) {
+      elem = M * dN;
+    }
+
+    if (dL != 0) {
+      elem *= L;
+    }
+
+    auto desc_aux = make_robust_desc(args.ptr_aux, elem);
+
+    return {
+      args.ptr_aux,
+      args.dAux,
+      desc_aux
+    };
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return 0;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args,
+                       void* workspace, musaStream_t stream) {
+    return mutlass::Status::kSuccess;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Mp22AuxStore() {}
+
+  MUTLASS_HOST_DEVICE
+  Mp22AuxStore(Params const& params, SharedStorage const& shared_storage)
+    : params_ptr(&params) {}
+
+  Params const* params_ptr;
+
+  MUTLASS_DEVICE bool
+  is_C_load_needed() const {
+    return false;
+  }
+
+  template <class GTensor,
+            class RTensor,
+            class CTensor,
+            class TiledCopyR2G,
+            class ResidueMN,
+            class TileShapeMNK>
+  struct Callbacks : EmptyCallbacks {
+    MUTLASS_DEVICE
+    Callbacks(
+      GTensor&& tRG_gAux,
+      RTensor&& tRG_rAux,
+      CTensor&& tRG_cAux,
+      TiledCopyR2G tiled_r2g,
+      ResidueMN residue_mn,
+      TileShapeMNK tile_shape,
+      Params const* params_ptr
+    ) :
+      tRG_gAux(mute::forward<GTensor>(tRG_gAux)),
+      tRG_rAux(mute::forward<RTensor>(tRG_rAux)),
+      tRG_cAux(mute::forward<CTensor>(tRG_cAux)),
+      tiled_r2g(tiled_r2g),
+      residue_mn(residue_mn),
+      tile_shape(tile_shape),
+      params_ptr(params_ptr) {}
+
+    GTensor tRG_gAux;
+    RTensor tRG_rAux;
+    CTensor tRG_cAux;
+    TiledCopyR2G tiled_r2g;
+    ResidueMN residue_mn;
+    TileShapeMNK tile_shape;
+    Params const* params_ptr;
+
+    template <typename ElementAccumulator, typename ElementInput, int FragmentSize>
+    MUTLASS_DEVICE auto
+    visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n,
+      Array<ElementInput, FragmentSize> const& frg_input) {
+      using ConvertInput = NumericArrayConverter<Element, ElementInput, FragmentSize, RoundStyle>;
+      ConvertInput convert_input{};
+
+      Tensor tRG_rAux_frg = recast<Array<Element, FragmentSize>>(coalesce(tRG_rAux));                          // (EPI_V)
+      tRG_rAux_frg(epi_v) = convert_input(frg_input);
+
+      return frg_input;
+    }
+
+    MUTLASS_DEVICE void
+    end_step(int epi_m, int epi_n, int step_idx) {
+      auto src_v = tRG_rAux;
+      auto dst_v = tRG_gAux(_, _, _, epi_m, epi_n);
+      auto coord_v = tRG_cAux(_, _, _, epi_m, epi_n);
+
+      if (get<0>(residue_mn) >= get<0>(tile_shape) &&
+          get<1>(residue_mn) >= get<1>(tile_shape)) {
+        copy(tiled_r2g.with(params_ptr->desc_aux), src_v, dst_v);
+      } else {
+        MUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < size(dst_v); ++i) {
+          if (elem_less(coord_v(i), residue_mn)) {
+            dst_v(i) = src_v(i);
+          }
+        }
+      }
+    }
+  };
+
+  template <class... Args>
+  MUTLASS_DEVICE auto
+  get_callbacks(VisitorArgs<Args...> const& args) {
+    auto [M, N, K, L] = args.problem_shape_mnkl;
+    auto [m, n, k, l] = args.tile_coord_mnkl;
+
+    Tensor mAux = make_tensor(
+        make_gmem_ptr(params_ptr->ptr_aux),
+        make_shape(M, N, L),
+        params_ptr->dAux); // (M, N, L)
+
+    Tensor gAux = local_tile(mAux, take<0, 2>(args.tile_shape_mnk), make_coord(m, n, l)); // (BLK_M, BLK_N)
+    Tensor gAux_epi = flat_divide(gAux, args.epi_tile); // (EPI_M, EPI_N, TILE_M, TILE_N)
+    Tensor cAux = make_identity_tensor(make_shape(size<0>(gAux), size<1>(gAux)));
+    Tensor cAux_epi = flat_divide(cAux, args.epi_tile);
+
+    auto thread_copy = args.tiled_copy.get_thread_slice(args.thread_idx);
+    Tensor tDgAux    = thread_copy.partition_D(gAux_epi);
+    constexpr int vec_elem = decltype(max_common_vector(tDgAux, args.tDgD))::value;
+    constexpr int vec_bits = vec_elem * sizeof_bits_v<Element>;
+    using CopyAtomR2G      = Copy_Atom<
+            UniversalCopy<uint_bit_t<mute::min(1024, vec_bits)>>, Element>;
+    using TiledCopyR2G = TiledCopy<CopyAtomR2G,
+                            typename decltype(args.tiled_copy)::TiledLayout_TV,
+                            typename decltype(args.tiled_copy)::Tiler_MN>;
+
+    TiledCopyR2G tiled_copy_r2g{};
+    auto thr_copy_r2g = tiled_copy_r2g.get_slice(args.thread_idx);
+    auto tRG_gAux     = thr_copy_r2g.partition_D(gAux_epi);
+    auto tRG_rAux     = thr_copy_r2g.retile_S(make_tensor<Element>(take<0, 3>(shape(tRG_gAux))));
+    auto tRG_cAux     = thr_copy_r2g.partition_D(cAux_epi);
+
+    return Callbacks<
+      decltype(tRG_gAux), decltype(tRG_rAux), decltype(tRG_cAux),
+      TiledCopyR2G, decltype(args.residue_mn), decltype(args.tile_shape_mnk)>(
+          mute::move(tRG_gAux),
+          mute::move(tRG_rAux),
+          mute::move(tRG_cAux),
+          tiled_copy_r2g,
+          args.residue_mn,
+          args.tile_shape_mnk,
+          params_ptr);
+  }
+};
+
+} // namespace mutlass::epilogue::fusion
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/operations.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/operations.hpp
new file mode 100644
index 000000000..8d6189345
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/operations.hpp
@@ -0,0 +1,320 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <mutlass/numeric_conversion.h>
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::fusion {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Fusion Operations
+// Template args must not be implementation dependent
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct FusionOperation {
+  // metadata types/queries that can be overrided
+  using ElementOutput = void;
+  using ElementCompute = void;
+
+  using ElementSource = void;
+  static constexpr bool IsSourceSupported = false;
+
+  using ElementScalar = void;
+  static constexpr int AlignmentScalar = 0;
+  static constexpr bool IsScaleFactorSupported = false;
+  static constexpr bool IsPerRowScaleSupported = false;
+  using ElementBias = void;
+  static constexpr int AlignmentBias = 0;
+  static constexpr bool IsPerRowBiasSupported = false;
+  static constexpr bool IsDePerRowBiasSupported = false;
+
+  using ActivationFn = void;
+  static constexpr bool IsEltActSupported = false;
+  static constexpr bool IsDeEltActSupported = false;
+
+  using ElementAux = void;
+  using GmemLayoutTagAux = void;
+  static constexpr int AlignmentAux = 0;
+  static constexpr bool IsAuxOutSupported = false;
+  static constexpr bool IsAuxInSupported = false;
+
+  using ElementAmax = void;
+  static constexpr bool IsAbsMaxSupported = false;
+
+};
+
+// D = alpha * acc
+template<
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementScalar_ = ElementCompute_,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct ScaledAcc : FusionOperation {
+  using ElementOutput = ElementOutput_;
+  using ElementCompute = ElementCompute_;
+  using ElementScalar = ElementScalar_;
+  static constexpr int AlignmentScalar = 1;
+  static constexpr auto RoundStyle = RoundStyle_;
+};
+
+// D = alpha * acc + beta * C
+template<
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinearCombination
+    : ScaledAcc<ElementOutput_, ElementCompute_, ElementScalar_, RoundStyle_> {
+  using ElementSource = ElementSource_;
+  static constexpr bool IsSourceSupported = true;
+};
+
+// D = activation(alpha * acc + beta * C)
+template<
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinCombEltAct
+    : LinearCombination<ElementOutput_, ElementCompute_, ElementSource_, ElementScalar_, RoundStyle_> {
+  using ActivationFn = ActivationFn_<ElementCompute_>;
+  static constexpr bool IsEltActSupported = true;
+};
+
+
+// D = alpha * acc + beta * C + per-row bias
+template<
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentBias_ = 128 / mute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinCombPerRowBias
+    : LinearCombination<ElementOutput_, ElementCompute_, ElementSource_, ElementScalar_, RoundStyle_> {
+  using ElementBias = ElementBias_;
+  static constexpr int AlignmentBias = AlignmentBias_;
+  static constexpr bool IsPerRowBiasSupported = true;
+};
+
+// D = activation(alpha * acc + beta * C + per-row bias)
+template<
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentBias_ = 128 / mute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinCombPerRowBiasEltAct
+    : LinCombPerRowBias<ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  using ActivationFn = ActivationFn_<ElementCompute_>;
+  static constexpr bool IsEltActSupported = true;
+};
+
+// D = activation(alpha * acc + beta * C + per-row bias)
+// aux = alpha * acc + beta * C + per-row bias
+template<
+  class GmemLayoutTagAux_,
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementAux_ = ElementOutput_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentAux_ = 128 / mute::sizeof_bits_v<ElementAux_>,
+  int AlignmentBias_ = 128 / mute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinCombPerRowBiasEltActAux
+    : LinCombPerRowBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  using ElementAux = ElementAux_;
+  using GmemLayoutTagAux = GmemLayoutTagAux_;
+  static constexpr int AlignmentAux = AlignmentAux_;
+  static constexpr bool IsAuxOutSupported = true;
+};
+
+// D = activation(per-row alpha * acc + per-row beta * C + per-row bias)
+template<
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_, // per-row alpha/beta
+  int AlignmentBias_ = 128 / mute::sizeof_bits_v<ElementBias_>,
+  int AlignmentScalar_ = 128 / mute::sizeof_bits_v<ElementScalar_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct PerRowLinCombPerRowBiasEltAct
+    : LinCombPerRowBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  static constexpr int AlignmentScalar = AlignmentScalar_;
+  static constexpr bool IsPerRowScaleSupported = true;
+};
+
+// Z = scale_a * scale_b * alpha * acc + beta * scale_c * C + per-row bias
+// if D is fp8
+//   D = scale_d * activation(Z)
+// else
+//   D = activation(Z)
+template<
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentBias_ = 128 / mute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct ScaledLinCombPerRowBiasEltAct
+    : LinCombPerRowBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  static constexpr bool IsScaleFactorSupported = true;
+};
+
+// Z = scale_a * scale_b * alpha * acc + scale_c * beta * C + per-row bias
+// if D is fp8
+//   amax_d = max(abs(elements in activation(Z)))
+//   D = scale_d * activation(Z)
+// else
+//   D = activation(Z)
+// if Aux is fp8
+//   amax_aux = max(abs(elements in Z))
+//   Aux = scale_aux * Z
+// else
+//   Aux = Z
+template<
+  class GmemLayoutTagAux_,
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementAux_ = ElementOutput_,
+  class ElementAmax_ = ElementCompute_,
+  class ElementBias_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentAux_ = 128 / mute::sizeof_bits_v<ElementAux_>,
+  int AlignmentBias_ = 128 / mute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct ScaledLinCombPerRowBiasEltActAmaxAux
+    : ScaledLinCombPerRowBiasEltAct<ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementBias_, ElementSource_, ElementScalar_, AlignmentBias_, RoundStyle_> {
+  using ElementAmax = ElementAmax_;
+  static constexpr bool IsAbsMaxSupported = true;
+
+  using ElementAux = ElementAux_;
+  using GmemLayoutTagAux = GmemLayoutTagAux_;
+  static constexpr int AlignmentAux = AlignmentAux_;
+  static constexpr bool IsAuxOutSupported = true;
+};
+
+// Z = Aux
+// dY = alpha * acc + beta * C
+// D = d_activation(dY, Z)
+template<
+  class GmemLayoutTagAux_,
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementAux_ = ElementOutput_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentAux_ = 128 / mute::sizeof_bits_v<ElementAux_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinCombDeEltAct
+    : LinearCombination<ElementOutput_, ElementCompute_, ElementSource_, ElementScalar_, RoundStyle_> {
+  using ActivationFn = ActivationFn_<ElementCompute_>;
+  static constexpr bool IsDeEltActSupported = true;
+
+  using ElementAux = ElementAux_;
+  using GmemLayoutTagAux = GmemLayoutTagAux_;
+  static constexpr int AlignmentAux = AlignmentAux_;
+  static constexpr bool IsAuxInSupported = true;
+};
+
+// Z = Aux
+// dY = alpha * acc + beta * C
+// D = d_activation(dY, Z)
+// dBias = sum of columns of D
+template<
+  class GmemLayoutTagAux_,
+  template <class> class ActivationFn_,
+  class ElementOutput_,
+  class ElementCompute_,
+  class ElementAux_ = ElementOutput_,
+  class ElementBias_ = ElementCompute_,
+  class ElementSource_ = ElementOutput_,
+  class ElementScalar_ = ElementCompute_,
+  int AlignmentAux_ = 128 / mute::sizeof_bits_v<ElementAux_>,
+  int AlignmentBias_ = 128 / mute::sizeof_bits_v<ElementBias_>,
+  FloatRoundStyle RoundStyle_ = FloatRoundStyle::round_to_nearest
+>
+struct LinCombDeEltActDePerRowBias
+    : LinCombDeEltAct<GmemLayoutTagAux_, ActivationFn_, ElementOutput_, ElementCompute_,
+        ElementAux_, ElementSource_, ElementScalar_, AlignmentAux_, RoundStyle_> {
+  using ElementBias = ElementBias_;
+  static constexpr int AlignmentBias = AlignmentBias_;
+  static constexpr bool IsDePerRowBiasSupported = true;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::fusion
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/visitor_base.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/visitor_base.hpp
new file mode 100644
index 000000000..effdd24e7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/fusion/visitor_base.hpp
@@ -0,0 +1,477 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/workspace.h"
+
+#include "mute/tensor.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::epilogue::fusion::detail {
+
+using namespace mute;
+using mute::tuple;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Visitor Implementation
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class... Ops>
+struct VisitorImplBase {
+  // Shared memory allocation
+  using SharedStorage = tuple<typename Ops::SharedStorage...>;
+  // Host side fusion arguments
+  using Arguments = tuple<typename Ops::Arguments...>;
+  // Device side fusion params (Kernel-entry API)
+  using Params = tuple<typename Ops::Params...>;
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    return transform_apply(tuple<Ops...>{}, args,
+      [&] (auto&& op, auto const& op_args) {
+        using Op = mute::remove_cvref_t<decltype(op)>;
+        return Op::to_underlying_arguments(problem_shape, op_args, workspace);
+      },
+      [] (auto&&... op_params) { return mute::make_tuple(op_params...); }
+    );
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    return transform_apply(tuple<Ops...>{}, args,
+      [&] (auto&& op, auto const& op_args) {
+        using Op = mute::remove_cvref_t<decltype(op)>;
+        size_t op_workspace_size = Op::get_workspace_size(problem_shape, op_args);
+        return round_nearest(op_workspace_size, MinWorkspaceAlignment);
+      },
+      [&] (auto&&... op_workspace_size) {
+        return (0 + ... + op_workspace_size);
+      }
+    );
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream) {
+    Status status = Status::kSuccess;
+    uint8_t* op_workspace = reinterpret_cast<uint8_t*>(workspace);
+    return transform_apply(tuple<Ops...>{}, args,
+      // Initialize each operation's workspace, stopping at the first error
+      [&] (auto&& op, auto const& op_args) {
+        if (status != Status::kSuccess) {
+          return status;
+        }
+
+        using Op = mute::remove_cvref_t<decltype(op)>;
+        status = Op::initialize_workspace(problem_shape, op_args, op_workspace, stream);
+        if (op_workspace != nullptr) {
+          size_t op_workspace_size = Op::get_workspace_size(problem_shape, op_args);
+          op_workspace += round_nearest(op_workspace_size, MinWorkspaceAlignment);
+        }
+        return status;
+      },
+      // Return the final status
+      [&] (auto const&...) { return status; }
+    );
+  }
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase() {}
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
+    : ops(transform_apply(tuple<Ops...>{}, params, shared_storage,
+        [] (auto&& op, auto const& op_params, auto&& op_storage) {
+          using Op = mute::remove_cvref_t<decltype(op)>;
+          return Op(op_params, op_storage);
+        },
+        [] (auto&&... ops) { return mute::make_tuple(ops...); }
+      )) {}
+
+  // Ops can store kernel persistent variables (e.g. descriptors, scalars, wave counters)
+  tuple<Ops...> ops;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Base specializations so we can have standard layout params and simple aggregate initializers
+template <class Op0>
+struct VisitorImplBase<Op0> {
+
+  // Retain tuple for SharedStorage because empty structs have 1B alignment
+  // tuples use multiple inheritance, avoids this problem
+  using SharedStorage = tuple<
+    typename Op0::SharedStorage
+  >;
+
+  struct Arguments {
+    typename Op0::Arguments op_0;
+  };
+
+  struct Params {
+    typename Op0::Params op_0;
+  };
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    return Params{
+      Op0::to_underlying_arguments(problem_shape, args.op_0, workspace)
+    };
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    size_t workspace_size = 0;
+    workspace_size += Op0::get_workspace_size(problem_shape, args.op_0);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    return workspace_size;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream) {
+    Status status = Status::kSuccess;
+    uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
+    size_t workspace_offset = 0;
+
+    status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op0::get_workspace_size(problem_shape, args.op_0);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return status;
+  }
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase() {}
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
+    : ops({
+        Op0(params.op_0, get<0>(shared_storage))
+      }) {}
+
+  tuple<Op0> ops;
+};
+
+template <class Op0, class Op1>
+struct VisitorImplBase<Op0, Op1> {
+
+  using SharedStorage = tuple<
+    typename Op0::SharedStorage,
+    typename Op1::SharedStorage
+  >;
+
+  struct Arguments {
+    typename Op0::Arguments op_0;
+    typename Op1::Arguments op_1;
+  };
+
+  struct Params {
+    typename Op0::Params op_0;
+    typename Op1::Params op_1;
+  };
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    return Params{
+      Op0::to_underlying_arguments(problem_shape, args.op_0, workspace),
+      Op1::to_underlying_arguments(problem_shape, args.op_1, workspace)
+    };
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    size_t workspace_size = 0;
+    workspace_size += Op0::get_workspace_size(problem_shape, args.op_0);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    workspace_size += Op1::get_workspace_size(problem_shape, args.op_1);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    return workspace_size;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream) {
+    Status status = Status::kSuccess;
+    uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
+    size_t workspace_offset = 0;
+
+    status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op0::get_workspace_size(problem_shape, args.op_0);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op1::get_workspace_size(problem_shape, args.op_1);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return status;
+  }
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase() {}
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
+    : ops({
+        Op0(params.op_0, get<0>(shared_storage)),
+        Op1(params.op_1, get<1>(shared_storage))
+      }) {}
+
+  tuple<Op0, Op1> ops;
+};
+
+template <class Op0, class Op1, class Op2>
+struct VisitorImplBase<Op0, Op1, Op2> {
+
+  using SharedStorage = tuple<
+    typename Op0::SharedStorage,
+    typename Op1::SharedStorage,
+    typename Op2::SharedStorage
+  >;
+
+  struct Arguments {
+    typename Op0::Arguments op_0;
+    typename Op1::Arguments op_1;
+    typename Op2::Arguments op_2;
+  };
+
+  struct Params {
+    typename Op0::Params op_0;
+    typename Op1::Params op_1;
+    typename Op2::Params op_2;
+  };
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    return Params{
+      Op0::to_underlying_arguments(problem_shape, args.op_0, workspace),
+      Op1::to_underlying_arguments(problem_shape, args.op_1, workspace),
+      Op2::to_underlying_arguments(problem_shape, args.op_2, workspace)
+    };
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    size_t workspace_size = 0;
+    workspace_size += Op0::get_workspace_size(problem_shape, args.op_0);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    workspace_size += Op1::get_workspace_size(problem_shape, args.op_1);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    workspace_size += Op2::get_workspace_size(problem_shape, args.op_2);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    return workspace_size;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream) {
+    Status status = Status::kSuccess;
+    uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
+    size_t workspace_offset = 0;
+
+    status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op0::get_workspace_size(problem_shape, args.op_0);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op1::get_workspace_size(problem_shape, args.op_1);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = Op2::initialize_workspace(problem_shape, args.op_2, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op2::get_workspace_size(problem_shape, args.op_2);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return status;
+  }
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase() {}
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
+    : ops({
+        Op0(params.op_0, get<0>(shared_storage)),
+        Op1(params.op_1, get<1>(shared_storage)),
+        Op2(params.op_2, get<2>(shared_storage))
+      }) {}
+
+  tuple<Op0, Op1, Op2> ops;
+};
+
+template <class Op0, class Op1, class Op2, class Op3>
+struct VisitorImplBase<Op0, Op1, Op2, Op3> {
+
+  using SharedStorage = tuple<
+    typename Op0::SharedStorage,
+    typename Op1::SharedStorage,
+    typename Op2::SharedStorage,
+    typename Op3::SharedStorage
+  >;
+
+  struct Arguments {
+    typename Op0::Arguments op_0;
+    typename Op1::Arguments op_1;
+    typename Op2::Arguments op_2;
+    typename Op3::Arguments op_3;
+  };
+
+  struct Params {
+    typename Op0::Params op_0;
+    typename Op1::Params op_1;
+    typename Op2::Params op_2;
+    typename Op3::Params op_3;
+  };
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+    return Params{
+      Op0::to_underlying_arguments(problem_shape, args.op_0, workspace),
+      Op1::to_underlying_arguments(problem_shape, args.op_1, workspace),
+      Op2::to_underlying_arguments(problem_shape, args.op_2, workspace),
+      Op3::to_underlying_arguments(problem_shape, args.op_3, workspace)
+    };
+  }
+
+  template <class ProblemShape>
+  static size_t
+  get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
+    size_t workspace_size = 0;
+    workspace_size += Op0::get_workspace_size(problem_shape, args.op_0);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    workspace_size += Op1::get_workspace_size(problem_shape, args.op_1);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    workspace_size += Op2::get_workspace_size(problem_shape, args.op_2);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    workspace_size += Op3::get_workspace_size(problem_shape, args.op_3);
+    workspace_size = round_nearest(workspace_size, MinWorkspaceAlignment);
+
+    return workspace_size;
+  }
+
+  template <class ProblemShape>
+  static mutlass::Status
+  initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, musaStream_t stream) {
+    Status status = Status::kSuccess;
+    uint8_t* workspace_ptr = reinterpret_cast<uint8_t*>(workspace);
+    size_t workspace_offset = 0;
+
+    status = Op0::initialize_workspace(problem_shape, args.op_0, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op0::get_workspace_size(problem_shape, args.op_0);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = Op1::initialize_workspace(problem_shape, args.op_1, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op1::get_workspace_size(problem_shape, args.op_1);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = Op2::initialize_workspace(problem_shape, args.op_2, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op2::get_workspace_size(problem_shape, args.op_2);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    status = Op3::initialize_workspace(problem_shape, args.op_3, workspace_ptr + workspace_offset, stream);
+    workspace_offset += Op3::get_workspace_size(problem_shape, args.op_3);
+    workspace_offset = round_nearest(workspace_offset, MinWorkspaceAlignment);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+
+    return status;
+  }
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase() {}
+
+  MUTLASS_HOST_DEVICE
+  VisitorImplBase(Params const& params, SharedStorage const& shared_storage)
+    : ops({
+        Op0(params.op_0, get<0>(shared_storage)),
+        Op1(params.op_1, get<1>(shared_storage)),
+        Op2(params.op_2, get<2>(shared_storage)),
+        Op3(params.op_3, get<3>(shared_storage))
+      }) {}
+
+  tuple<Op0, Op1, Op2, Op3> ops;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::epilogue::fusion::detail
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/activation.h b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/activation.h
new file mode 100644
index 000000000..f38ae1005
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/activation.h
@@ -0,0 +1,791 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 This extends the contents of mutlass/functional.h with frequently used activation functions.
+
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/constants.h"
+#include "mutlass/complex.h"
+#include "mutlass/array.h"
+#include "mutlass/half.h"
+#include "mutlass/functional.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace thread {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Identity operator
+template <typename T>
+struct Identity {
+  static const bool kIsHeavy = false;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T value) const {
+    return value;
+  }
+};
+
+template <typename T, int N>
+struct Identity<Array<T, N> > {
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> value) const {
+    return value;
+  }
+};
+
+/// Scale operator
+template <typename T>
+struct Scale {
+  struct Arguments {
+    using scale_type = T;
+    T scale = T(1);
+  };
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T value, T scale) const {
+    multiplies<T> mul;
+    return mul(scale, value);
+  }
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T value, Arguments args = Arguments()) const {
+    return this->operator()(value, args.scale);
+  }
+};
+
+template <typename T, int N>
+struct Scale<Array<T, N>> {
+  using Arguments = typename Scale<T>::Arguments;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> values, T scale) const {
+    multiplies<Array<T, N>> mul;
+    return mul(scale, values);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> values, Arguments args = Arguments()) const {
+    return this->operator()(values, args.scale);
+  }
+};
+
+/// Specialization to compose other activations with a defined unary operator
+/// e.g. Scale<Identity<T>>
+template <template <class> class Activation, typename T>
+struct Scale<Activation<T>> {
+  using Arguments = typename Scale<T>::Arguments;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T value, typename Arguments::scale_type scale) const {
+    multiplies<T> mul;
+    Activation<T> act;
+    return mul(scale, act(value));
+  }
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T value, Arguments args = Arguments()) const {
+    return this->operator()(value, args.scale);
+  }
+};
+
+/// ReLu operator - propagates NaNs
+/// Always put threshold in the right hand side of max to propagate NaN.
+template <typename T>
+struct ReLu {
+  static const bool kIsHeavy = false;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T threshold, T value) const {
+    maximum<T> mx;
+
+    return mx(value, threshold);
+  }
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T value) const {
+    maximum<T> mx;
+
+    return mx(value, T(0));
+  }
+};
+
+template <typename T>
+using ReLU = ReLu<T>;
+
+template <typename T, int N>
+struct ReLu<Array<T, N>> {
+  static const bool kIsHeavy = false;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(T const & threshold, Array<T, N> const &frag) const {
+    maximum<Array<T, N>> mx;
+
+    return mx(frag, threshold);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &frag) const {
+    maximum<Array<T, N>> mx;
+    return mx(frag, T(0));
+  }
+};
+
+// Generic clamp
+template <typename T>
+struct Clamp {
+  struct Arguments {
+    T lower_bound = MUTLASS_STL_NAMESPACE::numeric_limits<T>::lowest();
+    T upper_bound = MUTLASS_STL_NAMESPACE::numeric_limits<T>::max();
+  };
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const& value, T const& lower_bound, T const& upper_bound) const {
+    constexpr bool PropagateNaN = true;
+    maximum<T, PropagateNaN> mx;
+    minimum<T, PropagateNaN> mn;
+
+    return mn(mx(value, lower_bound), upper_bound);
+  }
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const& value, Arguments const& args = Arguments()) const {
+    return this->operator()(value, args.lower_bound, args.upper_bound);
+  }
+};
+
+template <typename T, int N>
+struct Clamp<Array<T,N>> {
+  using Arguments = typename Clamp<T>::Arguments;
+
+  MUTLASS_HOST_DEVICE
+  Array<T,N> operator()(Array<T,N> const& values, T const& lower_bound, T const& upper_bound) const {
+    constexpr bool PropagateNaN = true;
+    maximum<Array<T,N>, PropagateNaN> mx;
+    minimum<Array<T,N>, PropagateNaN> mn;
+
+    return mn(mx(values, lower_bound), upper_bound);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T,N> operator()(Array<T,N> const& values, Arguments const& args = Arguments()) const {
+    return this->operator()(values, args.lower_bound, args.upper_bound);
+  }
+};
+
+// Leaky Relu operator
+template <typename T>
+struct LeakyReLU {
+
+  static const bool kIsHeavy = false;
+
+  struct Arguments {
+    T leaky_alpha = T(0);
+  };
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const& value, T const& leaky_alpha) const {
+    T res = value > T(0) ? value : value * leaky_alpha;
+    return res;
+  }
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const& value, Arguments const& args = Arguments()) const {
+    return this->operator()(value, args.leaky_alpha);
+  }
+};
+
+template <typename T, int N>
+struct LeakyReLU<Array<T, N> > {
+
+  static const bool kIsHeavy = false;
+
+  using Arguments = typename LeakyReLU<T>::Arguments;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& values, T const& leaky_alpha) const {
+    Array<T, N> y;
+    LeakyReLU<T> leaky_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < int(values.size()); ++i) {
+      y[i] = leaky_op(values[i], leaky_alpha);
+    }
+
+    return y;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& values, Arguments const& args = Arguments()) const {
+    return this->operator()(values, args.leaky_alpha);
+  }
+};
+
+// Y = min((X <= threshold ? 0 : X), upper_bound)
+template <typename T>
+struct ThresholdReLU {
+  static constexpr bool kIsHeavy = false;
+
+  struct Arguments {
+    T threshold = T(0);
+    T upper_bound = MUTLASS_STL_NAMESPACE::numeric_limits<T>::max();
+  };
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T value, T threshold, T upper_bound) const {
+    minimum_with_nan_propagation<T> mn;
+
+    return mn((value <= threshold ? T(0) : value), upper_bound);
+  }
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T value, Arguments const& args = Arguments()) const {
+    return operator()(value, args.threshold, args.upper_bound);
+  }
+};
+
+template <typename T, int N>
+struct ThresholdReLU<Array<T,N>> {
+  static constexpr bool kIsHeavy = false;
+
+  using Arguments = typename ThresholdReLU<T>::Arguments;
+
+  MUTLASS_HOST_DEVICE
+  Array<T,N> operator()(Array<T,N> const& values, T threshold, T upper_bound) const {
+    ThresholdReLU<T> relu;
+
+    Array<T,N> retvals;
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      retvals[i] = relu(values[i], threshold, upper_bound);
+    }
+
+    return retvals;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T,N> operator()(Array<T,N> const& values, Arguments const& args = Arguments()) const {
+    return operator()(values, args.threshold, args.upper_bound);
+  }
+};
+
+// Tanh operator
+template <typename T>
+struct Tanh {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &value) const {
+    return fast_tanh(value);
+  }
+};
+
+template <typename T, int N>
+struct Tanh<Array<T, N> > {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &value) const {
+    Array<T, N> y;
+    Tanh<T> tanh_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = tanh_op(value[i]);
+    }
+
+    return y;
+  }
+};
+
+template <int N>
+struct Tanh<Array<half_t, N>> {
+  using T = half_t;
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& z) const {
+    fast_tanh_op<Array<T, N>> tanh;
+    return tanh(z);
+  }
+};
+
+// Sigmoid operator
+template <typename T>
+struct Sigmoid {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &value) const {
+    return T(1) / (T(1) + fast_exp(-value));
+  }
+};
+
+template <typename T, int N>
+struct Sigmoid<Array<T, N>> {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& z) const {
+    plus<Array<T, N>> add;
+
+#if defined(MUTLASS_USE_TANH_FOR_SIGMOID)
+    multiplies<Array<T, N>> mul;
+    fast_tanh_op<Array<T, N>> tanh;
+    return mul(add(tanh(mul(z, mutlass::constants::half<T>())), mutlass::constants::one<T>()),
+               mutlass::constants::half<T>());
+#else
+    divides<Array<T, N>> div;
+    negate<Array<T, N>> neg;
+    fast_exp_op<Array<T, N>> fast_exp;
+    return div(mutlass::constants::one<T>(),
+               add(mutlass::constants::one<T>(),
+                   fast_exp(neg(z))));
+#endif
+  }
+};
+
+// SiLu (swish) operator introduced by Elfwing et al. in the following paper
+// "Sigmoid-Weighted Linear Units for Neural Network Function Approximation in Reinforcement Learning" (2017)
+// https://arxiv.org/pdf/1702.03118.pdf
+// It is used in EfficientNet and YOLOv5, for example.
+// Reference: https://pytorch.org/docs/stable/generated/torch.nn.SiLU.html
+template <typename T>
+struct SiLu {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &value) const {
+    Sigmoid<T> sigmoid;
+    return value * sigmoid(value);
+  }
+};
+
+template <typename T, int N>
+struct SiLu<Array<T, N>> {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &value) const {
+    Sigmoid<Array<T, N>> sigmoid_op;
+    multiplies<Array<T, N>>     mul;
+    return mul(value, sigmoid_op(value));
+  }
+};
+
+template <typename T>
+using ScaledSiLu = Scale<SiLu<T>>;
+
+// Hardswish operator introduced by Howard et al. in the following paper
+// "Searching for MobileNetV3" (2019)
+// https://arxiv.org/pdf/1905.02244.pdf
+// It is used in models based on MobilenetNetV3.
+// Reference: https://pytorch.org/docs/stable/generated/torch.nn.Hardswish.html
+template <typename T>
+struct HardSwish {
+  static const bool kIsHeavy = false;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &x) const {
+    minimum<T> mn;
+    maximum<T> mx;
+    T relu6 = mn(mx(x + T(3), T(0)), T(6));
+    return x * relu6 / T(6);
+  }
+};
+
+template <>
+struct HardSwish<float> {
+  using T = float;
+  static const bool kIsHeavy = false;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &x) const {
+    minimum<T> mn;
+    maximum<T> mx;
+    T relu6 = mn(mx(x + T(3), T(0)), T(6));
+    return x * relu6 * 0.16666667f;
+  }
+};
+
+template <typename T, int N>
+struct HardSwish<Array<T, N> > {
+  static const bool kIsHeavy = false;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &value) const {
+    Array<T, N> y;
+    HardSwish<T> hardswish_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = hardswish_op(value[i]);
+    }
+
+    return y;
+  }
+};
+
+template <int N>
+struct HardSwish<Array<half_t, N> > {
+  using T = half_t;
+  static const bool kIsHeavy = false;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &value) const {
+    minimum<Array<T, N> > mn;
+    maximum<Array<T, N> > mx;
+    multiplies<Array<T, N> > mul;
+    plus<Array<T, N> > add;
+
+    return mul(mul(mn(mx(add(value, T(3)), T(0)), T(6)), value), T(0.16666667f));
+  }
+};
+
+template <typename T>
+using ScaledHardSwish = Scale<HardSwish<T>>;
+
+//
+// GELU function definitions implemented as described by
+//   Hendrycks, D., and Gimpel, K. in
+//   "Gaussian Error Linear Units (GELUs)." (2020)
+//   https://arxiv.org/pdf/1606.08415.pdf
+//
+// Floating-point constants are Taylor coefficients described in the paper.
+//
+
+// GELU operator
+template <typename T>
+struct GELU {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &value) const {
+    return T(mutlass::constants::half<T>() * value *
+      (mutlass::constants::one<T>() + (T)erff((float)(value * mutlass::constants::half_root_two<T>()))));
+  }
+};
+
+template <>
+struct GELU<float> {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  float operator()(float const &value) const {
+    return mutlass::constants::half<float>() * value *
+      (mutlass::constants::one<float>() + erff(value * mutlass::constants::half_root_two<float>() ));
+  }
+};
+
+template <>
+struct GELU<double> {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  double operator()(double const &value) const {
+    return mutlass::constants::half<double>() * value *
+      (mutlass::constants::one<double>() + erf( value * mutlass::constants::half_root_two<double>() ));
+  }
+};
+
+template <typename T, int N>
+struct GELU<Array<T, N> > {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &value) const {
+    Array<T, N> y;
+    GELU<T> gelu_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = gelu_op(value[i]);
+    }
+
+    return y;
+  }
+};
+
+template <typename T>
+using ScaledGELU = Scale<GELU<T>>;
+
+// GELU operator implemented using the Taylor series approximation
+template <typename T>
+struct GELU_taylor {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &z) const {
+
+    T k0 = T(0.7978845608028654);
+    T k1 = T(0.044715);
+
+    return T(mutlass::constants::half<T>() * z *
+      (mutlass::constants::one<T>() + fast_tanh(k0 * z * (mutlass::constants::one<T>() + k1 * z * z))));
+  }
+};
+
+template <int N>
+struct GELU_taylor<Array<half_t, N> > {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const &z) const {
+
+    using T = half_t;
+    Array<half_t, N> y;
+
+    half_t k0 = half_t(0.7978845608028654);
+    half_t k1 = half_t(0.044715);
+
+    multiply_add<Array<half_t, N>> fma;
+    multiplies<Array<half_t, N>>     mul;
+    plus<Array<half_t, N>>         add;
+
+    fast_tanh_op<Array<half_t, N>> tanh;
+
+    Array<half_t, N> u = mul(mul(k0, z), fma(mul(k1, z), z, mutlass::constants::one<T>()));
+
+    y = mul(mul(z, mutlass::constants::half<T>()), add(mutlass::constants::one<T>(), tanh(u)));
+
+    return y;
+  }
+};
+
+template <typename T, int N>
+struct GELU_taylor<Array<T, N> > {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &value) const {
+    Array<T, N> y;
+    GELU_taylor<T> gelu_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = gelu_op(value[i]);
+    }
+
+    return y;
+  }
+};
+
+template <typename T>
+using ScaledGELU_taylor = Scale<GELU_taylor<T>>;
+
+/// Computes backwards pass for GELU operator assuming d_t is the layer gradient and
+/// z is computed from the forward pass.
+template <typename T>
+struct dGELU {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &d_t, T const &z) const {
+
+    T k0 = T(0.7978845608028654);
+    T k1 = T(0.044715);
+    T k2 = T(0.1070322243);
+
+    T tanh_out = fast_tanh(k0 * z * (1 + k1 * z * z));
+
+    T ff = constants::half<T>() * z * ((1 - tanh_out * tanh_out) * (k0 + k2 * z * z)) +
+      constants::half<T>() * (1 + tanh_out);
+
+    return ff * d_t;
+  }
+};
+
+template <typename T, int N>
+struct dGELU<Array<T, N> > {
+  static const bool kIsHeavy = true;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &d_t, Array<T, N> const &z) const {
+    Array<T, N> y;
+    dGELU<T> gelu_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = gelu_op(d_t[i], z[i]);
+    }
+
+    return y;
+  }
+};
+
+template <typename T>
+struct dReLU {
+  MUTLASS_HOST_DEVICE
+  T operator()(T d_t, bool d_relu) const {
+    return d_relu ? d_t : T(0);
+  }
+
+  template <typename U>
+  MUTLASS_HOST_DEVICE
+  T operator()(T d_t, U d_relu) const {
+    return operator()(d_t, static_cast<bool>(d_relu));
+  }
+};
+
+template <typename T, int N>
+struct dReLU<Array<T, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& d_t, bool const (&d_relu)[N]) const {
+    Array<T, N> y;
+    dReLU<T> relu_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = relu_op(d_t[i], d_relu[i]);
+    }
+
+    return y;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& d_t, Array<uint1b_t, N> const& d_relu) const {
+    UnpackPredicates<N> unpack_op;
+
+    bool preds[N];
+    unpack_op(preds, d_relu);
+
+    return operator()(d_t, preds);
+  }
+
+  template <typename U>
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& d_t, Array<U, N> const& d_relu) const {
+    Array<T, N> y;
+    dReLU<T> relu_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = relu_op(d_t[i], d_relu[i]);
+    }
+
+    return y;
+  }
+};
+
+/// Computes backwards pass for ReLU operator assuming d_t is the layer gradient and
+/// z is computed from the forward pass.
+template <typename T>
+struct dReLU_Z {
+  MUTLASS_HOST_DEVICE
+  T operator()(T d_t, T z) const {
+    return z < 0 ? T(0) : d_t;
+  }
+};
+
+template <typename T, int N>
+struct dReLU_Z<Array<T, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& d_t, Array<T, N> const& z) const {
+    Array<T, N> y;
+    dReLU_Z<T> relu_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = relu_op(d_t[i], z[i]);
+    }
+
+    return y;
+  }
+};
+
+// ElementwiseFilter operator
+// Filters by a specific value and maps it to 0.0
+// Used in GEMM + comm
+template <typename T>
+struct ElementwiseFilter {
+
+  static const bool kIsHeavy = false;
+
+  struct Arguments {
+    T value_to_filter = T(-0.0);
+    T filtered_value = T(0.0);
+  };
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const& value, T const& value_to_filter, T const& filtered_value) const {
+    T res = value == value_to_filter ? filtered_value : value;
+    return res;
+  }
+
+  MUTLASS_HOST_DEVICE
+  T operator()(T const& value, Arguments const& args = Arguments()) const {
+    return this->operator()(value, args.value_to_filter, args.filtered_value);
+  }
+};
+
+template <typename T, int N>
+struct ElementwiseFilter<Array<T, N> > {
+
+  static const bool kIsHeavy = false;
+
+  using Arguments = typename ElementwiseFilter<T>::Arguments;
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& values, T const& value_to_filter, T const& filtered_value) const {
+    Array<T, N> y;
+    ElementwiseFilter<T> filter_op;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < int(values.size()); ++i) {
+      y[i] = filter_op(values[i], value_to_filter, filtered_value);
+    }
+
+    return y;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const& values, Arguments const& args = Arguments()) const {
+    return this->operator()(values, args.value_to_filter, args.filtered_value);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace thread
+} // namespace epilogue
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/linear_combination.h b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/linear_combination.h
new file mode 100644
index 000000000..47924e269
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/linear_combination.h
@@ -0,0 +1,524 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Functor performing linear combination operations used by epilogues.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/array.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/epilogue/thread/scale_type.h"
+#include "mutlass/epilogue/thread/linear_combination_params.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace thread {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Applies a linear combination operator to an array of elements.
+///
+/// D = alpha * accumulator + beta * source
+///
+template <
+  typename ElementOutput_,                             ///< Data type used to load and store tensors
+  int Count,                                           ///< Number of elements computed per operation.
+                                                       ///< Usually it is 128/sizeof_bits<ElementOutput_>,
+                                                       ///< but we use 64 or 32 sometimes when there are not enough data to store
+  typename ElementAccumulator_ = ElementOutput_,       ///< Accumulator data type
+  typename ElementCompute_ = ElementOutput_,           ///< Data type used to compute linear combination
+  ScaleType::Kind Scale = ScaleType::Default,          ///< Control Alpha and Beta scaling
+  FloatRoundStyle Round = FloatRoundStyle::round_to_nearest,
+  typename ElementSource_ = ElementOutput_
+>
+class LinearCombination {
+public:
+
+  using ElementOutput = ElementOutput_;
+  using ElementSource = ElementSource_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
+  using ElementScalar = ElementCompute;
+  using ElementC = ElementSource_;
+  using ElementD = ElementOutput_;
+
+  static int const kCount = Count;
+  static const ScaleType::Kind kScale = Scale;
+  using FragmentOutput = Array<ElementOutput, kCount>;
+  using FragmentSource = Array<ElementSource, kCount>;
+  using FragmentAccumulator = Array<ElementAccumulator, kCount>;
+  using FragmentCompute = Array<ElementCompute, kCount>;
+
+  static FloatRoundStyle const kRound = Round;
+
+  /// Host-constructable parameters structure
+  struct Params
+  {
+    ElementCompute alpha;                         ///< scales accumulators
+    ElementCompute beta;                          ///< scales source tensor
+    ElementCompute const *alpha_ptr;              ///< pointer to accumulator scalar - if not null, loads it from memory
+    ElementCompute const *beta_ptr;               ///< pointer to source scalar - if not null, loads it from memory
+    ElementCompute const* const* alpha_ptr_array; ///< array of pointers to accumulator scalar per group/batch
+    ElementCompute const* const* beta_ptr_array;  ///< array of pointers to source scalar per group/batch
+
+    MUTLASS_HOST_DEVICE
+    Params():
+      alpha(ElementCompute(1)),
+      beta(ElementCompute(0)),
+      alpha_ptr(nullptr),
+      beta_ptr(nullptr),
+      alpha_ptr_array(nullptr),
+      beta_ptr_array(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute alpha,
+      ElementCompute beta
+    ):
+      alpha(alpha), beta(beta),
+      alpha_ptr(nullptr), beta_ptr(nullptr),
+      alpha_ptr_array(nullptr), beta_ptr_array(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute alpha
+    ):
+      alpha(alpha), beta(0),
+      alpha_ptr(nullptr), beta_ptr(nullptr),
+      alpha_ptr_array(nullptr), beta_ptr_array(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute const *alpha_ptr,
+      ElementCompute const *beta_ptr
+    ):
+      alpha(0), beta(0),
+      alpha_ptr(alpha_ptr), beta_ptr(beta_ptr),
+      alpha_ptr_array(nullptr), beta_ptr_array(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute const *alpha_ptr
+    ):
+      alpha(0), beta(0),
+      alpha_ptr(alpha_ptr), beta_ptr(nullptr),
+      alpha_ptr_array(nullptr), beta_ptr_array(nullptr) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute const* const* alpha_ptr_array,
+      ElementCompute const* const* beta_ptr_array
+    ):
+      alpha(0), beta(0),
+      alpha_ptr(nullptr), beta_ptr(nullptr),
+      alpha_ptr_array(alpha_ptr_array), beta_ptr_array(beta_ptr_array) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute const* const* alpha_ptr_array
+    ):
+      alpha(0), beta(0),
+      alpha_ptr(nullptr), beta_ptr(nullptr),
+      alpha_ptr_array(alpha_ptr_array), beta_ptr_array(nullptr) { }
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  ElementCompute alpha_;
+  ElementCompute beta_;
+
+public:
+
+  /// Constructs the function object, possibly loading from pointers in host memory
+  MUTLASS_HOST_DEVICE
+  LinearCombination(Params const &params, int group_idx = 0) {
+    if (params.alpha_ptr_array != nullptr && params.alpha_ptr_array[group_idx] != nullptr) {
+      alpha_ = *(params.alpha_ptr_array[group_idx]);
+    }
+    else if (params.alpha_ptr != nullptr) {
+      alpha_ = *params.alpha_ptr;
+    }
+    else {
+      alpha_ = params.alpha;
+    }
+    if (params.beta_ptr_array != nullptr && params.beta_ptr_array[group_idx] != nullptr) {
+      beta_ = *(params.beta_ptr_array[group_idx]);
+    }
+    else if (params.beta_ptr != nullptr) {
+      beta_ = *params.beta_ptr;
+    }
+    else {
+      beta_ = params.beta;
+    }
+  }
+
+  /// Returns true if source is needed
+  MUTLASS_HOST_DEVICE
+  bool is_source_needed() const {
+    if (Scale == ScaleType::NoBetaScaling) return true;
+
+    if (Scale == ScaleType::OnlyAlphaScaling) return false;
+
+    if (Scale == ScaleType::Nothing) return false;
+
+    return beta_ != ElementCompute(0);
+  }
+
+  /// Functionally required for serial reduction in the epilogue
+  MUTLASS_HOST_DEVICE
+  void set_k_partition(int k_partition, int k_partition_count) {
+    if (k_partition) {
+      beta_ = ElementCompute(1);
+    }
+  }
+
+  /// Computes linear scaling with source: D = alpha * accumulator + beta * source
+  MUTLASS_HOST_DEVICE
+  FragmentOutput operator()(
+      FragmentAccumulator const &accumulator,
+      FragmentSource const &source) const {
+
+    // Convert source to internal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
+    NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
+
+    // Convert to destination numeric type
+    NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
+
+    FragmentCompute converted_source = source_converter(source);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    if (Scale == ScaleType::Nothing)
+      return destination_converter(converted_accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    multiplies<FragmentCompute> mul_add_source;
+    multiply_add<FragmentCompute> mul_add_accumulator;
+
+    if (Scale == ScaleType::NoBetaScaling)
+      intermediate = converted_source;
+    else
+      intermediate = mul_add_source(beta_, converted_source);                             // X =  beta * C + uniform
+
+    intermediate = mul_add_accumulator(alpha_, converted_accumulator, intermediate);    // D = alpha * Accum + X
+
+    return destination_converter(intermediate);
+  }
+
+  /// Computes linear scaling: D = alpha * accumulator
+  MUTLASS_HOST_DEVICE
+  FragmentOutput operator()(
+      FragmentAccumulator const &accumulator) const {
+
+    // Convert source to interal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
+
+    // Convert to destination numeric type
+    NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
+
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    if (Scale == ScaleType::Nothing)
+      return destination_converter(converted_accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+    multiplies<FragmentCompute> mul_accumulator;
+
+    intermediate = mul_accumulator(alpha_, converted_accumulator);    // D = alpha * Accum
+
+    return destination_converter(intermediate);
+  }
+
+  //
+  // Specializations for scalar (for use with mute::collective::DefaultEpilogue)
+  //
+  MUTLASS_HOST_DEVICE
+  ElementD operator()(ElementAccumulator const accumulator, ElementC const source) const {
+    // Convert everything to Compute type, do compute, and then store to output type
+    NumericConverter<ElementCompute, ElementAccumulator, Round> accumulator_converter;
+    [[maybe_unused]] NumericConverter<ElementCompute, ElementC, Round> source_converter;
+    NumericConverter<ElementD, ElementCompute, Round> destination_converter;
+
+    // Convert to destination numeric type
+
+    ElementCompute converted_accumulator = accumulator_converter(accumulator);
+    if constexpr (Scale == ScaleType::Nothing) {
+      return destination_converter(converted_accumulator);
+    }
+
+    // Perform binary operations
+    ElementCompute intermediate;
+    multiplies<ElementCompute> multiply;
+    multiply_add<ElementCompute> madd;
+
+    if constexpr (Scale == ScaleType::NoBetaScaling) {
+      intermediate = source_converter(source);
+    }
+    else {
+      intermediate = multiply(beta_, source);                            // X =  beta * C + uniform
+    }
+
+    intermediate = madd(alpha_, converted_accumulator, intermediate);    // D = alpha * Accum + X
+    return destination_converter(intermediate);
+  }
+
+  MUTLASS_HOST_DEVICE
+  ElementD operator()(ElementAccumulator const accumulator) const {
+    // Convert everything to Compute type, do compute, and then store to output type
+    NumericConverter<ElementCompute, ElementAccumulator, Round> accumulator_converter;
+    NumericConverter<ElementD, ElementCompute, Round> destination_converter;
+    ElementCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Convert to destination numeric type
+    if constexpr (Scale == ScaleType::Nothing) {
+      return destination_converter(converted_accumulator);
+    }
+
+    // Perform binary operations
+    ElementCompute intermediate;
+    multiplies<ElementCompute> multiply;
+
+    intermediate = multiply(alpha_, accumulator);    // D = alpha * Accum
+    return destination_converter(intermediate);
+  }
+};
+
+/// Applies a linear combination operator to an array of elements.
+///
+/// D = vector_alpha * accumulator + (optional) vector_beta/scalar_beta * source
+///
+template <
+  typename ElementOutput_,            ///< Data type used to load and store tensors
+  int Count,                          ///< Number of elements computed per operation.
+  typename ElementAccumulator_,       ///< Accumulator data type
+  typename ElementCompute_,           ///< Data type used to compute linear combination
+  FloatRoundStyle Round,
+  typename ElementSource_
+>
+class LinearCombination<ElementOutput_,
+                        Count,
+                        ElementAccumulator_,
+                        ElementCompute_,
+                        ScaleType::PerChannelScaling,
+                        Round,
+                        ElementSource_> {
+public:
+
+  using ElementOutput = ElementOutput_;
+  using ElementSource = ElementSource_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
+  using ElementC = ElementSource_;
+  using ElementD = ElementOutput_;
+
+  static int const kCount = Count;
+  static const ScaleType::Kind kScale = ScaleType::PerChannelScaling;
+  static constexpr bool IsPerChannelScalingSupported = true;
+
+  using FragmentOutput = Array<ElementOutput, kCount>;
+  using FragmentSource = Array<ElementSource, kCount>;
+  using FragmentAccumulator = Array<ElementAccumulator, kCount>;
+  using FragmentCompute = Array<ElementCompute, kCount>;
+
+  static FloatRoundStyle const kRound = Round;
+
+  /// Host-constructable parameters structure
+  struct Params
+  {
+    ElementCompute const *alpha_ptr;       ///< pointer to accumulator vector
+    ElementCompute const *beta_ptr;        ///< pointer to source vector
+    ElementCompute beta;                   ///< scales source tensor
+
+    MUTLASS_HOST_DEVICE
+    Params():
+      alpha_ptr(nullptr),
+      beta_ptr(nullptr),
+      beta(ElementCompute(0)) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute const *alpha_ptr,
+      ElementCompute const *beta_ptr
+    ):
+      alpha_ptr(alpha_ptr), beta_ptr(beta_ptr), beta(ElementCompute(0)) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute const *alpha_ptr
+    ):
+      alpha_ptr(alpha_ptr), beta_ptr(nullptr), beta(ElementCompute(0)) { }
+
+    MUTLASS_HOST_DEVICE
+    Params(
+      ElementCompute const *alpha_ptr,
+      ElementCompute beta
+    ):
+      alpha_ptr(alpha_ptr), beta_ptr(nullptr), beta(beta) { }
+
+  };
+
+private:
+
+  //
+  // Data members
+  //
+
+  ElementCompute const* beta_ptr_ = nullptr;
+  ElementCompute beta_ = 0;
+
+public:
+
+  /// Constructs the function object
+  MUTLASS_HOST_DEVICE
+  LinearCombination(Params const& params) {
+    if (params.beta_ptr) {
+      beta_ptr_ = params.beta_ptr;
+    }
+    else {
+      beta_ = params.beta;
+    }
+  }
+
+  /// Returns true if source is needed
+  MUTLASS_HOST_DEVICE
+  bool is_source_needed() const {
+    return beta_ptr_ != nullptr || beta_ != ElementCompute(0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  bool is_beta_vector() const {
+    return beta_ptr_ != nullptr;
+  }
+
+  /// Computes linear scaling with source: D = vector_alpha * accumulator + vector_beta * source
+  MUTLASS_HOST_DEVICE
+  FragmentOutput operator()(
+      FragmentAccumulator const& accumulator,
+      FragmentSource const& source,
+      FragmentCompute const& valpha,
+      FragmentCompute const& vbeta) const {
+    // Convert source to internal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
+    NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
+
+    // Convert to destination numeric type
+    NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
+
+    FragmentCompute converted_source = source_converter(source);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    multiplies<FragmentCompute> mul_add_source;
+    multiply_add<FragmentCompute> mul_add_accumulator;
+
+    intermediate = mul_add_source(vbeta, converted_source);                             // X = vector_beta * C + uniform
+
+    intermediate = mul_add_accumulator(valpha, converted_accumulator, intermediate);    // D = vector_alpha * Accum + X
+
+    return destination_converter(intermediate);
+  }
+
+  /// Computes linear scaling with source: D = vector_alpha * accumulator + scalar_beta(from host) * source
+  MUTLASS_HOST_DEVICE
+  FragmentOutput operator()(
+      FragmentAccumulator const& accumulator,
+      FragmentSource const& source,
+      FragmentCompute const& valpha) const {
+    // Convert source to internal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementSource, kCount, Round> source_converter;
+    NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
+
+    // Convert to destination numeric type
+    NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
+
+    FragmentCompute converted_source = source_converter(source);
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+
+    multiplies<FragmentCompute> mul_add_source;
+    multiply_add<FragmentCompute> mul_add_accumulator;
+
+
+    intermediate = mul_add_source(beta_, converted_source);                           // X =  scalar_beta * C + uniform
+
+    intermediate = mul_add_accumulator(valpha, converted_accumulator, intermediate);    // D = vector_alpha * Accum + X
+
+    return destination_converter(intermediate);
+  }
+
+  /// Computes linear scaling: D = vector_alpha * accumulator
+  MUTLASS_HOST_DEVICE
+  FragmentOutput operator()(
+      FragmentAccumulator const& accumulator,
+      FragmentCompute const& valpha) const {
+    // Convert source to interal compute numeric type
+    NumericArrayConverter<ElementCompute, ElementAccumulator, kCount, Round> accumulator_converter;
+
+    // Convert to destination numeric type
+    NumericArrayConverter<ElementOutput, ElementCompute, kCount, Round> destination_converter;
+
+    FragmentCompute converted_accumulator = accumulator_converter(accumulator);
+
+    // Perform binary operations
+    FragmentCompute intermediate;
+    multiplies<FragmentCompute> mul_accumulator;
+
+    intermediate = mul_accumulator(valpha, converted_accumulator);    // D = vector_alpha * Accum
+
+    return destination_converter(intermediate);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace thread
+} // namespace epilogue
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/linear_combination_params.h b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/linear_combination_params.h
new file mode 100644
index 000000000..dcb50013a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/linear_combination_params.h
@@ -0,0 +1,76 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace thread {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct LinearCombinationParams {
+  uint64_t alpha_data[2];
+  uint64_t beta_data[2];
+
+  MUTLASS_HOST_DEVICE
+  LinearCombinationParams()
+  : alpha_data {0lu, 0lu}, beta_data {0lu, 0lu}
+  { }
+
+  template <typename ElementCompute>
+  MUTLASS_HOST_DEVICE
+  LinearCombinationParams(ElementCompute alpha, ElementCompute beta)
+  : alpha_data {0lu, 0lu}, beta_data {0lu, 0lu}
+  {
+#if defined(__MUSA_ARCH__)
+    reinterpret_cast<ElementCompute&>(alpha_data) = alpha;
+    reinterpret_cast<ElementCompute&>(beta_data) = beta;
+#else
+    memcpy( alpha_data, &alpha, sizeof(ElementCompute) );
+    memcpy( beta_data, &beta, sizeof(ElementCompute) );
+#endif
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace thread
+} // namespace epilogue
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/scale_type.h b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/scale_type.h
new file mode 100644
index 000000000..790194ef6
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/epilogue/thread/scale_type.h
@@ -0,0 +1,67 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Enum defines the behaviors of the epilogue.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace epilogue {
+namespace thread {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Specifies internal data type for computation
+/// Note :
+///  1. Scalar means alpha/beta is a single value from host(constant param) or device memory.
+///  2. Vector means alpha/beta is a vector always from device memory.
+struct ScaleType {
+  enum Kind {
+    Default,                           // D = scalar_alpha x Acc + scalar_beta x C
+    NoBetaScaling,                     // D = scalar_alpha x Acc + C
+    OnlyAlphaScaling,                  // D = scalar_alpha x Acc
+    PerChannelScaling,                 // D = vector_alpha x Acc + vector_beta x C
+    OnlyAlphaPerChannelScaling,        // D = vector_alpha x Acc
+    Nothing                            // D = Acc
+  };
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace thread
+} // namespace epilogue
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/fast_math.h b/vllm_musa/csrc/flash_attn/include/mutlass/fast_math.h
new file mode 100644
index 000000000..b0412cb89
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/fast_math.h
@@ -0,0 +1,1028 @@
+/***************************************************************************************************
+* Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <cstdint>
+#include <cmath>
+#include <type_traits>
+#endif
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+#include "mutlass/uint128.h"
+#include "mutlass/coord.h"
+#include "mutlass/half.h"
+
+/**
+ * \file
+ * \brief Math utilities
+ */
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+MUTLASS_HOST_DEVICE void swap(T &lhs, T &rhs) {
+  T tmp = lhs;
+  lhs = rhs;
+  rhs = tmp;
+}
+
+/******************************************************************************
+ * Static math utilities
+ ******************************************************************************/
+
+/// Mixed precision dot product
+template <typename Index, typename LongIndex, int N>
+MUTLASS_HOST_DEVICE LongIndex dot(
+  Coord<N, Index> const &coord,
+  Coord<N, LongIndex> const &stride,
+  LongIndex acc = LongIndex()) {
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int n = 0; n < N; ++n) {
+    acc += LongIndex(coord[n]) * stride[n];
+  }
+  return acc;
+}
+
+/**
+ * Statically determine if N is a power-of-two
+ */
+template <int N>
+struct is_pow2 {
+  static bool const value = ((N & (N - 1)) == 0);
+};
+
+/**
+ * Statically determine log2(N), rounded down
+ */
+template <int N, int CurrentVal = N, int Count = 0>
+struct log2_down {
+  /// Static logarithm value
+  enum { value = log2_down<N, (CurrentVal >> 1), Count + 1>::value };
+};
+
+// Base case
+template <int N, int Count>
+struct log2_down<N, 1, Count> {
+  enum { value = Count };
+};
+
+/**
+ * Statically determine log2(N), rounded up
+ */
+template <int N, int CurrentVal = N, int Count = 0>
+struct log2_up {
+  /// Static logarithm value
+  enum { value = log2_up<N, (CurrentVal >> 1), Count + 1>::value };
+};
+
+// Base case
+template <int N, int Count>
+struct log2_up<N, 1, Count> {
+  enum { value = ((1 << Count) < N) ? Count + 1 : Count };
+};
+
+/**
+ * Statically estimate sqrt(N) to the nearest power-of-two
+ */
+template <int N>
+struct sqrt_est {
+  enum { value = 1 << (log2_up<N>::value / 2) };
+};
+
+/**
+ * For performing a constant-division with a compile-time assertion that the
+ * Divisor evenly-divides the Dividend.
+ */
+template <int Dividend, int Divisor>
+struct divide_assert {
+  enum { value = Dividend / Divisor };
+
+  static_assert((Dividend % Divisor == 0), "Not an even multiple");
+};
+
+/******************************************************************************
+ * Rounding
+ ******************************************************************************/
+
+/**
+ * Round dividend up to the nearest multiple of divisor
+ */
+template <typename dividend_t, typename divisor_t>
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+dividend_t round_nearest(dividend_t dividend, divisor_t divisor) {
+  return ((dividend + divisor - 1) / divisor) * divisor;
+}
+
+template <typename value_t>
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+value_t abs_for_integer(value_t a) {
+  return ((a > 0) ? a : -a);
+}
+/**
+ * Greatest common divisor
+ */
+template <typename value_t>
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+value_t gcd(value_t a, value_t b) {
+  for (;;) {
+    if (a == 0) return mutlass::abs_for_integer(b);
+    b %= a;
+    if (b == 0) return mutlass::abs_for_integer(a);
+    a %= b;
+  }
+}
+
+/**
+ * Least common multiple
+ */
+template <typename value_t>
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+value_t lcm(value_t a, value_t b) {
+  value_t temp = mutlass::gcd(a, b);
+  return (temp != 0) ? value_t(mutlass::abs_for_integer(a) / temp * mutlass::abs_for_integer(b)) : value_t{};
+}
+
+/**
+ * Greatest common divisor
+ */
+template <typename value_t>
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+value_t gcd_cxx11(value_t a, value_t b) {
+  return (a == 0 || b == 0) ? mutlass::abs_for_integer(a | b) : mutlass::gcd_cxx11(b, a % b);
+}
+
+/**
+ * Least common multiple
+ */
+template <typename value_t>
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+value_t lcm_cxx11(value_t a, value_t b) {
+  return mutlass::gcd_cxx11(a, b) ? (mutlass::abs_for_integer(a) / mutlass::gcd_cxx11(a, b) *
+                                    mutlass::abs_for_integer(b))
+                                  : value_t{};
+}
+
+/// Returns the smallest value in the half-open range [a, a+b) that is a multiple of b
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+int round_up(int a, int b) {
+  return ((a + b - 1) / b) * b;
+}
+
+/// Returns the ceiling of (a / b)
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+int ceil_div(int a, int b) {
+  return (a + b - 1) / b;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/**
+ * log2 computation, what's the
+ * difference between the below codes and
+ * log2_up/down codes?
+ */
+template <typename value_t>
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+value_t clz(value_t x) {
+  for (int i = 31; i >= 0; --i) {
+    if ((1 << i) & x)
+      return value_t(31 - i);
+  }
+  return value_t(32);
+}
+
+template <typename value_t>
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+value_t find_log2(value_t x) {
+  int a = int(31 - clz(x));
+  a += (x & (x - 1)) != 0;  // Round up, add 1 if not a power of 2.
+  return a;
+}
+
+
+/**
+ * Find divisor, using find_log2
+ */
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+void find_divisor(unsigned int& mul, unsigned int& shr, unsigned int denom) {
+  if (denom == 1) {
+    mul = 0;
+    shr = 0;
+  } else {
+    unsigned int p = 31 + find_log2(denom);
+    unsigned m = unsigned(((1ull << p) + unsigned(denom) - 1) / unsigned(denom));
+
+    mul = m;
+    shr = p - 32;
+  }
+}
+
+/**
+ * Find quotient and remainder using device-side intrinsics
+ */
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+void fast_divmod(int& quo, int& rem, int src, int div, unsigned int mul, unsigned int shr) {
+
+  #if defined(__MUSA_ARCH__)
+  // Use IMUL.HI if div != 1, else simply copy the source.
+  quo = (div != 1) ? __umulhi(src, mul) >> shr : src;
+  #else
+  quo = int((div != 1) ? int(((int64_t)src * mul) >> 32) >> shr : src);
+  #endif
+
+  // The remainder.
+  rem = src - (quo * div);
+}
+
+// For long int input
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+void fast_divmod(int& quo, int64_t& rem, int64_t src, int div, unsigned int mul, unsigned int shr) {
+
+  #if defined(__MUSA_ARCH__)
+  // Use IMUL.HI if div != 1, else simply copy the source.
+  quo = (div != 1) ? __umulhi(src, mul) >> shr : src;
+  #else
+  quo = int((div != 1) ? ((src * mul) >> 32) >> shr : src);
+  #endif
+  // The remainder.
+  rem = src - (quo * div);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Object to encapsulate the fast division+modulus operation.
+///
+/// This object precomputes two values used to accelerate the computation and is best used
+/// when the divisor is a grid-invariant. In this case, it may be computed in host code and
+/// marshalled along other kernel arguments using the 'Params' pattern.
+///
+/// Example:
+///
+///
+///   int quotient, remainder, dividend, divisor;
+///
+///   FastDivmod divmod(divisor);
+///
+///   divmod(quotient, remainder, dividend);
+///
+///   // quotient = (dividend / divisor)
+///   // remainder = (dividend % divisor)
+///
+struct FastDivmod {
+  using value_div_type = int;
+  using value_mod_type = int64_t;
+  int32_t divisor = 1;
+  uint32_t multiplier = 0u;
+  uint32_t shift_right = 0u;
+
+  // Find quotient and remainder using device-side intrinsics
+  MUTLASS_HOST_DEVICE
+  void fast_divmod(int& quotient, int& remainder, int dividend) const {
+
+#if defined(__MUSA_ARCH__)
+    // Use IMUL.HI if divisor != 1, else simply copy the source.
+    quotient = (divisor != 1) ? __umulhi(dividend, multiplier) >> shift_right : dividend;
+#else
+    quotient = int((divisor != 1) ? int(((int64_t)dividend * multiplier) >> 32) >> shift_right : dividend);
+#endif
+
+    // The remainder.
+    remainder = dividend - (quotient * divisor);
+  }
+
+  /// For long int input
+  MUTLASS_HOST_DEVICE
+  void fast_divmod(int& quotient, int64_t& remainder, int64_t dividend) const {
+
+#if defined(__MUSA_ARCH__)
+    // Use IMUL.HI if divisor != 1, else simply copy the source.
+    quotient = (divisor != 1) ? __umulhi(dividend, multiplier) >> shift_right : dividend;
+#else
+    quotient = int((divisor != 1) ? ((dividend * multiplier) >> 32) >> shift_right : dividend);
+#endif
+    // The remainder.
+    remainder = dividend - (quotient * divisor);
+  }
+
+
+  /// Construct the FastDivmod object, in host code ideally.
+  ///
+  /// This precomputes some values based on the divisor and is computationally expensive.
+
+  constexpr FastDivmod() = default;
+
+  MUTLASS_HOST_DEVICE
+  FastDivmod(int divisor_): divisor(divisor_) {
+    assert(divisor_ >= 0);
+    if (divisor != 1) {
+      unsigned int p = 31 + find_log2(divisor);
+      unsigned m = unsigned(((1ull << p) + unsigned(divisor) - 1) / unsigned(divisor));
+
+      multiplier = m;
+      shift_right = p - 32;
+    }
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  MUTLASS_HOST_DEVICE
+  void operator()(int &quotient, int &remainder, int dividend) const {
+    fast_divmod(quotient, remainder, dividend);
+  }
+
+  /// Computes integer division using precomputed values. This is computationally
+  /// inexpensive.
+  MUTLASS_HOST_DEVICE
+  int div(int dividend) const {
+    int quotient, remainder;
+    fast_divmod(quotient, remainder, dividend);
+    return quotient;
+  }
+
+  /// Alias for `div` to match the interface of FastDivmodU64
+  MUTLASS_HOST_DEVICE
+  int divide(int dividend) const {
+    return div(dividend);
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  ///
+  /// Simply returns the quotient
+  MUTLASS_HOST_DEVICE
+  int divmod(int &remainder, int dividend) const {
+    int quotient;
+    fast_divmod(quotient, remainder, dividend);
+    return quotient;
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  MUTLASS_HOST_DEVICE
+  void operator()(int &quotient, int64_t &remainder, int64_t dividend) const {
+    fast_divmod(quotient, remainder, dividend);
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  MUTLASS_HOST_DEVICE
+  int divmod(int64_t &remainder, int64_t dividend) const {
+    int quotient;
+    fast_divmod(quotient, remainder, dividend);
+    return quotient;
+  }
+
+  /// Returns the divisor when cast to integer
+  MUTLASS_HOST_DEVICE
+  operator int() const { return divisor; }
+
+};
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Object to encapsulate the fast division+modulus operation for 64b integer division.
+///
+/// This object precomputes two values used to accelerate the computation and is best used
+/// when the divisor is a grid-invariant. In this case, it may be computed in host code and
+/// marshalled along other kernel arguments using the 'Params' pattern.
+///
+/// Example:
+///
+///
+///   uint64_t quotient, remainder, dividend, divisor;
+///
+///   FastDivmodU64 divmod(divisor);
+///
+///   divmod(quotient, remainder, dividend);
+///
+///   // quotient = (dividend / divisor)
+///   // remainder = (dividend % divisor)
+///
+struct FastDivmodU64 {
+
+  uint64_t divisor;
+  uint64_t multiplier;
+  unsigned int shift_right;
+  unsigned int round_up;
+
+  //
+  // Static methods
+  //
+
+  /// Computes b, where 2^b is the greatest power of two that is less than or equal to x
+  MUTLASS_HOST_DEVICE
+  static uint32_t integer_log2(uint64_t x) {
+    uint32_t n = 0;
+    while (x >>= 1) {
+      ++n;
+    }
+    return n;
+  }
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  FastDivmodU64(): divisor(0), multiplier(0), shift_right(0), round_up(0) { }
+
+  /// Construct the FastDivmod object, in host code ideally.
+  ///
+  /// This precomputes some values based on the divisor and is computationally expensive.
+  MUTLASS_HOST_DEVICE
+  FastDivmodU64(uint64_t divisor_): divisor(divisor_), multiplier(1), shift_right(0), round_up(0) {
+
+    if (divisor) {
+      shift_right = integer_log2(divisor);
+
+      if ((divisor & (divisor - 1)) == 0) {
+        multiplier = 0;
+      }
+      else {
+        uint64_t power_of_two = (uint64_t(1) << shift_right);
+        uint64_t multiplier_lo = uint128_t(0, power_of_two) / divisor;
+        multiplier = uint128_t(power_of_two, power_of_two) / divisor;
+        round_up = (multiplier_lo == multiplier ? 1 : 0);
+      }
+    }
+  }
+
+  /// Returns the quotient of floor(dividend / divisor)
+  MUTLASS_HOST_DEVICE
+  uint64_t divide(uint64_t dividend) const {
+    uint64_t quotient = 0;
+
+    #ifdef __MUSA_ARCH__
+      uint64_t x = dividend;
+      if (multiplier) {
+        x = __umul64hi(dividend + round_up, multiplier);
+      }
+      quotient = (x >> shift_right);
+    #else
+      quotient = dividend / divisor;
+    #endif
+
+    return quotient;
+  }
+
+  /// Computes the remainder given a computed quotient and dividend
+  MUTLASS_HOST_DEVICE
+  uint64_t modulus(uint64_t quotient, uint64_t dividend) const {
+    return dividend - quotient * divisor;
+  }
+
+  /// Returns the quotient of floor(dividend / divisor) and computes the remainder
+  MUTLASS_HOST_DEVICE
+  uint64_t divmod(uint64_t &remainder, uint64_t dividend) const {
+    uint64_t quotient = divide(dividend);
+    remainder = modulus(quotient, dividend);
+    return quotient;
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  MUTLASS_HOST_DEVICE
+  void operator()(uint64_t &quotient, uint64_t &remainder, uint64_t dividend) const {
+    quotient = divmod(remainder, dividend);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Object to encapsulate the fast division+modulus operation for 64b integer division
+/// in which the divisor is a power of two.
+struct FastDivmodU64Pow2 {
+
+  uint64_t divisor;
+  unsigned int shift_right;
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  FastDivmodU64Pow2(): divisor(0), shift_right(0) { }
+
+  /// Construct the FastDivmod object, in host code ideally.
+  ///
+  /// This precomputes some values based on the divisor and is computationally expensive.
+  MUTLASS_HOST_DEVICE
+  FastDivmodU64Pow2(uint64_t divisor_): divisor(divisor_), shift_right(FastDivmodU64::integer_log2(divisor_)) { }
+
+  /// Returns the quotient of floor(dividend / divisor)
+  MUTLASS_HOST_DEVICE
+  uint64_t divide(uint64_t dividend) const {
+    return dividend >> shift_right;
+  }
+
+  /// Computes the remainder given a computed quotient and dividend
+  MUTLASS_HOST_DEVICE
+  uint64_t modulus(uint64_t dividend) const {
+    // See https://docs.nvidia.com/musa/musa-c-best-practices-guide/index.html#division-modulo-operations
+    return dividend & (divisor - 1);
+  }
+
+  /// Returns the quotient of floor(dividend / divisor) and computes the remainder
+  MUTLASS_HOST_DEVICE
+  uint64_t divmod(uint64_t &remainder, uint64_t dividend) const {
+    uint64_t quotient = divide(dividend);
+    remainder = modulus(dividend);
+    return quotient;
+  }
+
+  /// Computes integer division and modulus using precomputed values. This is computationally
+  /// inexpensive.
+  MUTLASS_HOST_DEVICE
+  void operator()(uint64_t &quotient, uint64_t &remainder, uint64_t dividend) const {
+    quotient = divmod(remainder, dividend);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes the coordinate decomposition from a linear index (64-bit linear index => coord<int32_t>)
+///
+/// This decomposition is accelerated by the FastDivmodU64 object. It is assumed that
+/// a coordinate of <Rank> indices can be decomposed by <Rank - 1> div/mod operations.
+/// Note, is assumed that element divmod[0] divides by extent[1].
+///
+/// For example, assume 4-D coordinate (n, p, q, c) is mapped to a linear index `npqc`. This
+/// can be decomposed via three divide and modulus operations:
+///
+///      c = npqc % C;         |  divmod[2] = FastDivmodU64(C)
+///    npq = npqc / C;         |   coord[3] = c
+///
+///      q =  npq % Q;         |  divmod[1] = FastDivmodU64(Q)
+///     np =  npq / Q;         |   coord[2] = q
+///
+///      p =   np % P;         |  divmod[0] = FastDivmodU64(P)
+///      n =   np / P;         |   coord[1] = p
+///
+///                            |   coord[0] = n
+///
+template <int Rank>
+MUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecomposition(
+  uint64_t linear_idx,                    ///< Linear index to decompose
+  FastDivmodU64 const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = Rank; i > 1; --i) {
+    uint64_t remainder;
+    linear_idx = divmod[i - 2].divmod(remainder, linear_idx);
+    coord[i - 1] = int(remainder);
+  }
+
+  coord[0] = int(linear_idx);
+
+  return coord;
+}
+
+/// Computes the coordinate decomposition from a linear index (32-bit linear index => coord<int32_t>)
+template <int Rank>
+MUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecomposition(
+  int linear_idx,                    ///< Linear index to decompose
+  FastDivmod const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = Rank; i > 1; --i) {
+    int remainder;
+    linear_idx = divmod[i - 2].divmod(remainder, linear_idx);
+    coord[i - 1] = int(remainder);
+  }
+
+  coord[0] = int(linear_idx);
+
+  return coord;
+}
+
+template <int Rank>
+MUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecompositionLittleEndian(
+  uint64_t linear_idx,                    ///< Linear index to decompose
+  FastDivmodU64 const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank - 1; ++i) {
+    uint64_t remainder;
+    linear_idx = divmod[i].divmod(remainder, linear_idx);
+    coord[i] = int(remainder);
+  }
+
+  coord[Rank - 1] = int(linear_idx);
+
+  return coord;
+}
+
+/// Computes the coordinate decomposition from a linear index (32-bit linear index => coord<int32_t>)
+template <int Rank>
+MUTLASS_HOST_DEVICE Coord<Rank> CoordinateDecompositionLittleEndian(
+  int linear_idx,                    ///< Linear index to decompose
+  FastDivmod const *divmod) {          ///< Pointer to array of Rank-1 FastDivmodU64 objects
+
+  static_assert(Rank > 0, "CoordinateDecomposition requires Rank=1 or greater.");
+
+  Coord<Rank> coord;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank - 1; ++i) {
+    int remainder;
+    linear_idx = divmod[i].divmod(remainder, linear_idx);
+    coord[i] = int(remainder);
+  }
+
+  coord[Rank - 1] = int(linear_idx);
+
+  return coord;
+}
+
+/// Safely computes the offset of a linear index in bytes for all types
+template <typename Element>
+MUTLASS_HOST_DEVICE int64_t OffsetBytes(int64_t index) {
+
+  static_assert(
+    (sizeof_bits<Element>::value >= 8 && !(sizeof_bits<Element>::value % 8)) ||
+    (sizeof_bits<Element>::value <  8 && !(8 % sizeof_bits<Element>::value)),
+    "Size of numeric type in bits must either be divisible by 8 bits, or 8 bits must be divisible by the size.");
+
+  if (sizeof_bits<Element>::value >= 8) {
+    return index * (sizeof_bits<Element>::value / 8);
+  }
+  else {
+    int const kElementsPerByte = ((8 / sizeof_bits<Element>::value) + ((sizeof_bits<Element>::value >= 8) ? 1 : 0));
+    return index / kElementsPerByte;
+  }
+}
+
+MUTLASS_HOST_DEVICE int64_t OffsetBytes(int64_t index, int64_t element_sizeof_bits) {
+  if (element_sizeof_bits >= 8) {
+    return index * (element_sizeof_bits / 8);
+  }
+  else {
+    int64_t const kElementsPerByte = ((8 / element_sizeof_bits) + ((element_sizeof_bits >= 8) ? 1 : 0));
+    return index / kElementsPerByte;
+  }
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// Min/Max
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int A, int B>
+struct Min {
+  static int const kValue = (A < B) ? A : B;
+};
+
+template <int A, int B>
+struct Max {
+  static int const kValue = (A > B) ? A : B;
+};
+
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17 int const_min(int a, int b) {
+    return (b < a ? b : a);
+}
+
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17 int const_max(int a, int b) {
+    return (b > a ? b : a);
+}
+
+template <typename T>
+MUTLASS_HOST_DEVICE
+T fast_min(T a, T b) {
+  return (b < a ? b : a);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+float fast_min(float a, float b) {
+  return fminf(a, b);
+}
+
+template <typename T>
+MUTLASS_HOST_DEVICE
+T fast_max(T a, T b) {
+  return (a < b ? b : a);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+float fast_max(float a, float b) {
+  return fmaxf(a, b);
+}
+
+MUTLASS_HOST_DEVICE
+float fast_cos(float theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::cosf(theta);
+  #else
+  return std::cos(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+double fast_cos(double theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::cos(theta);
+  #else
+  return std::cos(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+float fast_sin(float theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::sinf(theta);
+  #else
+  return std::sin(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+double fast_sin(double theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::sin(theta);
+  #else
+  return std::sin(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+float fast_acos(float theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::acosf(theta);
+  #else
+  return std::acos(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+double fast_acos(double theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::acos(theta);
+  #else
+  return std::acos(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+float fast_asin(float theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::asinf(theta);
+  #else
+  return std::asin(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+double fast_asin(double theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::asin(theta);
+  #else
+  return std::asin(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+float fast_sqrt(float theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::sqrtf(theta);
+  #else
+  return std::sqrt(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+double fast_sqrt(double theta) {
+  #if defined(__MUSA_ARCH__)
+  return ::sqrt(theta);
+  #else
+  return std::sqrt(theta);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+float fast_exp(float x) {
+  #if defined(__MUSA_ARCH__)
+  return ::expf(x);
+  #else
+  return std::exp(x);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+double fast_exp(double x) {
+  #if defined(__MUSA_ARCH__)
+  return ::exp(x);
+  #else
+  return std::exp(x);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t fast_exp(half_t x) {
+  #if defined(__MUSACC__) && (_cplusplus)
+      return (half_t)(::hexp(x.to_half()));
+  #else
+      return (half_t)(fast_exp(float(x)));
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+float fast_log(float x) {
+  #if defined(__MUSA_ARCH__)
+  return ::logf(x);
+  #else
+  return std::log(x);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+double fast_log(double x) {
+  #if defined(__MUSA_ARCH__)
+  return ::log(x);
+  #else
+  return std::log(x);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+float fast_tanh(float x) {
+  #if defined(__MUSA_ARCH__)
+    #if (__MUSA_ARCH__ >= 310)
+      return __musa_tanh_f(x);
+    #else
+      return ::tanhf(x);
+    #endif
+  #else
+  return std::tanh(x);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+double fast_tanh(double x) {
+  #if defined(__MUSA_ARCH__)
+  return ::tanh(x);
+  #else
+  return std::tanh(x);
+  #endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t fast_tanh(half_t x) {
+  return half_t(fast_tanh(float(x)));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct fast_exp_op {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &rhs) const {
+    return fast_exp(rhs);
+  }
+};
+
+#if defined(__MUSACC__) && (_cplusplus)
+template <int N>
+struct fast_exp_op<Array<half_t, N>> {
+  MUTLASS_DEVICE
+  Array<half_t, N> operator()(Array<half_t, N> const &rhs) const {
+
+    Array<half_t, N> result;
+
+    // use x2 specialization
+    __half2 const *in  = reinterpret_cast<__half2 const *>(&rhs);
+    __half2 *out = reinterpret_cast<__half2 *>(&result);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      out[i] = ::h2exp(in[i]);
+    }
+
+    // residual
+    if (N % 2) {
+      half_t last = rhs[N - 1];
+      result[N - 1] = half_t(::hexp(last.to_half()));
+    }
+
+    return result;
+  }
+};
+#endif // #if defined(__MUSACC__) && (_cplusplus)
+
+template <typename T, int N>
+struct fast_exp_op<Array<T, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &rhs) const {
+
+    fast_exp_op<T> fast_op;
+    Array<T, N> y;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = fast_op(rhs[i]);
+    }
+
+    return y;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct fast_tanh_op {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &rhs) const {
+    return fast_tanh(rhs);
+  }
+};
+
+template <typename T, int N>
+struct fast_tanh_op<Array<T, N>> {
+  MUTLASS_HOST_DEVICE
+  Array<T, N> operator()(Array<T, N> const &rhs) const {
+
+    fast_tanh_op<T> fast_op;
+    Array<T, N> y;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      y[i] = fast_op(rhs[i]);
+    }
+
+    return y;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Absolute value function
+template <typename T>
+MUTLASS_HOST_DEVICE
+T absolute_value(T x) {
+  if (x < T()) {
+    return -x;
+  }
+  return x;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/float8.h b/vllm_musa/csrc/flash_attn/include/mutlass/float8.h
new file mode 100644
index 000000000..aefb1a358
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/float8.h
@@ -0,0 +1,1236 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a class for using IEEE half-precision floating-point types in host or
+      device code.
+*/
+
+#pragma once
+
+#define MUSA_FP8_ENABLED 1
+
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 310)
+#define MUSA_FP8_CVT_ENABLED 1
+#endif // defined(__MUSA_ARCH__)
+
+#ifdef __GNUC__
+// Ignore checks on reinterpret-casts that are being used for bitcasts.
+#pragma GCC diagnostic ignored "-Wstrict-aliasing"
+#endif
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(__MUSACC_RTC__)
+
+#include "mutlass/floating_point_mtrtc.h"
+
+#else
+#include <cmath>
+#include <limits>
+#include <cstdint>
+#include <cstring>
+#endif
+
+#ifdef MUSA_FP8_ENABLED
+#include <musa_fp8.h>
+#endif
+#include <musa_fp16.h>
+
+#include "mutlass/mutlass.h"
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+//  FP8 Has 2 encodings possible : E4M3 and E5M2
+//
+//  E4M3 : 7  |  6 5 4 3  |  2 1 0
+//  E5M2 : 7  |  6 5 4 3 2  |  1 0
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+enum class FloatEncoding {
+    E4M3,
+    E5M2
+};
+
+template<FloatEncoding T>
+struct alignas(1) float8_base {
+
+    static constexpr bool IS_E4M3 = (T == FloatEncoding::E4M3);
+    static constexpr bool IS_E5M2 = (T == FloatEncoding::E5M2);
+
+    // Number of Bits representing mantissa and exponents
+    static constexpr int FP32_NUM_BITS = 32;
+    static constexpr int FP32_NUM_EXPONENT_BITS = 8;
+    static constexpr int FP32_NUM_MANTISSA_BITS = 23;
+    static constexpr uint32_t FP32_NAN = 0x7fffffff;
+    static constexpr uint32_t FP32_INFINITY_MASK = 0x7f800000;
+    static constexpr int FP32_MAX_EXPONENT  =  127;
+    static constexpr int FP32_MIN_EXPONENT  = -126;
+    static constexpr int FP32_EXPONENT_BIAS =  127;
+
+    static constexpr int FP16_NUM_BITS = 16;
+    static constexpr int FP16_NUM_EXPONENT_BITS = 5;
+    static constexpr int FP16_NUM_MANTISSA_BITS = 10;
+    static constexpr uint16_t FP16_NAN = 0x7fff;
+    static constexpr uint16_t FP16_INFINITY_MASK = 0x7c00;
+    static constexpr int FP16_MAX_EXPONENT  = 15;
+    static constexpr int FP16_MIN_EXPONENT  = -14;
+    static constexpr int FP16_EXPONENT_BIAS = 15;
+
+    static constexpr int FP8_NUM_BITS = 8;
+    static constexpr int FP8_NUM_EXPONENT_BITS = IS_E4M3 ? 4 : 5;
+    static constexpr int FP8_NUM_MANTISSA_BITS = IS_E4M3 ? 3 : 2;
+    static constexpr uint8_t  FP8_NAN = 0x7f; // Also F8_INF
+    static constexpr uint8_t  FP8_INFINITY_MASK = IS_E4M3 ? 0x78 : 0x7c;
+    static constexpr int FP8_MAX_EXPONENT  = IS_E4M3 ?  7 :  15;
+    static constexpr int FP8_MIN_EXPONENT  = IS_E4M3 ? -6 : -14;
+    static constexpr int FP8_EXPONENT_BIAS = IS_E4M3 ?  7 :  15;
+
+    static constexpr uint8_t  FP8_EXPONENT_MASK = (1 << FP8_NUM_EXPONENT_BITS) - 1;
+    static constexpr uint8_t  FP8_MANTISSA_MASK = (1 << FP8_NUM_MANTISSA_BITS) - 1;
+
+    static constexpr uint8_t FP8_MAX_FLT = (IS_E4M3 ? 0x7e : 0x7b);
+
+    // 256 in float
+    static constexpr uint32_t FP8_SAT_VAL_FP32 = 0x43800000;
+
+    //
+    // Data members
+    //
+
+    /// Data container
+    uint8_t storage;
+
+    /// Ctors.
+    MUTLASS_HOST_DEVICE
+    float8_base() = default;
+
+    /// Is finite implementation
+    MUTLASS_HOST_DEVICE
+    static bool isfinite(float flt) {
+        uint32_t s;
+
+        #if defined(__MUSA_ARCH__)
+        s = reinterpret_cast<uint32_t const &>(flt);
+        #else
+        std::memcpy(&s, &flt, sizeof(s));
+        #endif
+
+        return (s & 0x7f800000) < 0x7f800000;
+    }
+
+    /// Is NaN implementation
+    MUTLASS_HOST_DEVICE
+    static bool isnan(float flt) {
+        uint32_t s;
+
+        #if defined(__MUSA_ARCH__)
+        s = reinterpret_cast<uint32_t const &>(flt);
+        #else
+        std::memcpy(&s, &flt, sizeof(s));
+        #endif
+
+        return (s & 0x7fffffff) > 0x7f800000;
+    }
+
+    /// Is infinite implementation
+    MUTLASS_HOST_DEVICE
+    static bool isinf(float flt) {
+        uint32_t s;
+
+        #if defined(__MUSA_ARCH__)
+        s = reinterpret_cast<uint32_t const &>(flt);
+        #else
+        std::memcpy(&s, &flt, sizeof(s));
+        #endif
+
+        // Sign = 0 for +inf, 1 for -inf
+        // Exponent = all ones
+        // Mantissa = all zeros
+        return (s == 0x7f800000) || (s == 0xff800000);
+    }
+
+    /// FP32 -> FP8 conversion - rounds to nearest even
+    MUTLASS_HOST_DEVICE
+    static uint8_t convert_float_to_fp8(float const& flt) {
+
+        // software implementation rounds toward nearest even
+        uint32_t s;
+
+        #if defined(__MUSA_ARCH__)
+        s = reinterpret_cast<uint32_t const &>(flt);
+        #else
+        std::memcpy(&s, &flt, sizeof(s));
+        #endif
+
+        // Extract the bits in the FP32 type
+        uint8_t sign = uint8_t((s >> 24 & 0x80));
+        int32_t exp = int32_t((s >> FP32_NUM_MANTISSA_BITS) & 0xff) - FP32_EXPONENT_BIAS;
+        int mantissa = s & 0x7fffff;
+        uint8_t u = 0;
+
+        uint8_t const kF8_NaN = 0x7f;
+
+        // NaN => NaN
+        if (isnan(flt)) {
+            return kF8_NaN;
+        }
+
+        // Inf => MAX_FLT (satfinite)
+        if (isinf(flt)) {
+            return sign | FP8_MAX_FLT;
+        }
+
+        // Special handling
+        if (exp == -128) {
+            // int8 range is from -128 to 127
+            // So 255(inf) - 127(bias) = 128 - will show up as -128
+
+            // satfinite
+            return (sign | FP8_MAX_FLT);
+        }
+
+        int sticky_bit = 0;
+
+        bool skip_sign = false;
+        bool may_be_nan = false;
+
+        if ( (exp >= FP8_MIN_EXPONENT) && (exp <= FP8_MAX_EXPONENT) ) {
+            // normal fp32 to normal fp8
+            exp = exp + FP8_EXPONENT_BIAS;
+            u = uint8_t((uint32_t(exp) & FP8_EXPONENT_MASK) << FP8_NUM_MANTISSA_BITS);
+            u = uint8_t(u | (mantissa >> (FP32_NUM_MANTISSA_BITS - FP8_NUM_MANTISSA_BITS)));
+        } else if(exp < FP8_MIN_EXPONENT) {
+            // normal single-precision to subnormal float8-precision representation
+            int rshift = (FP8_MIN_EXPONENT - exp);
+            if (rshift < FP32_NUM_BITS) {
+                mantissa |= (1 << FP32_NUM_MANTISSA_BITS);
+
+                sticky_bit = ((mantissa & ((1 << rshift) - 1)) != 0);
+
+                mantissa = (mantissa >> rshift);
+                u = (uint8_t(mantissa >> (FP32_NUM_MANTISSA_BITS- FP8_NUM_MANTISSA_BITS)) & FP8_MANTISSA_MASK);
+            } else {
+                mantissa = 0;
+                u = 0;
+            }
+        // Exponent > FP8_MAX_EXPONENT - this is a special case done to match HW
+        // 0x4380_0000 to 0x43e0_0000 - maps from 256 to 448, and does not saturate / inf.
+        } else {
+            if( exp == (FP8_MAX_EXPONENT + 1) ) {
+                uint8_t mantissa_tmp = uint8_t(mantissa >> (FP32_NUM_MANTISSA_BITS - FP8_NUM_MANTISSA_BITS));
+                if( mantissa_tmp < FP8_MANTISSA_MASK) {
+                    exp = exp + FP8_EXPONENT_BIAS;
+                    u = uint8_t(uint32_t(exp) << FP8_NUM_MANTISSA_BITS) | mantissa_tmp;
+                    may_be_nan =  (mantissa_tmp == (FP8_MANTISSA_MASK-1));
+                } else {
+                    // satfinite
+                    return (sign | FP8_MAX_FLT);
+                }
+            } else{
+                // satfinite
+                return (sign | FP8_MAX_FLT);
+            }
+        }
+
+        // round to nearest even
+        int NUM_BITS_SHIFT = FP32_NUM_MANTISSA_BITS - (FP8_NUM_MANTISSA_BITS + 1);
+        int round_bit = ((mantissa >> NUM_BITS_SHIFT) & 1);
+        sticky_bit |= ((mantissa & ((1 << NUM_BITS_SHIFT) - 1)) != 0);
+
+        if ((round_bit && sticky_bit) || (round_bit && (u & 1))) {
+            u = uint8_t(u + 1);
+            if( may_be_nan ) {
+                skip_sign = true;
+            }
+        }
+
+        if (u > FP8_MAX_FLT) {
+            // satfinite
+            u = (sign | FP8_MAX_FLT);
+        }
+
+        if( ! skip_sign ) {
+            u |= sign;
+        }
+
+        return u;
+    }
+
+
+    /// Converts a fp8 value stored as a uint8_t to a float
+    MUTLASS_HOST_DEVICE
+    static float convert_fp8_to_float(uint8_t const& x) {
+
+        uint32_t constexpr kF32_NaN = 0x7fffffff;
+
+        uint8_t const &f8 = x;
+        uint32_t sign = (f8 >> (FP8_NUM_BITS - 1)) & 1;
+        uint32_t exp = (f8 >> FP8_NUM_MANTISSA_BITS) & FP8_EXPONENT_MASK;
+        uint32_t mantissa = f8 & FP8_MANTISSA_MASK;
+        unsigned f = (sign << (FP32_NUM_BITS-1));
+
+        if (IS_E4M3 && exp == 15 && mantissa == 0x7) {
+            f = kF32_NaN;
+        }
+        else if (exp > 0 && (IS_E4M3 || exp < (FP8_MAX_EXPONENT + FP8_EXPONENT_BIAS + 1))) {
+            // normal
+            exp += (FP32_EXPONENT_BIAS - FP8_EXPONENT_BIAS);
+            f = f |
+                (exp << FP32_NUM_MANTISSA_BITS) |
+                (mantissa << (FP32_NUM_MANTISSA_BITS-FP8_NUM_MANTISSA_BITS));
+        } else if (exp == 0) {
+            if (mantissa) {
+                // subnormal
+                exp += (FP32_EXPONENT_BIAS - FP8_EXPONENT_BIAS) + 1;
+                while ((mantissa & (1 << FP8_NUM_MANTISSA_BITS)) == 0) {
+                    mantissa <<= 1;
+                    exp--;
+                }
+                mantissa &= FP8_MANTISSA_MASK;
+                f = f |
+                    (exp << FP32_NUM_MANTISSA_BITS) |
+                    (mantissa << (FP32_NUM_MANTISSA_BITS-FP8_NUM_MANTISSA_BITS));
+            } else {
+                // sign-preserving zero
+            }
+        } else {
+            if(mantissa == 0){
+                // Sign-preserving infinity
+                f = (f | 0x7f800000);
+            } else {
+                // Canonical NaN
+                f = kF32_NaN;
+            }
+        }
+
+        #if defined(__MUSA_ARCH__)
+        return reinterpret_cast<float const&>(f);
+        #else
+        float flt;
+        std::memcpy(&flt, &f, sizeof(flt));
+        return flt;
+        #endif
+    }
+};
+
+
+// Forward declaration of float_e5m2_t to define float_e4m3_t <=> float_e5m2_t
+// conversions in class float_e4m3_t
+struct float_e5m2_t;
+
+
+///////////////////////////////////////////////////////////////
+///
+/// floating-point 8 type : E4M3
+///
+///////////////////////////////////////////////////////////////
+struct alignas(1) float_e4m3_t : float8_base<FloatEncoding::E4M3> {
+
+    using Base = float8_base<FloatEncoding::E4M3>;
+
+    static constexpr int MAX_EXPONENT = Base::FP8_MAX_EXPONENT;
+
+    //
+    // Static conversion operators
+    //
+
+    /// Constructs from an uint8_t
+    MUTLASS_HOST_DEVICE
+    static float_e4m3_t bitcast(uint8_t x) {
+        float_e4m3_t f;
+        f.storage = x;
+        return f;
+    }
+
+    /// FP32 -> FP8 conversion - rounds to nearest even
+    MUTLASS_HOST_DEVICE
+    static float_e4m3_t from_float(float const& flt) {
+    #if defined(MUSA_FP8_CVT_ENABLED)
+        uint8_t tmp = __musa_f2e4m3_rn(flt);
+
+        return *reinterpret_cast<float_e4m3_t *>(&tmp);
+    #else
+        return bitcast(Base::convert_float_to_fp8(flt));
+    #endif
+    }
+
+    /// FP16 -> E5M2 conversion - rounds to nearest even
+    MUTLASS_HOST_DEVICE
+    static float_e4m3_t from_half(half const& flt) {
+    #if defined(MUSA_FP8_CVT_ENABLED)
+        uint8_t tmp = __musa_f162e4m3_rn(flt);
+
+        return *reinterpret_cast<float_e4m3_t *>(&tmp);
+    #else
+        return bitcast(Base::convert_float_to_fp8(__half2float(flt)));
+    #endif
+    }
+
+    // E4M3 -> half
+    MUTLASS_HOST_DEVICE
+    static half to_half(float_e4m3_t const& x) {
+    #if defined(MUSA_FP8_CVT_ENABLED)
+        uint8_t bits = x.storage;
+        return __musa_e4m32f16_rn(bits);
+    #else
+        return __float2half(Base::convert_fp8_to_float(x.storage));
+    #endif
+    }
+
+    // E4M3 -> Float
+    MUTLASS_HOST_DEVICE
+    static float to_float(float_e4m3_t const& x) {
+    #if defined(MUSA_FP8_CVT_ENABLED)
+        uint8_t bits = x.storage;
+        return __half2float(__musa_e4m32f16_rn(bits));
+    #else
+        return Base::convert_fp8_to_float(x.storage);
+    #endif
+    }
+
+    //
+    // Methods
+    //
+
+    /// Constructor inheritance
+    using Base::Base;
+
+    /// Default constructor
+    float_e4m3_t() = default;
+
+#ifdef MUSA_FP8_ENABLED
+    /// Conversion from MUSA's FP8 type
+    MUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(__mt_fp8_e4m3 x) {
+        storage = x.__x;
+    }
+#endif
+
+    /// Floating point conversion
+    MUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(float x) {
+        storage = from_float(x).storage;
+    }
+
+    MUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(half x) {
+        storage = from_half(x).storage;
+    }
+
+    /// Floating point conversion
+    MUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(double x): float_e4m3_t(float(x)) {
+    }
+
+    /// Integer conversion
+    MUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(int x): float_e4m3_t(float(x)) {
+    }
+
+    MUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(unsigned x): float_e4m3_t(float(x)) {
+    }
+
+    /// E5M2 conversion. Defined after float_e5m2_t is defined.
+    MUTLASS_HOST_DEVICE
+    explicit float_e4m3_t(float_e5m2_t x);
+
+#ifdef MUSA_FP8_ENABLED
+    /// Assignment from MUSA's FP8 type
+    MUTLASS_HOST_DEVICE
+    float_e4m3_t & operator=(__mt_fp8_e4m3 x) {
+        storage = x.__x;
+        return *this;
+    }
+#endif
+
+    /// Converts to float
+    MUTLASS_HOST_DEVICE
+    operator float() const {
+        return to_float(*this);
+    }
+
+    /// Converts to half
+    MUTLASS_HOST_DEVICE
+    operator half() const {
+        return to_half(*this);
+    }
+
+    /// Converts to float
+    MUTLASS_HOST_DEVICE
+    explicit operator double() const {
+        return double(to_float(*this));
+    }
+
+    /// Converts to int
+    MUTLASS_HOST_DEVICE
+    explicit operator int() const {
+    #if defined(__MUSA_ARCH__)
+        return __half2int_rn(to_half(*this));
+    #else
+        return int(to_float(*this));
+    #endif
+    }
+
+    /// Casts to bool
+    MUTLASS_HOST_DEVICE
+    explicit operator bool() const {
+    #if defined(__MUSA_ARCH__)
+        return bool(__half2int_rn(to_half(*this)));
+    #else
+        return bool(int(to_float(*this)));
+    #endif
+    }
+
+    /// Accesses raw internal state
+    MUTLASS_HOST_DEVICE
+    uint8_t& raw() {
+        return storage;
+    }
+
+    /// Accesses raw internal state
+    MUTLASS_HOST_DEVICE
+    uint8_t raw() const {
+        return storage;
+    }
+
+    /// Returns the sign bit
+    MUTLASS_HOST_DEVICE
+    bool signbit() const {
+        return ((storage & (1 << (Base::FP8_NUM_BITS - 1))) != 0);
+    }
+
+    /// Returns the biased exponent
+    MUTLASS_HOST_DEVICE
+    int exponent_biased() const {
+        return int((storage >> FP8_NUM_MANTISSA_BITS) & Base::FP8_EXPONENT_MASK);
+    }
+
+    /// Returns the unbiased exponent
+    MUTLASS_HOST_DEVICE
+    int exponent() const {
+        return exponent_biased() - 15;
+    }
+
+    /// Returns the mantissa
+    MUTLASS_HOST_DEVICE
+    int mantissa() const {
+        return int(storage & Base::FP8_MANTISSA_MASK);
+    }
+};
+///////////////////////////////////////////////////////////////
+///
+/// floating-point 8 type : E5M2
+///
+///////////////////////////////////////////////////////////////
+struct alignas(1) float_e5m2_t : float8_base<FloatEncoding::E5M2> {
+
+    using Base = float8_base<FloatEncoding::E5M2>;
+
+    static constexpr int MAX_EXPONENT = Base::FP8_MAX_EXPONENT;
+
+    //
+    // Static conversion operators
+    //
+
+    /// Constructs from an uint8_t
+    MUTLASS_HOST_DEVICE
+    static float_e5m2_t bitcast(uint8_t x) {
+        float_e5m2_t f;
+        f.storage = x;
+        return f;
+    }
+
+    /// FP32 -> FP8 conversion - rounds to nearest even
+    MUTLASS_HOST_DEVICE
+    static float_e5m2_t from_float(float const& flt) {
+    #if defined(MUSA_FP8_CVT_ENABLED)
+        uint8_t tmp = __musa_f2e5m2_rn(flt);
+
+        return *reinterpret_cast<float_e5m2_t *>(&tmp);
+    #else
+        return bitcast(Base::convert_float_to_fp8(flt));
+    #endif
+    }
+
+    /// FP16 -> E5M2 conversion - rounds to nearest even
+    MUTLASS_HOST_DEVICE
+    static float_e5m2_t from_half(half const& flt) {
+    #if defined(MUSA_FP8_CVT_ENABLED)
+        uint8_t tmp = __musa_f162e5m2_rn(flt);
+
+        return *reinterpret_cast<float_e5m2_t *>(&tmp);
+    #else
+        return bitcast(Base::convert_float_to_fp8(__half2float(flt)));
+    #endif
+    }
+
+    // E5M2 -> half
+    MUTLASS_HOST_DEVICE
+    static half to_half(float_e5m2_t const& x) {
+    #if defined(MUSA_FP8_CVT_ENABLED)
+        uint8_t bits = x.storage;
+        return __musa_e5m22f16_rn(bits);
+    #else
+        return __float2half(Base::convert_fp8_to_float(x.storage));
+    #endif
+    }
+
+    // E5M2 -> Float
+    MUTLASS_HOST_DEVICE
+    static float to_float(float_e5m2_t const& x) {
+    #if defined(MUSA_FP8_CVT_ENABLED)
+        uint8_t bits = x.storage;
+        return __half2float(__musa_e5m22f16_rn(bits));
+    #else
+        return Base::convert_fp8_to_float(x.storage);
+    #endif
+    }
+
+    //
+    // Methods
+    //
+
+    /// Constructor inheritance
+    using Base::Base;
+
+    /// Default constructor
+    float_e5m2_t() = default;
+
+#ifdef MUSA_FP8_ENABLED
+    /// Conversion from MUSA's FP8 type
+    MUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(__mt_fp8_e5m2 x) {
+        storage = x.__x;
+    }
+#endif
+
+    /// Floating point conversion
+    MUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(float x) {
+        storage = from_float(x).storage;
+    }
+
+    MUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(half x) {
+      storage = from_half(x).storage;
+    }
+
+    /// Floating point conversion
+    MUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(double x): float_e5m2_t(float(x)) {
+    }
+
+    /// Integer conversion
+    MUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(int x): float_e5m2_t(float(x)) {
+    }
+
+    MUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(unsigned x): float_e5m2_t(float(x)) {
+    }
+
+    /// E4M3 conversion
+    MUTLASS_HOST_DEVICE
+    explicit float_e5m2_t(float_e4m3_t x);
+
+#ifdef MUSA_FP8_ENABLED
+    /// Assignment from MUSA's FP8 type
+    MUTLASS_HOST_DEVICE
+    float_e5m2_t & operator=(__mt_fp8_e5m2 x) {
+        storage = x.__x;
+        return *this;
+    }
+#endif
+
+    /// Converts to float
+    MUTLASS_HOST_DEVICE
+    operator float() const {
+        return to_float(*this);
+    }
+
+    /// Converts to half
+    MUTLASS_HOST_DEVICE
+    operator half() const {
+      return to_half(*this);
+    }
+
+    /// Converts to float
+    MUTLASS_HOST_DEVICE
+    explicit operator double() const {
+        return double(to_float(*this));
+    }
+
+    /// Converts to int
+    MUTLASS_HOST_DEVICE
+    explicit operator int() const {
+    #if defined(__MUSA_ARCH__)
+        return __half2int_rn(to_half(*this));
+    #else
+        return int(to_float(*this));
+    #endif
+    }
+
+    /// Casts to bool
+    MUTLASS_HOST_DEVICE
+    explicit operator bool() const {
+    #if defined(__MUSA_ARCH__)
+        return bool(__half2int_rn(to_half(*this)));
+    #else
+        return bool(int(to_float(*this)));
+    #endif
+    }
+
+    /// Accesses raw internal state
+    MUTLASS_HOST_DEVICE
+    uint8_t& raw() {
+        return storage;
+    }
+
+    /// Accesses raw internal state
+    MUTLASS_HOST_DEVICE
+    uint8_t raw() const {
+        return storage;
+    }
+
+    /// Returns the sign bit
+    MUTLASS_HOST_DEVICE
+    bool signbit() const {
+        return ((storage & (1 << (Base::FP8_NUM_BITS - 1))) != 0);
+    }
+
+    /// Returns the biased exponent
+    MUTLASS_HOST_DEVICE
+    int exponent_biased() const {
+        return int((storage >> FP8_NUM_MANTISSA_BITS) & Base::FP8_EXPONENT_MASK);
+    }
+
+    /// Returns the unbiased exponent
+    MUTLASS_HOST_DEVICE
+    int exponent() const {
+        return exponent_biased() - 15;
+    }
+
+    /// Returns the mantissa
+    MUTLASS_HOST_DEVICE
+    int mantissa() const {
+        return int(storage & Base::FP8_MANTISSA_MASK);
+    }
+};
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Arithmetic operators
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+MUTLASS_HOST_DEVICE
+bool operator==(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) == float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator!=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) != float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) < float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) <= float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) > float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>=(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float(lhs) >= float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t operator+(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float_e4m3_t(float(lhs) + float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t operator-(float_e4m3_t const& lhs) {
+    return float_e4m3_t(-float(lhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t operator-(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float_e4m3_t(float(lhs) - float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t operator*(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float_e4m3_t(float(lhs) * float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t operator/(float_e4m3_t const& lhs, float_e4m3_t const& rhs) {
+    return float_e4m3_t(float(lhs) / float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t& operator+=(float_e4m3_t & lhs, float_e4m3_t const& rhs) {
+    lhs = float_e4m3_t(float(lhs) + float(rhs));
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t& operator-=(float_e4m3_t & lhs, float_e4m3_t const& rhs) {
+    lhs = float_e4m3_t(float(lhs) - float(rhs));
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t& operator*=(float_e4m3_t & lhs, float_e4m3_t const& rhs) {
+    lhs = float_e4m3_t(float(lhs) * float(rhs));
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t& operator/=(float_e4m3_t & lhs, float_e4m3_t const& rhs) {
+    lhs = float_e4m3_t(float(lhs) / float(rhs));
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t& operator++(float_e4m3_t & lhs) {
+    float tmp(lhs);
+    ++tmp;
+    lhs = float_e4m3_t(tmp);
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t& operator--(float_e4m3_t & lhs) {
+    float tmp(lhs);
+    --tmp;
+    lhs = float_e4m3_t(tmp);
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t operator++(float_e4m3_t & lhs, int) {
+    float_e4m3_t ret(lhs);
+    float tmp(lhs);
+    tmp++;
+    lhs = float_e4m3_t(tmp);
+    return ret;
+}
+
+MUTLASS_HOST_DEVICE
+float_e4m3_t operator--(float_e4m3_t & lhs, int) {
+    float_e4m3_t ret(lhs);
+    float tmp(lhs);
+    tmp--;
+    lhs = float_e4m3_t(tmp);
+    return ret;
+}
+
+MUTLASS_HOST_DEVICE
+bool operator==(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) == float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator!=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) != float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) < float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) <= float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) > float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>=(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float(lhs) >= float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t operator+(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float_e5m2_t(float(lhs) + float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t operator-(float_e5m2_t const& lhs) {
+    return float_e5m2_t(-float(lhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t operator-(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float_e5m2_t(float(lhs) - float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t operator*(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float_e5m2_t(float(lhs) * float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t operator/(float_e5m2_t const& lhs, float_e5m2_t const& rhs) {
+    return float_e5m2_t(float(lhs) / float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t& operator+=(float_e5m2_t & lhs, float_e5m2_t const& rhs) {
+    lhs = float_e5m2_t(float(lhs) + float(rhs));
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t& operator-=(float_e5m2_t & lhs, float_e5m2_t const& rhs) {
+    lhs = float_e5m2_t(float(lhs) - float(rhs));
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t& operator*=(float_e5m2_t & lhs, float_e5m2_t const& rhs) {
+    lhs = float_e5m2_t(float(lhs) * float(rhs));
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t& operator/=(float_e5m2_t & lhs, float_e5m2_t const& rhs) {
+    lhs = float_e5m2_t(float(lhs) / float(rhs));
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t& operator++(float_e5m2_t & lhs) {
+    float tmp(lhs);
+    ++tmp;
+    lhs = float_e5m2_t(tmp);
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t& operator--(float_e5m2_t & lhs) {
+    float tmp(lhs);
+    --tmp;
+    lhs = float_e5m2_t(tmp);
+    return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t operator++(float_e5m2_t & lhs, int) {
+    float_e5m2_t ret(lhs);
+    float tmp(lhs);
+    tmp++;
+    lhs = float_e5m2_t(tmp);
+    return ret;
+}
+
+MUTLASS_HOST_DEVICE
+float_e5m2_t operator--(float_e5m2_t & lhs, int) {
+    float_e5m2_t ret(lhs);
+    float tmp(lhs);
+    tmp--;
+    lhs = float_e5m2_t(tmp);
+    return ret;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// float_e4m3_t <=> float_e5m2_t conversions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// float_e4m3_t <= float_e5m2_t
+MUTLASS_HOST_DEVICE
+float_e4m3_t::float_e4m3_t(float_e5m2_t x) {
+    storage = from_float(float_e5m2_t::to_float(x)).storage;
+}
+
+/// float_e5m2_t <= float_e4m3_t
+MUTLASS_HOST_DEVICE
+float_e5m2_t::float_e5m2_t(float_e4m3_t x) {
+    storage = from_float(float_e4m3_t::to_float(x)).storage;
+}
+
+///////////////////////////////////////////////////////////////
+///
+/// Umbrella floating-point 8-bit data type : type_erased_dynamic_float8_t
+/// This umbrella datatype can be enabled when a user provides a specific
+/// datatype in runtime argument list.
+///
+/// Currently supported runtime datatypes compatible with type_erased_dynamic_float8_t:
+///   QMMAFormat::E5M2
+///   QMMAFormat::E4M3
+///
+///////////////////////////////////////////////////////////////
+
+union type_erased_dynamic_float8_t {
+  uint8_t data;
+  mutlass::float_e5m2_t e5m2;
+  mutlass::float_e4m3_t e4m3;
+  MUTLASS_HOST_DEVICE
+  explicit operator mutlass::float_e5m2_t() const {
+    return e5m2;
+  }
+
+  MUTLASS_HOST_DEVICE
+  explicit operator mutlass::float_e4m3_t() const {
+    return e4m3;
+  }
+
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Standard Library operations and definitions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if !defined(__MUSACC_RTC__)
+namespace std {
+
+/// Numeric limits common to all float8 types
+template <typename T>
+struct float8_base_numeric_limits {
+private:
+  using F8Type = T;
+public:
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+  static std::float_denorm_style const has_denorm = std::denorm_present;
+  static bool const has_denorm_loss = true;
+  static std::float_round_style const round_style = std::round_to_nearest;
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = F8Type::FP8_NUM_MANTISSA_BITS;
+
+  /// Least positive value
+  MUTLASS_HOST_DEVICE
+  static F8Type min() { return F8Type::bitcast(0x01); }
+
+  /// Maximum finite value
+  MUTLASS_HOST_DEVICE
+  static F8Type max() { return F8Type::bitcast(F8Type::FP8_MAX_FLT); }
+
+  /// Returns maximum rounding error
+  MUTLASS_HOST_DEVICE
+  static F8Type round_error() { return F8Type(0.5f); }
+
+  /// Returns positive infinity value
+  MUTLASS_HOST_DEVICE
+  static F8Type infinity() { return F8Type::bitcast(F8Type::FP8_INFINITY_MASK); }
+
+  /// Returns quiet NaN value
+  MUTLASS_HOST_DEVICE
+  static F8Type quiet_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); }
+
+  /// Returns signaling NaN value
+  MUTLASS_HOST_DEVICE
+  static F8Type signaling_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); }
+
+  /// Returns smallest positive subnormal value
+  MUTLASS_HOST_DEVICE
+  static F8Type denorm_min() { return F8Type::bitcast(0x01); }
+};
+
+/// Numeric limits for float_e4m3_t
+template <>
+struct numeric_limits<mutlass::float_e4m3_t> :
+    public float8_base_numeric_limits<mutlass::float_e4m3_t> {
+  static bool const has_infinity = false;
+
+  /// Minimum finite value
+  static mutlass::float_e4m3_t lowest() { return mutlass::float_e4m3_t::bitcast(0xfe); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value
+  static mutlass::float_e4m3_t epsilon() { return mutlass::float_e4m3_t::bitcast(0x20); }
+};
+
+/// Numeric limits for float_e5m2_t
+template <>
+struct numeric_limits<mutlass::float_e5m2_t>  :
+    public float8_base_numeric_limits<mutlass::float_e5m2_t> {
+  static bool const has_infinity = true;
+
+  /// Minimum finite value
+  static mutlass::float_e5m2_t lowest() { return mutlass::float_e5m2_t::bitcast(0xfb); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value
+  static mutlass::float_e5m2_t epsilon() { return mutlass::float_e5m2_t::bitcast(0x34); }
+};
+
+}  // namespace std
+#endif
+
+namespace platform {
+
+/// Numeric limits common to all float8 types
+template <typename T>
+struct float8_base_numeric_limits {
+private:
+  using F8Type = T;
+public:
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+#if !defined(__MUSACC_RTC__)
+  static std::float_denorm_style const has_denorm = std::denorm_present;
+#endif
+  static bool const has_denorm_loss = true;
+#if !defined(__MUSACC_RTC__)
+  static std::float_round_style const round_style = std::round_to_nearest;
+#endif
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = F8Type::FP8_NUM_MANTISSA_BITS;
+
+  /// Least positive value
+  MUTLASS_HOST_DEVICE
+  static F8Type min() { return F8Type::bitcast(0x01); }
+
+  /// Maximum finite value
+  MUTLASS_HOST_DEVICE
+  static F8Type max() { return F8Type::bitcast(F8Type::FP8_MAX_FLT); }
+
+  /// Returns maximum rounding error
+  MUTLASS_HOST_DEVICE
+  static F8Type round_error() { return F8Type(0.5f); }
+
+  /// Returns positive infinity value
+  MUTLASS_HOST_DEVICE
+  static F8Type infinity() { return F8Type::bitcast(F8Type::FP8_INFINITY_MASK); }
+
+  /// Returns quiet NaN value
+  MUTLASS_HOST_DEVICE
+  static F8Type quiet_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); }
+
+  /// Returns signaling NaN value
+  MUTLASS_HOST_DEVICE
+  static F8Type signaling_NaN() { return F8Type::bitcast(F8Type::FP8_NAN); }
+
+  /// Returns smallest positive subnormal value
+  MUTLASS_HOST_DEVICE
+  static F8Type denorm_min() { return F8Type::bitcast(0x01); }
+};
+
+/// std::numeric_limits
+template <class T>
+struct numeric_limits;
+
+/// Numeric limits for float_e4m3_t
+template <>
+struct numeric_limits<mutlass::float_e4m3_t> :
+    public float8_base_numeric_limits<mutlass::float_e4m3_t> {
+  static bool const has_infinity = false;
+
+  /// Minimum finite value
+  static mutlass::float_e4m3_t lowest() { return mutlass::float_e4m3_t::bitcast(0xfe); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value
+  static mutlass::float_e4m3_t epsilon() { return mutlass::float_e4m3_t::bitcast(0x20); }
+};
+
+/// Numeric limits for float_e5m2_t
+template <>
+struct numeric_limits<mutlass::float_e5m2_t>  :
+    public float8_base_numeric_limits<mutlass::float_e5m2_t> {
+  static bool const has_infinity = true;
+
+  /// Minimum finite value
+  static mutlass::float_e5m2_t lowest() { return mutlass::float_e5m2_t::bitcast(0xfb); }
+
+  /// Machine epsilon, that is, the difference between 1.0 and the next representable value
+  static mutlass::float_e5m2_t epsilon() { return mutlass::float_e5m2_t::bitcast(0x34); }
+};
+
+}  // namespace platform
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// User-defined literals
+//
+
+MUTLASS_HOST_DEVICE
+mutlass::float_e4m3_t operator "" _fe4m3(long double x) {
+  return mutlass::float_e4m3_t(float(x));
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::float_e4m3_t operator "" _fe4m3(unsigned long long int x) {
+  return mutlass::float_e4m3_t(int(x));
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::float_e5m2_t operator "" _fe5m2(long double x) {
+  return mutlass::float_e5m2_t(float(x));
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::float_e5m2_t operator "" _fe5m2(unsigned long long int x) {
+  return mutlass::float_e5m2_t(int(x));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/floating_point_mtrtc.h b/vllm_musa/csrc/flash_attn/include/mutlass/floating_point_mtrtc.h
new file mode 100644
index 000000000..2ad670547
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/floating_point_mtrtc.h
@@ -0,0 +1,62 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+// All floating-point numbers can be put in one of these categories.
+enum  {
+    FP_NAN =
+# define FP_NAN 0
+      FP_NAN,
+    FP_INFINITE =
+# define FP_INFINITE 1
+      FP_INFINITE,
+    FP_ZERO =
+# define FP_ZERO 2
+      FP_ZERO,
+    FP_SUBNORMAL =
+# define FP_SUBNORMAL 3
+      FP_SUBNORMAL,
+    FP_NORMAL =
+# define FP_NORMAL 4
+      FP_NORMAL
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/functional.h b/vllm_musa/csrc/flash_attn/include/mutlass/functional.h
new file mode 100644
index 000000000..215ff6574
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/functional.h
@@ -0,0 +1,640 @@
+  /***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Define basic numeric operators
+
+    This is inspired by the Standard Library's <functional> header.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/numeric_types.h"
+
+#include <musa_runtime.h>
+
+#ifdef _MSC_VER
+// Provides support for alternate operators such as 'and', 'or', ...
+#include <iso646.h>
+#endif // _MSC_VER
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct absolute_value_op {
+  MUTLASS_HOST_DEVICE
+  T operator()(T lhs) const {
+    return abs(lhs);
+  }
+};
+
+template <>
+struct absolute_value_op<float> {
+  MUTLASS_HOST_DEVICE
+  float operator()(float lhs) const { return fabs(lhs); }
+};
+
+template <typename T>
+struct plus {
+  MUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    lhs += rhs;
+    return lhs;
+  }
+};
+
+template <typename T>
+struct minus {
+  MUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    lhs -= rhs;
+    return lhs;
+  }
+};
+
+template <typename T>
+struct multiplies {
+  MUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    lhs *= rhs;
+    return lhs;
+  }
+};
+
+template <typename T>
+struct scale {
+  T const scaling_factor_;
+
+  MUTLASS_HOST_DEVICE
+  scale(float scaling_factor) : scaling_factor_(scaling_factor) {
+  }
+
+  T operator()(T const &rhs) const {
+    T result = rhs * scaling_factor_;
+    return result;
+  }
+};
+
+#if defined(__MUSA_ARCH__)
+/// Partial specializations needed when __MUSA_NO_HALF2_OPERATORS__ is set
+template<>
+struct plus<__half2> {
+  MUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hadd2(lhs, rhs);
+  }
+};
+
+template<>
+struct minus<__half2> {
+  MUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hsub2(lhs, rhs);
+  }
+};
+
+template<>
+struct multiplies<__half2> {
+  MUTLASS_HOST_DEVICE
+  __half2 operator()(__half2 lhs, __half2 const &rhs) const {
+    return __hmul2(lhs, rhs);
+  }
+};
+
+/// Partial specializations needed when __MUSA_NO_HALF_OPERATORS__ is set
+template<>
+struct plus<__half> {
+  MUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hadd(lhs, rhs);
+  }
+};
+
+template<>
+struct minus<__half> {
+  MUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hsub(lhs, rhs);
+  }
+};
+
+template<>
+struct multiplies<__half> {
+  MUTLASS_HOST_DEVICE
+  __half operator()(__half lhs, __half const &rhs) const {
+    return __hmul(lhs, rhs);
+  }
+};
+#endif // defined(__MUSA_ARCH__)
+
+
+/// Squares with optional conversion
+template <typename T, typename Output = T>
+struct square {
+  MUTLASS_HOST_DEVICE
+  Output operator()(T lhs) const {
+    multiplies<Output> mul_op;
+
+    Output y = Output(lhs);
+    return mul_op(y, y);
+  }
+};
+
+/// Returns the magnitude squared of an element.
+template <typename T, typename Output = T>
+struct magnitude_squared {
+  MUTLASS_HOST_DEVICE
+  Output operator()(T lhs) const {
+    multiplies<Output> mul_op;
+
+    Output y = Output(lhs);
+    return mul_op(y, y);
+  }
+};
+
+/// Computes the square of a difference with optional conversion
+template <typename T, typename Output = T>
+struct square_difference {
+  MUTLASS_HOST_DEVICE
+  Output operator()(T lhs, T rhs) const {
+    multiplies<Output> mul_op;
+
+    Output y = Output(lhs) - Output(rhs);
+    return mul_op(y, y);
+  }
+};
+
+/// Computes the square of a difference with optional conversion
+template <typename T, typename Output = T>
+struct magnitude_squared_difference {
+  MUTLASS_HOST_DEVICE
+  Output operator()(T lhs, T rhs) const {
+    multiplies<Output> mul_op;
+
+    Output y = Output(lhs) - Output(rhs);
+    return mul_op(y, y);
+  }
+};
+
+// Computes the reciprocal square root
+template <typename T>
+struct inverse_square_root;
+
+template <>
+struct inverse_square_root<float> {
+  MUTLASS_HOST_DEVICE
+  float operator()(float const &lhs) const {
+#if defined(__MUSA_ARCH__)
+    return rsqrtf(lhs);
+#else
+    return 1.f / std::sqrt(lhs);
+#endif
+  }
+};
+
+template <>
+struct inverse_square_root<half_t> {
+  MUTLASS_HOST_DEVICE
+  half_t operator()(half_t const &lhs) const {
+#if defined(__MUSA_ARCH__)
+    auto result = hrsqrt(reinterpret_cast<__half const &>(lhs));
+    return reinterpret_cast<half_t const &>(result);
+#else
+    return half_t(1.f / std::sqrt(half_t::convert(lhs)));
+#endif
+  }
+};
+
+/// Divides
+template <typename T>
+struct divides {
+  MUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    lhs /= rhs;
+    return lhs;
+  }
+};
+
+/// reciprocal_approximate
+template <typename T>
+struct reciprocal_approximate {
+  MUTLASS_HOST_DEVICE
+  T operator()(T lhs) const {
+    return divides<T>{}(T(1), lhs);
+  }
+};
+
+template <>
+struct reciprocal_approximate <float> {
+  MUTLASS_HOST_DEVICE
+  float operator()(float lhs) const {
+    float ret;
+      ret = 1.0f / lhs;
+    return ret;
+  }
+};
+
+/// Negate
+template <typename T>
+struct negate {
+  MUTLASS_HOST_DEVICE
+  T operator()(T lhs) const {
+    return -lhs;
+  }
+};
+
+/// Greater equal
+template <typename T>
+struct greater_equal {
+  MUTLASS_HOST_DEVICE
+  bool operator()(T const &lhs, T const &rhs) const {
+    return (lhs >= rhs);
+  }
+};
+
+/// Greater
+template <typename T>
+struct greater {
+  MUTLASS_HOST_DEVICE
+  bool operator()(T const &lhs, T const &rhs) const {
+    return (lhs > rhs);
+  }
+};
+
+/// Less equal
+template <typename T>
+struct less_equal {
+  MUTLASS_HOST_DEVICE
+  bool operator()(T const &lhs, T const &rhs) const {
+    return (lhs <= rhs);
+  }
+};
+
+/// Less
+template <typename T>
+struct less {
+  MUTLASS_HOST_DEVICE
+  bool operator()(T const &lhs, T const &rhs) const {
+    return (lhs < rhs);
+  }
+};
+
+template <typename T, bool PropagateNaN = false>
+struct maximum {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &lhs, T const &rhs) const {
+    if constexpr (PropagateNaN && mutlass::platform::is_floating_point<T>::value) {
+      using MUTLASS_CMATH_NAMESPACE :: isnan;
+
+      // Call isnan unqualified, so argument-dependent lookup (ADL)
+      // will find overloads such as mutlass::isnan(half_t).
+      // Calling ::isnan or std::isnan directly would force
+      // implicit conversions to float of custom number types
+      // in the mutlass namespace (e.g., mutlass::half_t).
+      return lhs > rhs || isnan(lhs) ? lhs : rhs;
+    }
+    else {
+      return (lhs < rhs ? rhs : lhs);
+    }
+  }
+};
+
+// This is a subclass and not an alias
+// in order to work around a known Clang issue,
+// where a template template parameter with one template parameter
+// does not match classes that take multiple template parameters
+// but have defaults for all but the first.
+template<typename T>
+struct maximum_with_default_nan_propagation : public maximum<T>
+{};
+
+// Maximum with nan propagation
+// To propagate NANs, the "max" of a two element that contains NaNs should also return a NaN
+template <typename T>
+struct maximum<T, true> {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &lhs, T const &rhs) const {
+#if defined(__MUSA_ARCH__)
+    return lhs > rhs or isnan(lhs) ? lhs : rhs;
+#else
+    return lhs > rhs or std::isnan(lhs) ? lhs : rhs;
+#endif
+  }
+};
+
+template <>
+struct maximum<float, false> {
+  MUTLASS_HOST_DEVICE
+  float operator()(float const &lhs, float const &rhs) const {
+    return fmaxf(lhs, rhs);
+  }
+};
+
+template <>
+struct maximum<float, true> {
+  MUTLASS_HOST_DEVICE
+  float operator()(float const lhs, float const rhs) const {
+    float res;
+#if defined(__MUSA_ARCH__)
+    res = lhs > rhs or isnan(lhs) ? lhs : rhs;
+#else
+    res = lhs > rhs or std::isnan(lhs) ? lhs : rhs;
+#endif
+    return res;
+  }
+};
+
+// This is a subclass and not an alias
+// in order to work around a known Clang issue,
+// where a template template parameter with one template parameter
+// does not match classes that take multiple template parameters
+// but have defaults for all but the first.
+template <typename T>
+struct maximum_with_nan_propagation : maximum<T, true>
+{};
+
+// This alias exists for backwards compatibility only.
+// Please use the correctly spelled class template above.
+template <typename T>
+using maximum_with_nan_propogation = maximum_with_nan_propagation<T>;
+
+template <typename T, bool PropagateNaN = false>
+struct minimum{
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &lhs, T const &rhs) const {
+    if constexpr (PropagateNaN && mutlass::platform::is_floating_point<T>::value) {
+      using MUTLASS_CMATH_NAMESPACE :: isnan;
+
+      return lhs < rhs || isnan(lhs) ? lhs : rhs;
+    }
+    else {
+      return (rhs < lhs ? rhs : lhs);
+    }
+  }
+};
+
+template <>
+struct minimum<float, false> {
+  MUTLASS_HOST_DEVICE
+  float operator()(float const &lhs, float const &rhs) const {
+    return fminf(lhs, rhs);
+  }
+};
+
+template <typename T>
+struct minimum_with_nan_propagation : minimum<T, true>
+{};
+
+template <typename T, bool PropagateNaN = false>
+struct maximum_absolute_value {
+  MUTLASS_HOST_DEVICE
+  float operator()(T const &lhs, T const &rhs) const {
+    absolute_value_op<T> abs_op;
+    maximum<T, PropagateNaN> max_op;
+
+    return max_op(abs_op(lhs), abs_op(rhs));
+  }
+};
+
+// assumes the left operand is already an absolute value
+template <typename T, bool PropagateNaN = false>
+struct maximum_absolute_value_reduction {
+  MUTLASS_HOST_DEVICE
+  float operator()(T const &lhs, T const &rhs) const {
+    absolute_value_op<T> abs_op;
+    maximum<T, PropagateNaN> max_op;
+
+    return max_op(lhs, abs_op(rhs));
+  }
+};
+
+/// Fused multiply-add
+template <typename A, typename B = A, typename C = A>
+struct multiply_add {
+  MUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    return C(a) * C(b) + c;
+  }
+};
+
+template <typename T>
+struct square_and_plus {
+  MUTLASS_HOST_DEVICE
+  T operator()(T lhs, T const &rhs) const {
+    multiply_add<T> multiply_add_op;
+    return multiply_add_op(rhs, rhs, lhs);
+  }
+};
+
+// Fused multiply-add that takes exactly one template parameter.
+// This is useful for working around a known Clang issue,
+// where a template template parameter with one template parameter
+// does not match classes that take multiple template parameters
+// but have defaults for all but the first.
+template <typename A>
+struct homogeneous_multiply_add : public multiply_add<A, A, A>
+{};
+
+/// Fused multiply-add
+template <typename A, typename B = A, typename C = A>
+struct multiply_add_relu0 {
+  MUTLASS_HOST_DEVICE
+  C operator()(A const &a, B const &b, C const &c) const {
+    maximum<C> mx;
+    return mx(C(a) * C(b) + c, C(0));
+  }
+};
+
+/// Fused multiply-add
+template <typename T>
+struct and_add {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b, T const &c) const {
+    return ((a & b) + c);
+  }
+};
+
+
+/// Fused multiply-add
+template <typename T>
+struct xor_add {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b, T const &c) const {
+    return ((a ^ b) + c);
+  }
+};
+
+template <typename T>
+struct conjugate {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a) const {
+    return a;
+  }
+};
+
+template <typename T>
+struct first {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const & first, T const &...) const {
+    return first;
+  }
+  MUTLASS_HOST_DEVICE
+  T operator()(T const & first) const {
+    return first;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct logical_and {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return ((static_cast<bool>(a) && static_cast<bool>(b)) ? T(1) : T());
+  }
+};
+
+template <typename T>
+struct logical_or {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return ((static_cast<bool>(a) || static_cast<bool>(b)) ? T(1) : T());
+  }
+};
+
+template <typename T>
+struct logical_not {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a) const {
+    return T(!(a));
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct bit_and {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return a & b;
+  }
+};
+
+template <typename T>
+struct bit_or {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return a | b;
+  }
+};
+
+template <typename T>
+struct bit_not {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a) const {
+    return ~a;
+  }
+};
+
+template <typename T>
+struct bit_xor {
+  MUTLASS_HOST_DEVICE
+  T operator()(T const &a, T const &b) const {
+    return a ^ b;
+  }
+};
+
+//////////////////////////////////////////////////////////////////////////////////////////////////
+/// Atomic reductions
+
+template <typename T>
+struct atomic_add
+{
+  MUTLASS_DEVICE
+  void operator()(T *ptr, const T &data)
+  {
+#if defined(__MUSA_ARCH__)
+    atomicAdd(ptr, data);
+#else
+    MUTLASS_UNUSED(ptr);
+    MUTLASS_UNUSED(data);
+    MUTLASS_NOT_IMPLEMENTED();
+#endif
+  }
+};
+
+template <typename T>
+struct atomic_maximum {
+  MUTLASS_DEVICE
+  T operator()(T *ptr, T value) const {
+#if defined(__MUSA_ARCH__)
+    return atomicMax(ptr, value);
+#else
+    MUTLASS_UNUSED(ptr);
+    MUTLASS_UNUSED(value);
+    MUTLASS_NOT_IMPLEMENTED();
+    return 0;
+#endif
+  }
+};
+
+template <>
+struct atomic_maximum<float> {
+  MUTLASS_DEVICE
+  float operator()(float *ptr, float value) const {
+#if defined(__MUSA_ARCH__)
+    return !signbit(value) ?
+      __int_as_float(atomicMax((int*)ptr, __float_as_int(value))) :
+      __uint_as_float(atomicMin((unsigned int*)ptr, __float_as_uint(value)));
+#else
+    MUTLASS_UNUSED(ptr);
+    MUTLASS_UNUSED(value);
+    MUTLASS_NOT_IMPLEMENTED();
+    return 0;
+#endif
+  }
+};
+
+// is_atomic
+template <class Fn>
+struct is_atomic : platform::false_type {};
+template <class T>
+struct is_atomic<atomic_add<T>> : platform::true_type {};
+template <class T>
+struct is_atomic<atomic_maximum<T>> : platform::true_type {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/common.inl b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/common.inl
new file mode 100644
index 000000000..636e1aeab
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/common.inl
@@ -0,0 +1,87 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mute/atom/copy_atom.hpp"
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+
+template<int ThreadCount, class Element, int Alignment, class StrideType, int TileSizeMN, int TileSizeK, class CopyOperation>
+constexpr auto
+make_simt_tiled_copy() {
+  // Maximize the number of threads along the gmem major mode to promote coalesced reads
+  // While making sure our thread layout tiles the threadblock tile evenly
+
+  if constexpr (mutlass::gemm::detail::is_k_major<StrideType>()) {
+    // K major thread layout for K major gmem
+    constexpr int threads_major = TileSizeK   / Alignment;
+    constexpr int threads_minor = ThreadCount / threads_major;
+    static_assert(threads_major > 0);
+    static_assert(ThreadCount % threads_major == 0);
+    static_assert(threads_minor == 0 || (TileSizeMN % threads_minor == 0));
+    return make_tiled_copy(
+      Copy_Atom<CopyOperation, Element>{},
+      Layout<Shape <Int<threads_minor>, Int<threads_major>>,
+             Stride<Int<threads_major>,                _1>>{},
+      Layout<Shape<_1, Int<Alignment>>>{});
+  }
+  else if constexpr (mutlass::gemm::detail::is_mn_major<StrideType>()) {
+    // MN major thread layout for MN major gmem
+    constexpr int threads_major = TileSizeMN  / Alignment;
+    constexpr int threads_minor = ThreadCount / threads_major;
+    static_assert(threads_major > 0);
+    static_assert(ThreadCount % threads_major == 0);
+    static_assert(threads_minor == 0 || (TileSizeK % threads_minor == 0));
+    return make_tiled_copy(
+      Copy_Atom<CopyOperation, Element>{},
+      Layout<Shape <Int<threads_major>, Int<threads_minor>>,
+             Stride<                _1, Int<threads_major>>>{},
+      Layout<Shape<Int<Alignment>, _1>>{});
+  }
+  else {
+    static_assert(mute::is_void_v<Element>, "Unsupported gmem layout for automatic gmem tiled copy builder.");
+  }
+}
+
+} // namespace detail
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::collective
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp22_gemm_builder.inl b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp22_gemm_builder.inl
new file mode 100644
index 000000000..1bdfd65b7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp22_gemm_builder.inl
@@ -0,0 +1,410 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mute/atom/mma_atom.hpp"
+#include "mute/atom/copy_atom.hpp"
+
+#include "mutlass/mutlass.h"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/arch/arch.h"
+#include "mutlass/arch/mma.h"
+#include "mutlass/layout/layout.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+#include "mutlass/gemm/collective/collective_mma.hpp"
+#include "mutlass/gemm/collective/builders/common.inl"
+
+namespace mutlass::gemm::collective {
+
+namespace detail {
+
+template <class Element, class StrideAB>
+constexpr auto make_mp22_smem_atom_layout() {
+  constexpr int size = sizeof(Element);
+
+  using half_inst_mn = Int<16>;
+  using half_inst_k = Int<16/size>;
+  using inst_k = Int<32/size>;
+
+  if constexpr (gemm::detail::is_mn_major<StrideAB>()) {
+    using shape = Shape<half_inst_mn, inst_k>;
+    return make_ordered_layout(shape{}, Step<_0, _1>{});
+  } else {
+    using shape = Shape<half_inst_mn, Shape<half_inst_k, _2>>;
+    return make_ordered_layout(shape{}, Step<_1, Step<_0, _2>>{});
+  }
+};
+
+template <
+  class ElementA, class StrideA,
+  class ElementB, class StrideB>
+constexpr auto mp22_mma_operation_select()
+{
+  static_assert(is_same_v<ElementA, ElementB>, "ElementA and ElementB must be same");
+  static_assert((gemm::detail::is_mn_major<StrideA>() || gemm::detail::is_k_major<StrideA>()) &&
+                (gemm::detail::is_mn_major<StrideB>() || gemm::detail::is_k_major<StrideB>()),
+                 "StrideA/B must be MN or K major");
+  // tfloat32
+  if constexpr (is_same_v<ElementA, tfloat32_t>) {
+    if constexpr (gemm::detail::is_mn_major<StrideA>()) {
+      if constexpr (gemm::detail::is_k_major<StrideB>()) {
+        return MP22_32x32x8_F32TF32TF32F32_NN{};
+      }
+      else {
+        return MP22_32x32x8_F32TF32TF32F32_NT{};
+      }
+    } else {
+      if constexpr (gemm::detail::is_k_major<StrideB>()) {
+        return MP22_32x32x8_F32TF32TF32F32_TN{};
+      }
+      else {
+        return MP22_32x32x8_F32TF32TF32F32_TT{};
+      }
+    }
+  }
+  // half
+  else if constexpr (is_same_v<ElementA, half_t>) {
+    if constexpr (gemm::detail::is_mn_major<StrideA>()) {
+      if constexpr (gemm::detail::is_k_major<StrideB>()) {
+        return MP22_32x32x16_F32F16F16F32_NN{};
+      }
+      else {
+        return MP22_32x32x16_F32F16F16F32_NT{};
+      }
+    } else {
+      if constexpr (gemm::detail::is_k_major<StrideB>()) {
+        return MP22_32x32x16_F32F16F16F32_TN{};
+      }
+      else {
+        return MP22_32x32x16_F32F16F16F32_TT{};
+      }
+    }
+  }
+  // bfloat16
+  else if constexpr (is_same_v<ElementA, bfloat16_t>) {
+    if constexpr (gemm::detail::is_mn_major<StrideA>()) {
+      if constexpr (gemm::detail::is_k_major<StrideB>()) {
+        return MP22_32x32x16_F32BF16BF16F32_NN{};
+      }
+      else {
+        return MP22_32x32x16_F32BF16BF16F32_NT{};
+      }
+    } else {
+      if constexpr (gemm::detail::is_k_major<StrideB>()) {
+        return MP22_32x32x16_F32BF16BF16F32_TN{};
+      }
+      else {
+        return MP22_32x32x16_F32BF16BF16F32_TT{};
+      }
+    }
+  }
+  // int8
+  else if constexpr (is_same_v<ElementA, int8_t>) {
+    if constexpr (gemm::detail::is_mn_major<StrideA>()) {
+      if constexpr (gemm::detail::is_k_major<StrideB>()) {
+        return MP22_32x32x32_S32S8S8S32_NN{};
+      }
+      else {
+        return MP22_32x32x32_S32S8S8S32_NT{};
+      }
+    } else {
+      if constexpr (gemm::detail::is_k_major<StrideB>()) {
+        return MP22_32x32x32_S32S8S8S32_TN{};
+      }
+      else {
+        return MP22_32x32x32_S32S8S8S32_TT{};
+      }
+    }
+  } else {
+    static_assert(sizeof(ElementA) == 0, "No eligible MMA operator for request configuration.");
+  }
+}
+
+template <
+  class TileShape_MNK,
+  class ElementA,
+  class ElementB,
+  class ElementAccumulator>
+constexpr auto
+mp22_make_simt_tiled_mma() {
+  static_assert(sizeof(ElementA) == 4 && sizeof(ElementB) == 4);
+  static_assert(sizeof(ElementAccumulator) == 4);
+
+  constexpr int TileM = size<0>(TileShape_MNK{});
+  constexpr int TileN = size<1>(TileShape_MNK{});
+  constexpr int TileK = size<2>(TileShape_MNK{});
+
+  constexpr int MaxAccumPerThread = 64;
+  constexpr int MaxFmaInstInHotloop = 256;
+  constexpr int PreferedLdsWidth = 4;
+
+  constexpr int ThreadCountForOccupancy = mute::max(TileM * TileN / MaxAccumPerThread, NumThreadsPerWarpBeforeMP31);
+  constexpr int ThreadCount =
+    ThreadCountForOccupancy * mute::max(TileM * TileN * TileK / (ThreadCountForOccupancy * MaxFmaInstInHotloop), 1);
+
+  constexpr int MaxLdsTile = ThreadCount == 512 ? 128 : 64;
+  static_assert(ThreadCount * PreferedLdsWidth <= (TileM * TileN), "Too small tile shape");
+  if constexpr (TileM > TileN) {
+    constexpr int ThreadAtomM = TileM > MaxLdsTile ? (MaxLdsTile / PreferedLdsWidth) : (TileM / PreferedLdsWidth);
+    constexpr int ThreadAtomN = ThreadCount / ThreadAtomM;
+    constexpr int ValueN = (TileN / ThreadAtomN) > PreferedLdsWidth ? PreferedLdsWidth : (TileN / ThreadAtomN);
+
+    using AtomLayout = Layout<Shape<Int<ThreadAtomM>, Int<ThreadAtomN>, _1>>;
+    using PermuteLayout = mute::Tile<
+            Layout<Shape<Int<ThreadAtomM>, Int<PreferedLdsWidth>>, Stride<          Int<ValueN>, _1>>,
+            Layout<Shape<Int<ThreadAtomN>,           Int<ValueN>>, Stride<Int<PreferedLdsWidth>, _1>>,
+            X>;
+
+    return TiledMMA<MMA_Atom<UniversalFMA<
+                                ElementAccumulator,
+                                ElementA,
+                                ElementB,
+                                ElementAccumulator>>,
+                    AtomLayout, PermuteLayout>{};
+  } else {
+    constexpr int ThreadAtomN = TileN > MaxLdsTile ? (MaxLdsTile / PreferedLdsWidth) : (TileN / PreferedLdsWidth);
+    constexpr int ThreadAtomM = ThreadCount / ThreadAtomN;
+    constexpr int ValueM = (TileM / ThreadAtomM) > PreferedLdsWidth ? PreferedLdsWidth : (TileM / ThreadAtomM);
+    using AtomLayout = Layout<Shape<Int<ThreadAtomM>, Int<ThreadAtomN>, _1>>;
+
+    using PermuteLayout = mute::Tile<
+            Layout<Shape<Int<ThreadAtomM>,           Int<ValueM>>, Stride<          Int<ValueM>, _1>>,
+            Layout<Shape<Int<ThreadAtomN>, Int<PreferedLdsWidth>>, Stride<Int<PreferedLdsWidth>, _1>>,
+            X>;
+
+    return TiledMMA<MMA_Atom<UniversalFMA<
+                                ElementAccumulator,
+                                ElementA,
+                                ElementB,
+                                ElementAccumulator>>,
+                    AtomLayout, PermuteLayout>{};
+  }
+}
+
+template <class TileShape_MNK, class MmaOp>
+constexpr auto
+mp22_make_tensorop_tiled_mma() {
+  constexpr int TileM = size<0>(TileShape_MNK{});
+  constexpr int TileN = size<1>(TileShape_MNK{});
+
+  constexpr int MmaOpShapeM = size<0>(typename MMA_Traits<MmaOp>::Shape_MNK{});
+  constexpr int MmaOpShapeN = size<1>(typename MMA_Traits<MmaOp>::Shape_MNK{});
+
+  static_assert(TileM % MmaOpShapeM == 0 && TileN % MmaOpShapeN == 0, "Invalid TileShape");
+
+  constexpr int MaxWarpTile = 128;
+
+  constexpr int ValueLayoutM = mute::min(MaxWarpTile, TileM) / MmaOpShapeM;
+  constexpr int ValueLayoutN = mute::min(MaxWarpTile, TileN) / MmaOpShapeN;
+
+  constexpr int AtomTileM = ValueLayoutM * MmaOpShapeM;
+  constexpr int AtomTileN = ValueLayoutN * MmaOpShapeN;
+
+  constexpr int AtomExtendM = TileM / AtomTileM;
+  constexpr int AtomExtendN = TileN / AtomTileN;
+
+  using AtomLayout = decltype(make_layout(Shape<Int<AtomExtendM>, Int<AtomExtendN>, _1>{}, LayoutLeft{}));
+
+  using PermuteLayoutK = X;
+
+  using ShapeM0 = Int<MmaOpShapeM>;
+  using ShapeM1 = Int<size<0>(AtomLayout{})>;
+  using ShapeM2 = Int<ValueLayoutM>;
+
+  using StrideM0 = _1;
+  using StrideM1 = Int<MmaOpShapeM * ValueLayoutM>;
+  using StrideM2 = Int<MmaOpShapeM>;
+
+  using PermuteLayoutM = Layout<Shape < ShapeM0,  ShapeM1,  ShapeM2>,
+                                Stride<StrideM0, StrideM1, StrideM2>>;
+
+  using ShapeN0 = Int<MmaOpShapeN>;
+  using ShapeN1 = Int<size<1>(AtomLayout{})>;
+  using ShapeN2 = Int<ValueLayoutN>;
+
+  using StrideN0 = _1;
+  using StrideN1 = Int<MmaOpShapeN * ValueLayoutN>;
+  using StrideN2 = Int<MmaOpShapeN>;
+
+  using PermuteLayoutN = Layout<Shape < ShapeN0,  ShapeN1,  ShapeN2>,
+                                Stride<StrideN0, StrideN1, StrideN2>>;
+
+  using PermuteLayout = mute::Tile<PermuteLayoutM, PermuteLayoutN, PermuteLayoutK>;
+
+  return TiledMMA<MMA_Atom<MmaOp>, AtomLayout, PermuteLayout>{};
+}
+
+} // namespace detail
+
+
+template <
+  class ArchTag,
+  class ElementA,
+  class GmemLayoutA,
+  int AlignmentA,
+  class ElementB,
+  class GmemLayoutB,
+  int AlignmentB,
+  class ElementAccumulator,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class StageCountType,
+  class KernelScheduleType
+>
+struct CollectiveBuilder<
+  ArchTag,
+  arch::OpClassSimt,
+  ElementA,
+  GmemLayoutA,
+  AlignmentA,
+  ElementB,
+  GmemLayoutB,
+  AlignmentB,
+  ElementAccumulator,
+  TileShape_MNK,
+  ClusterShape_MNK,
+  StageCountType,
+  KernelScheduleType
+> {
+  static_assert(is_static<TileShape_MNK>::value, "TileShape must be static");
+  static_assert(is_static<ClusterShape_MNK>::value, "ClusterShape must be static");
+
+  static constexpr int BlockM = shape<0>(TileShape_MNK{});
+  static constexpr int BlockN = shape<1>(TileShape_MNK{});
+  static constexpr int BlockK = shape<2>(TileShape_MNK{});
+
+  using TiledMma = decltype(detail::mp22_make_simt_tiled_mma<TileShape_MNK, ElementA, ElementB, ElementAccumulator>());
+
+  static constexpr int ThreadCount = size(TiledMma{});
+  using GmemTiledCopyA = decltype(detail::make_simt_tiled_copy<
+                                    ThreadCount, ElementA, AlignmentA, TagToStrideA_t<GmemLayoutA>,
+                                    BlockM, BlockK,
+                                    UniversalCopy<uint_bit_t<AlignmentA*sizeof_bits_v<ElementA>>>>());
+  using GmemTiledCopyB = decltype(detail::make_simt_tiled_copy<
+                                    ThreadCount, ElementB, AlignmentB, TagToStrideB_t<GmemLayoutB>,
+                                    BlockN, BlockK,
+                                    UniversalCopy<uint_bit_t<AlignmentB*sizeof_bits_v<ElementB>>>>());
+
+  using SmemLayoutAtomA = Layout<Shape <Int<BlockM>, Int<BlockK>>,
+                                 Stride<         _1, Int<BlockM>>>;
+  using SmemLayoutAtomB = Layout<Shape <Int<BlockN>, Int<BlockK>>,
+                                 Stride<         _1, Int<BlockM>>>;
+
+  using SmemCopyAtomA = Copy_Atom<DefaultCopy, ElementA>;
+  using SmemCopyAtomB = Copy_Atom<DefaultCopy, ElementB>;
+  using DispatchPolicy = MainloopMp22TwoStage;
+
+  using CollectiveOp = collective::CollectiveMma<
+    DispatchPolicy, TileShape_MNK,
+    ElementA, TagToStrideA_t<GmemLayoutA>,
+    ElementB, TagToStrideB_t<GmemLayoutB>,
+    TiledMma,
+    GmemTiledCopyA, SmemLayoutAtomA, SmemCopyAtomA, mute::identity,  // A
+    GmemTiledCopyB, SmemLayoutAtomB, SmemCopyAtomB, mute::identity   // B
+  >;
+};
+
+template <
+  class ElementA,
+  class GmemLayoutA,
+  int AlignmentA,
+  class ElementB,
+  class GmemLayoutB,
+  int AlignmentB,
+  class ElementAccumulator,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class StageCountType,
+  class KernelScheduleType
+>
+struct CollectiveBuilder<
+  arch::Mp22,
+  arch::OpClassTensorOp,
+  ElementA,
+  GmemLayoutA,
+  AlignmentA,
+  ElementB,
+  GmemLayoutB,
+  AlignmentB,
+  ElementAccumulator,
+  TileShape_MNK,
+  ClusterShape_MNK,
+  StageCountType,
+  KernelScheduleType
+> {
+  static_assert(is_static<TileShape_MNK>::value, "TileShape must be static");
+  static_assert(is_static<ClusterShape_MNK>::value, "ClusterShape must be static");
+
+  static constexpr int BlockM = shape<0>(TileShape_MNK{});
+  static constexpr int BlockN = shape<1>(TileShape_MNK{});
+  static constexpr int BlockK = shape<2>(TileShape_MNK{});
+
+  using StrideA = TagToStrideA_t<GmemLayoutA>;
+  using StrideB = TagToStrideB_t<GmemLayoutB>;
+
+  // For fp32 types, map to tf32 MMA value type
+  using MmaElementA = mute::conditional_t<mute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
+  using MmaElementB = mute::conditional_t<mute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
+
+  using MmaOp = decltype(detail::mp22_mma_operation_select<MmaElementA, StrideA, MmaElementB, StrideB>());
+
+  using TiledMma = decltype(detail::mp22_make_tensorop_tiled_mma<TileShape_MNK, MmaOp>());
+
+  static constexpr int ThreadCount = size(TiledMma{});
+
+  // A
+  using SmemLayoutAtomA = decltype(detail::make_mp22_smem_atom_layout<MmaElementA, StrideA>());
+  using SmemCopyAtomA = Copy_Atom<DefaultCopy, MmaElementA>;
+  using GmemTiledCopyA = decltype(detail::make_simt_tiled_copy<
+                                    ThreadCount, MmaElementA, AlignmentA, StrideA,
+                                    BlockM, BlockK,
+                                    UniversalCopy<uint_bit_t<AlignmentA*sizeof_bits_v<MmaElementA>>>>());
+  // B
+  using SmemLayoutAtomB = decltype(detail::make_mp22_smem_atom_layout<MmaElementB, StrideB>());
+  using SmemCopyAtomB = Copy_Atom<DefaultCopy, MmaElementB>;
+  using GmemTiledCopyB = decltype(detail::make_simt_tiled_copy<
+                                    ThreadCount, MmaElementB, AlignmentB, StrideB,
+                                    BlockN, BlockK,
+                                    UniversalCopy<uint_bit_t<AlignmentB*sizeof_bits_v<MmaElementB>>>>());
+  using DispatchPolicy = MainloopMp22TwoStage;
+
+  using CollectiveOp = collective::CollectiveMma<
+    DispatchPolicy, TileShape_MNK,
+    MmaElementA, StrideA,
+    MmaElementB, StrideB,
+    TiledMma,
+    GmemTiledCopyA, SmemLayoutAtomA, SmemCopyAtomA, mute::identity,  // A
+    GmemTiledCopyB, SmemLayoutAtomB, SmemCopyAtomB, mute::identity   // B
+  >;
+};
+
+} // namespace mutlass::gemm::collective
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp31_sqmma_builder.inl b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp31_sqmma_builder.inl
new file mode 100644
index 000000000..e4065efb8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp31_sqmma_builder.inl
@@ -0,0 +1,236 @@
+#pragma once
+
+#include "mutlass/gemm/collective/builders/mp31_sqmma_common.inl"
+#include "mutlass/gemm/dispatch_policy.hpp"
+
+#include "mute/arch/mma_mp31.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::collective {
+
+namespace detail {
+
+template <class TileShape_MNK, class SqmmaOp>
+constexpr auto
+mp31_make_sqmma_tiled_mma() {
+  constexpr int TileM = size<0>(TileShape_MNK{});
+  constexpr int TileN = size<1>(TileShape_MNK{});
+
+  constexpr int InstM = size<0>(typename MMA_Traits<SqmmaOp>::Shape_MNK{});
+  constexpr int InstN = size<1>(typename MMA_Traits<SqmmaOp>::Shape_MNK{});
+
+  constexpr int AtomM = TileM / InstM;
+  constexpr int AtomN = TileN / InstN;
+
+  static_assert(TileM % InstM == 0 && TileN % InstN == 0, "Invalid Tile shape");
+  using AtomLayout = decltype(make_layout(Shape<Int<AtomM>, Int<AtomN>, Int<1>>{}, LayoutRight{}));
+
+  return mute::make_tiled_mma(SqmmaOp{}, AtomLayout{});
+}
+
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template<int CapacityBytes, class ElementA, class ElementB, class TileShapeMNK, int stages>
+constexpr int
+compute_stage_count_or_override(StageCount<stages> stage_count) {
+  return stages;
+}
+
+template<int CapacityBytes, class ElementA, class ElementB, class TileShapeMNK, int stages>
+constexpr int
+compute_stage_count_or_override(mute::Int<stages> stage_count) {
+  return stages;
+}
+
+template<int CapacityBytes, class ElementA, class ElementB, class TileShapeMNK, int carveout_bytes>
+constexpr int
+compute_stage_count_or_override(StageCountAutoCarveout<carveout_bytes> stage_count) {
+  constexpr auto a_bits = mute::sizeof_bits_v<ElementA>;
+  constexpr auto b_bits = mute::sizeof_bits_v<ElementB>;
+  constexpr int stage_bytes =
+    mutlass::bits_to_bytes(a_bits * size<0>(TileShapeMNK{}) * size<2>(TileShapeMNK{})) +
+    mutlass::bits_to_bytes(b_bits * size<1>(TileShapeMNK{}) * size<2>(TileShapeMNK{}));
+
+  constexpr int MaximumStages = (CapacityBytes - carveout_bytes) / stage_bytes;
+
+  static_assert(MaximumStages >= 2, "Mainloop need at least 2-stage");
+  // Currently we just use 2-stage pipeline
+  return 2;
+}
+
+} // namespace detail
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// SQMMA_TME_SS
+template <
+  class ElementA,
+  class GmemLayoutATag,
+  int AlignmentA,
+  class ElementB,
+  class GmemLayoutBTag,
+  int AlignmentB,
+  class ElementAccumulator,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class StageCountType,
+  class KernelScheduleType
+>
+struct CollectiveBuilder<
+  arch::Mp31,
+  arch::OpClassTensorOp,
+  ElementA,
+  GmemLayoutATag,
+  AlignmentA,
+  ElementB,
+  GmemLayoutBTag,
+  AlignmentB,
+  ElementAccumulator,
+  TileShape_MNK,
+  ClusterShape_MNK,
+  StageCountType,
+  KernelScheduleType,
+  mute::enable_if_t<mute::is_same_v<KernelScheduleType, KernelTme>>
+> {
+  static_assert(is_static<TileShape_MNK>::value);
+  static_assert(is_static<ClusterShape_MNK>::value);
+
+  static_assert(detail::is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, detail::tme_alignment_bytes>(),
+      "Should meet TME alignment requirement\n");
+
+  // For fp32 types, map to tf32 MMA value type
+  using ElementAMma = mute::conditional_t<mute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
+  using ElementBMma = mute::conditional_t<mute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
+
+  static constexpr mute::TCE::Major SqmmaMajorA = detail::sqmma_ss_tag_to_major_A<GmemLayoutATag>();
+  static constexpr mute::TCE::Major SqmmaMajorB = detail::sqmma_ss_tag_to_major_B<GmemLayoutBTag>();
+
+  using SqmmaOp = decltype(mute::MP31::SQMMA::ss_op_selector<ElementA, ElementB, ElementAccumulator, TileShape_MNK, SqmmaMajorA, SqmmaMajorB>());
+
+  using TiledMma = decltype(detail::mp31_make_sqmma_tiled_mma<TileShape_MNK, SqmmaOp>());
+
+  using GmemTiledCopyA = MP31_TME_LOAD;
+  using GmemTiledCopyB = MP31_TME_LOAD;
+
+  using SmemLayoutAtomA = decltype(detail::ss_smem_selector_A<SqmmaMajorA, ElementAMma, SqmmaOp>());
+  using SmemLayoutAtomB = decltype(detail::ss_smem_selector_B<SqmmaMajorB, ElementBMma, SqmmaOp>());
+
+  using SmemCopyAtomA = void;
+  using SmemCopyAtomB = void;
+
+  static constexpr int PipelineStages = detail::compute_stage_count_or_override<detail::mp31_smem_capacity_bytes, ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{});
+  using DispatchPolicy = MainloopMp31TmeSqmma<PipelineStages>;
+
+  using ClusterShape_MNK_ = Shape<_1, _1, _1>;
+
+  using CollectiveOp = CollectiveMma<
+      DispatchPolicy,
+      TileShape_MNK,
+      ElementA,
+      TagToStrideA_t<GmemLayoutATag>,
+      ElementB,
+      TagToStrideB_t<GmemLayoutBTag>,
+      TiledMma,
+      GmemTiledCopyA,
+      SmemLayoutAtomA,
+      SmemCopyAtomA,
+      mute::identity,
+      GmemTiledCopyB,
+      SmemLayoutAtomB,
+      SmemCopyAtomB,
+      mute::identity
+  >;
+};
+
+// SQMMA_TME_SS_GROUP_SCALING
+template <
+  class ElementA,
+  class GmemLayoutATag,
+  int AlignmentA,
+  class ElementB,
+  class GmemLayoutBTag,
+  int AlignmentB,
+  class ElementAccumulator,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class StageCountType,
+  int ScaleGranularityM,
+  int ScaleGranularityN,
+  int ScaleGranularityK
+>
+struct CollectiveBuilder<
+  arch::Mp31,
+  arch::OpClassTensorOp,
+  ElementA,
+  GmemLayoutATag,
+  AlignmentA,
+  ElementB,
+  GmemLayoutBTag,
+  AlignmentB,
+  ElementAccumulator,
+  TileShape_MNK,
+  ClusterShape_MNK,
+  StageCountType,
+  KernelTmeGroupScaledAccum<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>
+> {
+  using KernelScheduleType = KernelTmeGroupScaledAccum<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>;
+
+  static_assert(is_static<TileShape_MNK>::value);
+  static_assert(is_static<ClusterShape_MNK>::value);
+
+  static_assert(detail::is_aligned<ElementA, AlignmentA, ElementB, AlignmentB, detail::tme_alignment_bytes>(),
+      "Should meet TME alignment requirement\n");
+
+  static constexpr bool IsFP8Input = detail::is_input_fp8<ElementA, ElementB>();
+
+  static_assert(IsFP8Input, "KernelTmeGroupScaledAccum is only compatible with FP8 now.");
+
+  // For fp32 types, map to tf32 MMA value type
+  using ElementAMma = mute::conditional_t<mute::is_same_v<ElementA, float>, tfloat32_t, ElementA>;
+  using ElementBMma = mute::conditional_t<mute::is_same_v<ElementB, float>, tfloat32_t, ElementB>;
+
+  static constexpr mute::TCE::Major SqmmaMajorA = detail::sqmma_ss_tag_to_major_A<GmemLayoutATag>();
+  static constexpr mute::TCE::Major SqmmaMajorB = detail::sqmma_ss_tag_to_major_B<GmemLayoutBTag>();
+
+  using MaxInstructionM = Int<64>;
+  using SqmmaOp = decltype(mute::MP31::SQMMA::ss_op_selector<ElementA, ElementB, ElementAccumulator, TileShape_MNK, SqmmaMajorA, SqmmaMajorB, MaxInstructionM>());
+  using TiledMma = decltype(detail::mp31_make_sqmma_tiled_mma<TileShape_MNK, SqmmaOp>());
+
+  using GmemTiledCopyA = MP31_TME_LOAD;
+  using GmemTiledCopyB = MP31_TME_LOAD;
+
+  using SmemLayoutAtomA = decltype(detail::ss_smem_selector_A<SqmmaMajorA, ElementAMma, SqmmaOp>());
+  using SmemLayoutAtomB = decltype(detail::ss_smem_selector_B<SqmmaMajorB, ElementBMma, SqmmaOp>());
+
+  using SmemCopyAtomA = void;
+  using SmemCopyAtomB = void;
+
+  static constexpr int PipelineStages = detail::compute_stage_count_or_override<detail::mp31_smem_capacity_bytes, ElementAMma, ElementBMma, TileShape_MNK>(StageCountType{});
+
+  using DispatchPolicy = MainloopMp31TmeSqmmaBlockScalingFP8<PipelineStages, KernelScheduleType, ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>;
+
+  using CollectiveOp = CollectiveMma<
+      DispatchPolicy,
+      TileShape_MNK,
+      ElementA,
+      TagToStrideA_t<GmemLayoutATag>,
+      ElementB,
+      TagToStrideB_t<GmemLayoutBTag>,
+      TiledMma,
+      GmemTiledCopyA,
+      SmemLayoutAtomA,
+      SmemCopyAtomA,
+      mute::identity,
+      GmemTiledCopyB,
+      SmemLayoutAtomB,
+      SmemCopyAtomB,
+      mute::identity
+  >;
+};
+
+} // namespace mutlass::gemm::collective
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp31_sqmma_common.inl b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp31_sqmma_common.inl
new file mode 100644
index 000000000..2a16ffd2f
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/builders/mp31_sqmma_common.inl
@@ -0,0 +1,90 @@
+#pragma once
+
+#include "mutlass/gemm/gemm.h"
+#include "mute/atom/mma_traits_mp31_sqmma.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+//
+// Some named constants
+//
+constexpr int tme_alignment_bytes = 4;
+constexpr int mp31_smem_capacity_bytes = 196608;
+
+template <class LayoutA>
+constexpr mute::TCE::Major
+sqmma_ss_tag_to_major_A() {
+  if constexpr (mutlass::gemm::detail::is_mn_major_A<LayoutA>()) {
+    return mute::TCE::Major::MN;
+  }
+  else {
+    return mute::TCE::Major::K;
+  }
+}
+
+template <class LayoutB>
+constexpr mute::TCE::Major
+sqmma_ss_tag_to_major_B() {
+  if constexpr (mutlass::gemm::detail::is_mn_major_B<LayoutB>()) {
+    return mute::TCE::Major::MN;
+  }
+  else {
+    return mute::TCE::Major::K;
+  }
+}
+
+template <mute::TCE::Major major, class ElementType, class SqmmaOp>
+MUTE_HOST_DEVICE constexpr
+auto
+ss_smem_selector_A()
+{
+  using AtomOpTraits = MMA_Traits<SqmmaOp>;
+  using AtomOpShape = typename AtomOpTraits::Shape_MNK;
+
+  using AtomM = decltype(get<0>(AtomOpShape{}));
+  using AtomK = decltype(get<2>(AtomOpShape{}));
+
+  return mute::MP31::SQMMA::make_canonical_gemm_smem_atom_layout<ElementType, major, AtomM, AtomK>();
+}
+
+template <mute::TCE::Major major, class ElementType, class SqmmaOp>
+MUTE_HOST_DEVICE constexpr
+auto
+ss_smem_selector_B()
+{
+  using AtomOpTraits = MMA_Traits<SqmmaOp>;
+  using AtomOpShape = typename AtomOpTraits::Shape_MNK;
+
+  using AtomN = decltype(get<1>(AtomOpShape{}));
+  using AtomK = decltype(get<2>(AtomOpShape{}));
+
+  return mute::MP31::SQMMA::make_canonical_gemm_smem_atom_layout<ElementType, major, AtomN, AtomK>();
+}
+
+template <class ElementA, int AlignmentA, class ElementB, int AlignmentB, int RequiredAlignment>
+constexpr bool
+is_aligned() {
+  return ((mute::sizeof_bits_v<ElementA> * AlignmentA / 8) % RequiredAlignment == 0) &&
+         ((mute::sizeof_bits_v<ElementB> * AlignmentB / 8) % RequiredAlignment == 0);
+}
+
+template <class ElementA, class ElementB>
+constexpr bool
+is_input_fp8() {
+  return ((mute::is_same_v<ElementA, float_e4m3_t> || mute::is_same_v<ElementA, float_e5m2_t>) &&
+          (mute::is_same_v<ElementB, float_e4m3_t> || mute::is_same_v<ElementB, float_e5m2_t>));
+}
+
+} // namespace detail
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/collective_builder.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/collective_builder.hpp
new file mode 100644
index 000000000..57db05735
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/collective_builder.hpp
@@ -0,0 +1,88 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+#include "mutlass/gemm/collective/collective_mma.hpp"
+
+namespace mutlass::gemm::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Used to specify stage counts or dispatch to automatic computation of stage count
+template<int num_stages>
+struct StageCount { static constexpr int value = num_stages; };
+
+template<int carveout_bytes>
+struct StageCountAutoCarveout { static constexpr int bytes = carveout_bytes; };
+
+using StageCountAuto = StageCountAutoCarveout<0>;
+
+// Used to automatically let the builder pick the kernel schedule.
+// Can be overridden with kernel schedule tags in mutlass/gemm/dispatch_policy.hpp
+struct KernelScheduleAuto {};
+
+// Used to automatically let the builder pick the TiledMMA permute layout.
+// Can be overridden with mute::Tile
+struct PermuteLayoutAuto {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class ArchTag,
+  class OpClass,
+  class ElementA,
+  class GmemLayoutA,
+  int AlignmentA,
+  class ElementB,
+  class GmemLayoutB,
+  int AlignmentB,
+  class ElementAccumulator,
+  class TileShape_MNK,
+  class ClusterShape_MNK,
+  class StageCountType,
+  class KernelScheduleType,
+  class Enable = void
+>
+struct CollectiveBuilder {
+  static_assert(sizeof(ElementA) == 0, "Could not build a collective for given parameters.");
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "mutlass/gemm/collective/builders/mp22_gemm_builder.inl"
+#include "mutlass/gemm/collective/builders/mp31_sqmma_builder.inl"
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/collective_mma.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/collective_mma.hpp
new file mode 100644
index 000000000..2553d4abd
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/collective_mma.hpp
@@ -0,0 +1,72 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/detail/dependent_false.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::collective {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class DispatchPolicy,
+  class TileShape,
+  class ElementA,
+  class StrideA,
+  class ElementB,
+  class StrideB,
+  class TiledMma,
+  class GmemTiledCopyA,
+  class SmemLayoutAtomA,
+  class SmemCopyAtomA,
+  class TransformA,
+  class GmemTiledCopyB,
+  class SmemLayoutAtomB,
+  class SmemCopyAtomB,
+  class TransformB
+>
+struct CollectiveMma {
+  static_assert(mutlass::detail::dependent_false<ElementA> == 0, "Could not find a mainloop specialization.");
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "mutlass/gemm/collective/mp22_mma_twostage.hpp"
+#include "mutlass/gemm/collective/mp31_mma_tme_sqmma_ss.hpp"
+#include "mutlass/gemm/collective/mp31_mma_tme_sqmma_ss_scaling.hpp"
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp22_mma_twostage.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp22_mma_twostage.hpp
new file mode 100644
index 000000000..d0885eac8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp22_mma_twostage.hpp
@@ -0,0 +1,649 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/trace.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+
+#include "mute/algorithm/functional.hpp"
+#include "mute/atom/mma_atom.hpp"
+#include "mute/algorithm/gemm.hpp"
+#include "mute/atom/mma_atom.hpp"
+#include "mute/tensor_predicate.hpp"
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::collective {
+using namespace mute;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class TileShape_,
+  class ElementA_,
+  class StrideA_,
+  class ElementB_,
+  class StrideB_,
+  class TiledMma_,
+  class GmemTiledCopyA_,
+  class SmemLayoutAtomA_,
+  class SmemCopyAtomA_,
+  class TransformA_,
+  class GmemTiledCopyB_,
+  class SmemLayoutAtomB_,
+  class SmemCopyAtomB_,
+  class TransformB_>
+struct CollectiveMma<
+    MainloopMp22TwoStageUnpredicated,
+    TileShape_,
+    ElementA_,
+    StrideA_,
+    ElementB_,
+    StrideB_,
+    TiledMma_,
+    GmemTiledCopyA_,
+    SmemLayoutAtomA_,
+    SmemCopyAtomA_,
+    TransformA_,
+    GmemTiledCopyB_,
+    SmemLayoutAtomB_,
+    SmemCopyAtomB_,
+    TransformB_>
+{
+  //
+  // Type Aliases
+  //
+  using DispatchPolicy = MainloopMp22TwoStageUnpredicated;
+  using TileShape = TileShape_;
+  using ElementA = ElementA_;
+  using StrideA = StrideA_;
+  using ElementB = ElementB_;
+  using StrideB = StrideB_;
+  using TiledMma = TiledMma_;
+  using ElementAccumulator = typename TiledMma::ValTypeC;
+  using GmemTiledCopyA = GmemTiledCopyA_;
+  using GmemTiledCopyB = GmemTiledCopyB_;
+  using SmemLayoutAtomA = SmemLayoutAtomA_;
+  using SmemLayoutAtomB = SmemLayoutAtomB_;
+  using SmemCopyAtomA = SmemCopyAtomA_;
+  using SmemCopyAtomB = SmemCopyAtomB_;
+  using TransformA = TransformA_;
+  using TransformB = TransformB_;
+  using ArchTag = typename DispatchPolicy::ArchTag;
+
+  static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  using SmemLayoutA = decltype(tile_to_shape(
+      SmemLayoutAtomA{},
+      make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}))));
+  using SmemLayoutB = decltype(tile_to_shape(
+      SmemLayoutAtomB{},
+      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}))));
+
+  static constexpr int SmemAlignmentBytes = 128;
+
+  struct SharedStorage
+  {
+    mute::array_aligned<ElementA, mute::cosize_v<SmemLayoutA>> smem_a;
+    mute::array_aligned<ElementB, mute::cosize_v<SmemLayoutB>> smem_b;
+  };
+
+  // Host side kernel arguments
+  struct Arguments {
+    ElementA const* ptr_A;
+    StrideA dA;
+    ElementB const* ptr_B;
+    StrideB dB;
+  };
+
+  // Device side kernel params
+  using Params = Arguments;
+
+  //
+  // Methods
+  //
+
+  CollectiveMma() = default;
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) {
+    (void) workspace;
+    return args;
+  }
+
+  template <class ProblemShape>
+  MUTLASS_HOST_DEVICE static bool
+  can_implement(
+    ProblemShape problem_shapes,
+    Arguments const& args) {
+    int TileShapeM = int(size<0>(TileShape{}));
+    int TileShapeN = int(size<1>(TileShape{}));
+    int TileShapeK = int(size<2>(TileShape{}));
+
+    int problem_m = int(size<0>(problem_shapes));
+    int problem_n = int(size<1>(problem_shapes));
+    int problem_k = int(size<2>(problem_shapes));
+
+    bool implementable = true;
+    // Check alignment for all problem sizes
+    implementable = implementable && (problem_m % TileShapeM == 0 &&
+                                      problem_n % TileShapeN == 0 &&
+                                      problem_k % TileShapeK == 0);
+
+    if (!implementable) {
+      MUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements.\n");
+    }
+    return implementable;
+  }
+
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  template <
+    class FrgTensorD,
+    class TensorA,
+    class TensorB,
+    class FrgTensorC,
+    class KTileIterator,
+    class ResidueMNK
+  >
+  MUTLASS_DEVICE void
+  operator() (
+      FrgTensorD &accum,
+      TensorA gA,
+      TensorB gB,
+      FrgTensorC const &src_accum,
+      KTileIterator k_tile_iter, int k_tile_count,
+      ResidueMNK residue_mnk,
+      int thread_idx,
+      char *smem_buf)
+  {
+    using namespace mute;
+
+    (void)residue_mnk;
+
+    static_assert(is_rmem<FrgTensorD>::value, "D tensor must be rmem resident.");
+    static_assert(is_gmem<TensorA>::value, "A tensor must be gmem resident.");
+    static_assert(is_gmem<TensorB>::value, "B tensor must be gmem resident.");
+    static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
+    static_assert(rank(SmemLayoutA{}) == 2,
+      "MainloopTwoStage must not have a smem shape with a pipeline mode.");
+    static_assert(rank(SmemLayoutB{}) == 2,
+      "MainloopTwoStage must not have a smem shape with a pipeline mode.");
+
+    // Construct shared memory tiles
+    SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+    Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
+
+    // Partition the copying of A and B tiles across the threads
+    GmemTiledCopyA gmem_tiled_copy_a;
+    GmemTiledCopyB gmem_tiled_copy_b;
+    auto copy_a_thr = gmem_tiled_copy_a.get_slice(thread_idx);
+    auto copy_b_thr = gmem_tiled_copy_b.get_slice(thread_idx);
+
+    Tensor tAgA = copy_a_thr.partition_S(gA);                                  // (ACPY,ACPY_M,ACPY_K,k)
+    Tensor tAsA = copy_a_thr.partition_D(sA);                                  // (ACPY,ACPY_M,ACPY_K)
+    Tensor tBgB = copy_b_thr.partition_S(gB);                                  // (BCPY,BCPY_N,BCPY_K,k)
+    Tensor tBsB = copy_b_thr.partition_D(sB);                                  // (BCPY,BCPY_N,BCPY_K)
+
+    // Allocate the register tiles for double buffering -- same shape as partitioned data
+    Tensor tArA = make_fragment_like(tAsA);                                    // (ACPY,ACPY_M,ACPY_K)
+    Tensor tBrB = make_fragment_like(tBsB);                                    // (BCPY,BCPY_N,BCPY_K)
+
+    // Tile MMA compute thread partitions and allocate accumulators
+    TiledMma tiled_mma;
+    auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
+    Tensor tCrA  = thr_mma.partition_fragment_A(sA);                           // (MMA,MMA_M,MMA_K)
+    Tensor tCrB  = thr_mma.partition_fragment_B(sB);                           // (MMA,MMA_M,MMA_K)
+
+    MUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum));                     // MMA_M
+    MUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(src_accum));                 // MMA_M
+    MUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum));                     // MMA_N
+    MUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(src_accum));                 // MMA_N
+    MUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));                      // MMA_K
+
+    //
+    // Copy Atom retiling
+    //
+
+    auto thr_copy_A       = make_tiled_copy_A(SmemCopyAtomA{}, tiled_mma).get_thread_slice(thread_idx);
+    Tensor tCsA           = thr_copy_A.partition_S(sA);
+    Tensor tCrA_copy_view = thr_copy_A.retile_D(tCrA);
+    MUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));            // M
+
+    auto thr_copy_B       = make_tiled_copy_B(SmemCopyAtomB{}, tiled_mma).get_thread_slice(thread_idx);
+    Tensor tCsB           = thr_copy_B.partition_S(sB);
+    Tensor tCrB_copy_view = thr_copy_B.retile_D(tCrB);
+    MUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));            // N
+
+    //
+    // Prologue
+    //
+
+    // Copy gmem to rmem for the first k_tile
+    copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tArA);
+    copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBrB);
+    if (--k_tile_count > 0) ++k_tile_iter;
+    // Copy rmem to smem
+    copy(tArA, tAsA);
+    copy(tBrB, tBsB);
+    // Clear accumulators
+    __syncthreads();
+
+    // Load A, B smem->rmem for k=0
+    copy(tCsA(_,_,0), tCrA_copy_view(_,_,0));
+    copy(tCsB(_,_,0), tCrB_copy_view(_,_,0));
+    //
+    // Mainloop
+    //
+
+    // Size of the k-tiles's outer product mode (k)
+    auto K_BLOCK_MAX = size<2>(tCrA);
+
+    MUTLASS_PRAGMA_NO_UNROLL
+    while (k_tile_count > -1)
+    {
+      // Pipeline the outer products with a static for loop
+      for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block)
+      {
+        if (k_block == K_BLOCK_MAX - 1)
+        {
+          __syncthreads();
+
+          // Copy rmem to smem
+          copy(tArA, tAsA);
+          copy(tBrB, tBsB);
+          __syncthreads();
+        }
+
+        // Load A, B smem->rmem for k+1
+        int k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX;     // static
+        copy(tCsA(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next));
+        copy(tCsB(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next));
+        if (k_block == 0)
+        {
+          // Copy gmem to rmem
+          copy(gmem_tiled_copy_a, tAgA(_,_,_,*k_tile_iter), tArA);
+          copy(gmem_tiled_copy_b, tBgB(_,_,_,*k_tile_iter), tBrB);
+          if (--k_tile_count > 0) ++k_tile_iter;
+        }
+
+        // transform before compute
+        mute::transform(tCrA(_,_,k_block), TransformA{});
+        mute::transform(tCrB(_,_,k_block), TransformB{});
+
+        // Thread-level register gemm for k
+        // disambiguate gemm (shared with the namespace name)
+        mute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), src_accum);
+      });
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  class TileShape_,
+  class ElementA_,
+  class StrideA_,
+  class ElementB_,
+  class StrideB_,
+  class TiledMma_,
+  class GmemTiledCopyA_,
+  class SmemLayoutAtomA_,
+  class SmemCopyAtomA_,
+  class TransformA_,
+  class GmemTiledCopyB_,
+  class SmemLayoutAtomB_,
+  class SmemCopyAtomB_,
+  class TransformB_>
+struct CollectiveMma<
+    MainloopMp22TwoStage,
+    TileShape_,
+    ElementA_,
+    StrideA_,
+    ElementB_,
+    StrideB_,
+    TiledMma_,
+    GmemTiledCopyA_,
+    SmemLayoutAtomA_,
+    SmemCopyAtomA_,
+    TransformA_,
+    GmemTiledCopyB_,
+    SmemLayoutAtomB_,
+    SmemCopyAtomB_,
+    TransformB_>
+{
+  //
+  // Type Aliases
+  //
+  using DispatchPolicy = MainloopMp22TwoStage;
+  using TileShape = TileShape_;
+  using ElementA = ElementA_;
+  using StrideA = StrideA_;
+  using ElementB = ElementB_;
+  using StrideB = StrideB_;
+  using TiledMma = TiledMma_;
+  using ElementAccumulator = typename TiledMma::ValTypeC;
+  using GmemTiledCopyA = GmemTiledCopyA_;
+  using GmemTiledCopyB = GmemTiledCopyB_;
+  using SmemLayoutAtomA = SmemLayoutAtomA_;
+  using SmemLayoutAtomB = SmemLayoutAtomB_;
+  using SmemCopyAtomA = SmemCopyAtomA_;
+  using SmemCopyAtomB = SmemCopyAtomB_;
+  using TransformA = TransformA_;
+  using TransformB = TransformB_;
+  using ArchTag = typename DispatchPolicy::ArchTag;
+
+  static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  using SmemLayoutA = decltype(tile_to_shape(
+      SmemLayoutAtomA{},
+      make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}))));
+  using SmemLayoutB = decltype(tile_to_shape(
+      SmemLayoutAtomB{},
+      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}))));
+
+  static constexpr int SmemAlignmentBytes = 128;
+
+  struct SharedStorage
+  {
+    mute::array_aligned<ElementA, mute::cosize_v<SmemLayoutA>> smem_a;
+    mute::array_aligned<ElementB, mute::cosize_v<SmemLayoutB>> smem_b;
+  };
+
+  // Host side kernel arguments
+  struct Arguments {
+    ElementA const* ptr_A;
+    StrideA dA;
+    ElementB const* ptr_B;
+    StrideB dB;
+  };
+
+  // Device side kernel params
+  using Params = Arguments;
+
+  //
+  // Methods
+  //
+
+  CollectiveMma() = default;
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& _, Arguments const& args, void* workspace) {
+    (void) workspace;
+    return args;
+  }
+
+  template <class ProblemShape>
+  MUTLASS_HOST_DEVICE static bool
+  can_implement(
+    ProblemShape problem_shapes,
+    Arguments const& args) {
+    const int alignmentA = mutlass::detail::get_alignment_count_from_gmem_tiled_copy<GmemTiledCopyA, ElementA>();
+    const int alignmentB = mutlass::detail::get_alignment_count_from_gmem_tiled_copy<GmemTiledCopyB, ElementB>();
+    int problem_m = int(size<0>(problem_shapes));
+    int problem_n = int(size<1>(problem_shapes));
+    int problem_k = int(size<2>(problem_shapes));
+
+    bool implementable = true;
+
+    implementable = implementable && mutlass::detail::check_alignment<alignmentA>(mute::make_shape(problem_m, problem_k), StrideA{});
+    implementable = implementable && mutlass::detail::check_alignment<alignmentB>(mute::make_shape(problem_n, problem_k), StrideB{});
+
+    if (!implementable) {
+      MUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements.\n");
+    }
+    return implementable;
+  }
+
+  /// Perform a threadblock-scoped matrix multiply-accumulate
+  template <
+    class FrgTensorD,
+    class TensorA,
+    class TensorB,
+    class FrgTensorC,
+    class KTileIterator,
+    class ResidueMNK
+  >
+  MUTLASS_DEVICE void
+  operator() (
+      FrgTensorD &accum,
+      TensorA gA,
+      TensorB gB,
+      FrgTensorC const &src_accum,
+      KTileIterator k_tile_iter, int k_tile_count,
+      ResidueMNK residue_mnk,
+      int thread_idx,
+      char *smem_buf)
+  {
+    using namespace mute;
+
+    static_assert(is_rmem<FrgTensorD>::value, "D tensor must be rmem resident.");
+    static_assert(is_gmem<TensorA>::value, "A tensor must be gmem resident.");
+    static_assert(is_gmem<TensorB>::value, "B tensor must be gmem resident.");
+    static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
+    static_assert(rank(SmemLayoutA{}) == 2,
+      "MainloopTwoStage must not have a smem shape with a pipeline mode.");
+    static_assert(rank(SmemLayoutB{}) == 2,
+      "MainloopTwoStage must not have a smem shape with a pipeline mode.");
+
+    // Construct shared memory tiles
+    SharedStorage& storage = *reinterpret_cast<SharedStorage*>(smem_buf);
+    Tensor sA = make_tensor(make_smem_ptr(storage.smem_a.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(storage.smem_b.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
+
+    // Shift tensor so residue_k is at origin (Can't read any k_coord < residue_k)
+    // This aligns the tensor with BLK_K for all but the 0th k_tile
+    gA.data() = &gA(0, get<2>(residue_mnk), 0);
+    gB.data() = &gB(0, get<2>(residue_mnk), 0);
+
+    // Partition the copying of A and B tiles across the threads
+    GmemTiledCopyA gmem_tiled_copy_a;
+    GmemTiledCopyB gmem_tiled_copy_b;
+    auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(thread_idx);
+    auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(thread_idx);
+
+    Tensor tAgA = gmem_thr_copy_a.partition_S(gA);                             // (ACPY,ACPY_M,ACPY_K,k)
+    Tensor tAsA = gmem_thr_copy_a.partition_D(sA);                             // (ACPY,ACPY_M,ACPY_K,PIPE)
+    Tensor tBgB = gmem_thr_copy_b.partition_S(gB);                             // (BCPY,BCPY_N,BCPY_K,k)
+    Tensor tBsB = gmem_thr_copy_b.partition_D(sB);                             // (BCPY,BCPY_N,BCPY_K,PIPE)
+
+    // Allocate the register tiles for double buffering -- same shape as partitioned data
+    Tensor tArA = make_fragment_like(tAsA);                                    // (ACPY,ACPY_M,ACPY_K)
+    Tensor tBrB = make_fragment_like(tBsB);                                    // (BCPY,BCPY_N,BCPY_K)
+
+    //
+    // PREDICATES
+    //
+
+    // Allocate predicate tensors for m and n
+    Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1,_0>{});
+    Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1,_0>{});
+
+    // Construct identity layout for sA and sB
+    Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA)));    // (BLK_M,BLK_K) -> (blk_m,blk_k)
+    Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB)));    // (BLK_N,BLK_K) -> (blk_n,blk_k)
+
+    // Repeat the partitioning with identity layouts
+    Tensor tAcA = gmem_thr_copy_a.partition_S(cA);                             // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k)
+    Tensor tBcB = gmem_thr_copy_b.partition_S(cB);                             // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k)
+
+    // Set predicates for m bounds
+    MUTLASS_PRAGMA_UNROLL
+    for (int m = 0; m < size<0>(tApA); ++m) {
+      tApA(m,0) = get<0>(tAcA(0,m,0)) < get<0>(residue_mnk);  // blk_m coord < residue_m
+    }
+    // Set predicates for n bounds
+    MUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < size<0>(tBpB); ++n) {
+      tBpB(n,0) = get<0>(tBcB(0,n,0)) < get<1>(residue_mnk);  // blk_n coord < residue_n
+    }
+
+    //
+    // PREFETCH
+    //
+
+    // Clear the rmem tiles to account for predicated off loads
+    clear(tArA);
+    clear(tBrB);
+
+    // Start async loads for 0th k-tile, where we take care of the k residue
+    {
+      Tensor tAgAk = tAgA(_,_,_,*k_tile_iter);
+      MUTLASS_PRAGMA_UNROLL
+      for (int k = 0; k < size<2>(tArA); ++k) {
+        if (get<1>(tAcA(0,0,k)) >= -get<2>(residue_mnk)) {      // blk_k coord < residue_k (gA shifted)
+          copy_if(gmem_tiled_copy_a, tApA(_,k), tAgAk(_,_,k), tArA(_,_,k));
+        }
+      }
+      Tensor tBgBk = tBgB(_,_,_,*k_tile_iter);
+      MUTLASS_PRAGMA_UNROLL
+      for (int k = 0; k < size<2>(tBrB); ++k) {
+        if (get<1>(tBcB(0,0,k)) >= -get<2>(residue_mnk)) {      // blk_k coord < residue_k (gB shifted)
+          copy_if(gmem_tiled_copy_b, tBpB(_,k), tBgBk(_,_,k), tBrB(_,_,k));
+        }
+      }
+      ++k_tile_iter;
+      --k_tile_count;
+    }
+
+    // Tile MMA compute thread partitions and allocate accumulators
+    TiledMma tiled_mma;
+    auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
+    Tensor tCrA  = thr_mma.make_fragment_A(thr_mma.partition_A(sA));           // (MMA,MMA_M,MMA_K)
+    Tensor tCrB  = thr_mma.make_fragment_B(thr_mma.partition_B(sB));           // (MMA,MMA_M,MMA_K)
+
+    MUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum));                     // MMA_M
+    MUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(src_accum));                 // MMA_M
+    MUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum));                     // MMA_N
+    MUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(src_accum));                 // MMA_N
+    MUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));                      // MMA_K
+
+    //
+    // Copy Atom retiling
+    //
+
+    auto thr_copy_A       = make_tiled_copy_A(SmemCopyAtomA{}, tiled_mma).get_thread_slice(thread_idx);
+    Tensor tCsA           = thr_copy_A.partition_S(sA);
+    Tensor tCrA_copy_view = thr_copy_A.retile_D(tCrA);
+    MUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));            // M
+
+    auto thr_copy_B       = make_tiled_copy_B(SmemCopyAtomB{}, tiled_mma).get_thread_slice(thread_idx);
+    Tensor tCsB           = thr_copy_B.partition_S(sB);
+    Tensor tCrB_copy_view = thr_copy_B.retile_D(tCrB);
+    MUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));            // N
+
+    //
+    // Prologue
+    //
+
+    // Copy rmem to smem
+    copy(tArA, tAsA);
+    copy(tBrB, tBsB);
+    // Clear accumulators
+    __syncthreads();
+
+    // Load A, B smem->rmem for k=0
+    copy(tCsA(_,_,0), tCrA_copy_view(_,_,0));
+    copy(tCsB(_,_,0), tCrB_copy_view(_,_,0));
+    //
+    // Mainloop
+    //
+
+    // Size of the k-tiles's outer product mode (k)
+    auto K_BLOCK_MAX = size<2>(tCrA);
+
+    MUTLASS_PRAGMA_NO_UNROLL
+    while (k_tile_count > -1)
+    {
+      // Pipeline the outer products with a static for loop
+      for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block)
+      {
+        if (k_block == K_BLOCK_MAX - 1)
+        {
+          __syncthreads();
+
+          // Copy rmem to smem
+          copy(tArA, tAsA);
+          copy(tBrB, tBsB);
+          __syncthreads();
+        }
+
+        // Load A, B smem->rmem for k+1
+        int k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX;    // static
+        copy(tCsA(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next));
+        copy(tCsB(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next));
+        if (k_block == 0)
+        {
+          if (k_tile_count <= 0) {
+            clear(tApA);
+            clear(tBpB);
+          }
+          copy_if(gmem_tiled_copy_a, tApA, tAgA(_,_,_,*k_tile_iter), tArA);
+          copy_if(gmem_tiled_copy_b, tBpB, tBgB(_,_,_,*k_tile_iter), tBrB);
+          ++k_tile_iter;
+          --k_tile_count;
+        }
+
+        // transform before compute
+        mute::transform(tCrA(_,_,k_block), TransformA{});
+        mute::transform(tCrB(_,_,k_block), TransformB{});
+
+        // Thread-level register gemm for k
+        // disambiguate gemm (shared with the namespace name)
+        mute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), src_accum);
+      });
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp31_mma_tme_sqmma_ss.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp31_mma_tme_sqmma_ss.hpp
new file mode 100644
index 000000000..9cb66d008
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp31_mma_tme_sqmma_ss.hpp
@@ -0,0 +1,461 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+#include "mutlass/pipeline/pipeline.hpp"
+
+#include "mute/atom/mma_atom.hpp"
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::collective {
+
+using namespace mute;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  int Stages,
+  class TileShape_,
+  class ElementA_,
+  class StrideA_,
+  class ElementB_,
+  class StrideB_,
+  class TiledMma_,
+  class GmemTiledCopyA_,
+  class SmemLayoutAtomA_,
+  class SmemCopyAtomA_,
+  class TransformA_,
+  class GmemTiledCopyB_,
+  class SmemLayoutAtomB_,
+  class SmemCopyAtomB_,
+  class TransformB_>
+struct CollectiveMma<
+    MainloopMp31TmeSqmma<Stages>,
+    TileShape_,
+    ElementA_,
+    StrideA_,
+    ElementB_,
+    StrideB_,
+    TiledMma_,
+    GmemTiledCopyA_,
+    SmemLayoutAtomA_,
+    SmemCopyAtomA_,
+    TransformA_,
+    GmemTiledCopyB_,
+    SmemLayoutAtomB_,
+    SmemCopyAtomB_,
+    TransformB_>
+{
+  //
+  // Type Aliases
+  //
+  using DispatchPolicy = MainloopMp31TmeSqmma<Stages>;
+  using TileShape = TileShape_;
+  using ElementA = ElementA_;
+  using StrideA = StrideA_;
+  using ElementB = ElementB_;
+  using StrideB = StrideB_;
+  using TiledMma = TiledMma_;
+  using ElementAccumulator = typename TiledMma::ValTypeC;
+  using GmemTiledCopyA = GmemTiledCopyA_;
+  using GmemTiledCopyB = GmemTiledCopyB_;
+  using SmemLayoutAtomA = SmemLayoutAtomA_;
+  using SmemLayoutAtomB = SmemLayoutAtomB_;
+  using SmemCopyAtomA = SmemCopyAtomA_;
+  using SmemCopyAtomB = SmemCopyAtomB_;
+  using TransformA = TransformA_;
+  using TransformB = TransformB_;
+  using ArchTag = typename DispatchPolicy::ArchTag;
+
+  using MainloopPipeline = mutlass::Mp31PipelineTmeAsync<DispatchPolicy::Stages>;
+
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState  = typename mutlass::PipelineState<DispatchPolicy::Stages>;
+
+  static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  static_assert(DispatchPolicy::Stages >= 2, "Stages should set to value 2 or more.");
+
+  static_assert(mute::is_base_of<mute::MP31::SQMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
+                mute::is_base_of<mute::MP31::SQMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
+                "MMA atom must source both A and B operand from smem_desc for this mainloop.");
+
+  static_assert(mute::is_same_v<GmemTiledCopyA, MP31_TME_LOAD>, "GmemTiledCopy - invalid MP31 TME copy atom specified.");
+  static_assert(mute::is_same_v<GmemTiledCopyB, MP31_TME_LOAD>, "GmemTiledCopy - invalid MP31 TME copy atom specified.");
+
+  using SmemLayoutA = decltype(tile_to_shape(
+      SmemLayoutAtomA{},
+      make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
+  using SmemLayoutB = decltype(tile_to_shape(
+      SmemLayoutAtomB{},
+      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
+
+  static constexpr int NumWarps = size(TiledMma{}) / NumThreadsPerWarp;
+
+  static constexpr int NumBarriers = MainloopPipeline::NumBarriers;
+
+  static constexpr int SmemAlignmentBytes = 256;
+
+  struct SharedStorage
+  {
+    mute::array_aligned<typename TiledMma::ValTypeA, mute::cosize_v<SmemLayoutA>, SmemAlignmentBytes> smem_A;
+    mute::array_aligned<typename TiledMma::ValTypeB, mute::cosize_v<SmemLayoutB>, SmemAlignmentBytes> smem_B;
+  };
+
+  // Host side kernel arguments
+  struct Arguments {
+    ElementA const* ptr_A;
+    StrideA dA;
+    ElementB const* ptr_B;
+    StrideB dB;
+  };
+
+  // Device side kernel params
+  struct Params {
+    using TME_A = decltype(make_tme_copy(
+        GmemTiledCopyA{},
+        make_tensor(static_cast<ElementA const*>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
+        take<0,2>(SmemLayoutA{}),
+        make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}))));
+
+    using TME_B = decltype(make_tme_copy(
+        GmemTiledCopyB{},
+        make_tensor(static_cast<ElementB const*>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
+        take<0,2>(SmemLayoutB{}),
+        make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}))));
+    TME_A tme_load_a;
+    TME_B tme_load_b;
+  };
+
+
+  //
+  // Methods
+  //
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+
+    // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    Tensor tensor_a = make_tensor(args.ptr_A, make_layout(make_shape(M,K,L), args.dA));
+    Tensor tensor_b = make_tensor(args.ptr_B, make_layout(make_shape(N,K,L), args.dB));
+
+    typename Params::TME_A tme_load_a = make_tme_copy(
+        GmemTiledCopyA{},
+        tensor_a,
+        take<0,2>(SmemLayoutA{}),
+        make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})));
+
+    typename Params::TME_B tme_load_b = make_tme_copy(
+        GmemTiledCopyB{},
+        tensor_b,
+        take<0,2>(SmemLayoutB{}),
+        make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})));
+    return {
+      tme_load_a,
+      tme_load_b
+    };
+  }
+
+  template <class ProblemShape>
+  MUTLASS_HOST_DEVICE static bool
+  can_implement(
+      ProblemShape const& problem_shape,
+      [[maybe_unused]] Arguments const& args) {
+    constexpr int tme_alignment_bits = 32;
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    bool implementable = true;
+    constexpr int min_tme_aligned_elements_A = tme_alignment_bits / mutlass::sizeof_bits<ElementA>::value;
+    implementable = implementable && mutlass::detail::check_alignment<min_tme_aligned_elements_A>(mute::make_shape(M,K,L), StrideA{});
+    constexpr int min_tme_aligned_elements_B = tme_alignment_bits / mutlass::sizeof_bits<ElementB>::value;
+    implementable = implementable && mutlass::detail::check_alignment<min_tme_aligned_elements_B>(mute::make_shape(N,K,L), StrideB{});
+
+    if (!implementable) {
+      MUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TME.\n");
+    }
+    return implementable;
+  }
+
+  /// Issue Tme Descriptor Prefetch
+  MUTLASS_DEVICE
+  static void prefetch_tme_descriptors(Params const& mainloop_params)
+  {
+    mute::prefetch_tme_descriptor(mainloop_params.tme_load_a.get_tme_descriptor());
+    mute::prefetch_tme_descriptor(mainloop_params.tme_load_b.get_tme_descriptor());
+  }
+
+  template <class ProblemShape_MNKL>
+  MUTLASS_DEVICE auto
+  load_init(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params) const {
+    using X = Underscore;
+    // Separate out problem shape for convenience
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    // TME requires special handling of strides to deal with coord codomain mapping
+    // Represent the full tensors -- get these from TME
+    Tensor mA_mkl = mainloop_params.tme_load_a.get_tme_tensor(make_shape(M,K,L));                            // (m,k,l)
+    Tensor mB_nkl = mainloop_params.tme_load_b.get_tme_tensor(make_shape(N,K,L));                            // (n,k,l)
+
+    auto blk_shape = TileShape{};                                                                // (BLK_M,BLK_N,BLK_K)
+    auto blk_coord = make_coord(_,_,_);                                                   // (m,n,k) -- defer the slice
+
+    // Makke tiled views
+    Tensor gA_mkl = local_tile(mA_mkl, blk_shape, blk_coord, Step<_1, X,_1>{});                  // (BLK_M,BLK_K,m,k,l)
+    Tensor gB_nkl = local_tile(mB_nkl, blk_shape, blk_coord, Step< X,_1,_1>{});                  // (BLK_N,BLK_K,n,k,l)
+
+    return mute::make_tuple(gA_mkl, gB_nkl);
+  }
+
+  /// Perform a collective-scoped matrix multiply-accumulate
+  template <
+    class TensorA, class TensorB,
+    class BlockCoord,
+    class FrgTensorC,
+    class KTileIterator
+  >
+  MUTLASS_DEVICE void
+  operator() (
+      mute::tuple<TensorA, TensorB> const& load_inputs,
+      BlockCoord const& blk_coord,
+      FrgTensorC& accum,
+      KTileIterator k_tile_iter, int k_tile_count,
+      int thread_idx,
+      char* shared_memory,
+      Params const& mainloop_params)
+  {
+    using namespace mute;
+
+    static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
+    static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2.");
+    static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2.");
+    static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
+    static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
+    static_assert(mute::is_void_v<SmemCopyAtomA>,
+      "MP31 SQMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
+    static_assert(mute::is_void_v<SmemCopyAtomB>,
+      "MP31 SQMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
+
+    SharedStorage& storage = *reinterpret_cast<SharedStorage*>(shared_memory);
+    Tensor sA = make_tensor(make_smem_ptr(storage.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(storage.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
+
+    auto cta_tme_a = mainloop_params.tme_load_a.get_slice(0);
+    auto cta_tme_b = mainloop_params.tme_load_b.get_slice(0);
+
+    Tensor gA_mkl = get<0>(load_inputs);
+    Tensor gB_nkl = get<1>(load_inputs);
+
+    auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
+
+    Tensor gA = gA_mkl(_,_,m_coord,_,l_coord);
+    Tensor gB = gB_nkl(_,_,n_coord,_,l_coord);
+
+
+    // Applies the mapping frome tme
+    Tensor tAgA = cta_tme_a.partition_S(gA);
+    Tensor tAsA = cta_tme_a.partition_D(sA);
+
+    Tensor tBgB = cta_tme_b.partition_S(gB);
+    Tensor tBsB = cta_tme_b.partition_D(sB);
+
+    // Set the bytes transferred in this TME transaction (may involve multiple issues)
+    constexpr uint32_t TmeTransactionBytes = static_cast<uint32_t>(
+        (size<0>(sA) * size<1>(sA) * sizeof(ElementA)) +
+        (size<0>(sB) * size<1>(sB) * sizeof(ElementB)));
+
+    // Obtain warp index
+    int warp_idx = canonical_warp_idx();
+
+    PipelineParams params;
+    params.transaction_bytes = TmeTransactionBytes;
+    params.num_consumers = NumWarps;
+
+    // Init pipeline
+    MainloopPipeline pipeline(params);
+
+    // State variables used for iterating the circular buffer
+    // smem_pipe_read / release is used by the consumer of SMEM data - i.e MMA
+    // smem_pipe_write is used by the producer of SMEM data - i.e. TME
+    PipelineState smem_pipe_read;
+    PipelineState smem_pipe_release;
+    PipelineState smem_pipe_write = mutlass::make_producer_start_state<MainloopPipeline>();
+
+    // We need this to gurantee that the Pipeline init is visible
+    // to all warps in the CTA
+    __syncthreads();
+
+
+    //
+    // Prologue TME
+    //
+
+    // Keep a copy to know when to stop issuing loads
+    int k_tile_count_tme = k_tile_count;
+
+    // Start async tme loads for all pipes but the last
+    MUTLASS_PRAGMA_UNROLL
+    for (int k_pipe = 0; k_pipe < DispatchPolicy::Stages-1; ++k_pipe) {
+      if (warp_idx == 0 && k_tile_count_tme > 0) {
+        pipeline.producer_acquire(smem_pipe_write);
+        uint32_t bar_id = pipeline.producer_get_barrier_id(smem_pipe_write);
+        copy(mainloop_params.tme_load_a.with(bar_id), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,k_pipe));
+        copy(mainloop_params.tme_load_b.with(bar_id), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,k_pipe));
+
+      }
+      // advance producer pipeline
+      ++smem_pipe_write;
+      ++k_tile_iter;
+      --k_tile_count_tme;
+    }
+
+    //
+    // Define C accumulators and A/B partitioning
+    //
+
+    TiledMma tiled_mma;
+    auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
+
+    Tensor tCsA = thread_mma.partition_A(sA);                                  // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCsB = thread_mma.partition_B(sB);                                  // (MMA,MMA_N,MMA_K,PIPE)
+    // Allocate "fragments/descriptors"
+    Tensor tCrA = thread_mma.make_fragment_A(tCsA);                            // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCrB = thread_mma.make_fragment_B(tCsB);                            // (MMA,MMA_N,MMA_K,PIPE)
+
+    MUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum));                     // M
+    MUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum));                     // N
+    MUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));                      // K
+    MUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));                      // PIPE
+    MUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tAsA));                      // PIPE
+    MUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tBsB));                      // PIPE
+    MUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA));        // PIPE
+    MUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB));        // PIPE
+    // Currently we prefer one instruction per squad atom
+    MUTE_STATIC_ASSERT_V(size<1>(accum) == Int<1>{});
+    MUTE_STATIC_ASSERT_V(size<2>(accum) == Int<1>{});
+    MUTE_STATIC_ASSERT_V(size<2>(tCsA)  == Int<1>{});
+
+    // Prologue MMA
+    {
+      pipeline.consumer_wait(smem_pipe_read);
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+        mute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accum);
+      }
+
+      // Last stage tme
+      if (warp_idx == 0 && k_tile_count_tme > 0) {
+        pipeline.producer_acquire(smem_pipe_write);
+        uint32_t bar_id = pipeline.producer_get_barrier_id(smem_pipe_write);
+        copy(mainloop_params.tme_load_a.with(bar_id), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,DispatchPolicy::Stages-1));
+        copy(mainloop_params.tme_load_b.with(bar_id), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,DispatchPolicy::Stages-1));
+      }
+      // advance producer pipeline
+      ++smem_pipe_write;
+      ++k_tile_iter;
+      --k_tile_count_tme;
+
+      warpsquad_wait();
+      pipeline.consumer_release(smem_pipe_release);
+
+      // advance consumer pipeline
+      ++smem_pipe_read;
+      ++smem_pipe_release;
+      --k_tile_count;
+    }
+
+    if (k_tile_count == 0) return;
+
+    // PIPELINED MAIN LOOP
+
+    MUTLASS_PRAGMA_NO_UNROLL
+    for ( ; k_tile_count > 1; --k_tile_count)
+    {
+      pipeline.consumer_wait(smem_pipe_read);
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+        mute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accum);
+      }
+
+      if (warp_idx == 0 && k_tile_count_tme > 0) {
+        pipeline.producer_acquire(smem_pipe_write);
+        uint32_t bar_id = pipeline.producer_get_barrier_id(smem_pipe_write);
+        copy(mainloop_params.tme_load_a.with(bar_id), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write.index()));
+        copy(mainloop_params.tme_load_b.with(bar_id), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write.index()));
+      }
+      // advance producer pipeline
+      ++smem_pipe_write;
+      ++k_tile_iter;
+      --k_tile_count_tme;
+
+      warpsquad_wait();
+      pipeline.consumer_release(smem_pipe_release);
+
+      // advance consumer pipeline
+      ++smem_pipe_read;
+      ++smem_pipe_release;
+    }
+
+    {
+      pipeline.consumer_wait(smem_pipe_read);
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+        mute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accum);
+      }
+
+      warpsquad_wait();
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp31_mma_tme_sqmma_ss_scaling.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp31_mma_tme_sqmma_ss_scaling.hpp
new file mode 100644
index 000000000..ef55c8fad
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/mp31_mma_tme_sqmma_ss_scaling.hpp
@@ -0,0 +1,687 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+
+#include "mutlass/mutlass.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+#include "mutlass/gemm/collective/scaling_accumulation.hpp"
+#include "mutlass/pipeline/pipeline.hpp"
+
+#include "mute/atom/mma_atom.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::collective {
+
+using namespace mute;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  int Stages,
+  class KernelSchedule,
+  int ScaleGranularityM_,
+  int ScaleGranularityN_,
+  int ScaleGranularityK_,
+  class TileShape_,
+  class ElementA_,
+  class StrideA_,
+  class ElementB_,
+  class StrideB_,
+  class TiledMma_,
+  class GmemTiledCopyA_,
+  class SmemLayoutAtomA_,
+  class SmemCopyAtomA_,
+  class TransformA_,
+  class GmemTiledCopyB_,
+  class SmemLayoutAtomB_,
+  class SmemCopyAtomB_,
+  class TransformB_>
+struct CollectiveMma<
+    MainloopMp31TmeSqmmaBlockScalingFP8<Stages, KernelSchedule, ScaleGranularityM_, ScaleGranularityN_, ScaleGranularityK_>,
+    TileShape_,
+    ElementA_,
+    StrideA_,
+    ElementB_,
+    StrideB_,
+    TiledMma_,
+    GmemTiledCopyA_,
+    SmemLayoutAtomA_,
+    SmemCopyAtomA_,
+    TransformA_,
+    GmemTiledCopyB_,
+    SmemLayoutAtomB_,
+    SmemCopyAtomB_,
+    TransformB_>
+{
+  //
+  // Type Aliases
+  //
+  using DispatchPolicy = MainloopMp31TmeSqmmaBlockScalingFP8<Stages, KernelSchedule, ScaleGranularityM_, ScaleGranularityN_, ScaleGranularityK_>;
+  using TileShape = TileShape_;
+  using ElementA = ElementA_;
+  using StrideA = StrideA_;
+  using ElementB = ElementB_;
+  using StrideB = StrideB_;
+  using TiledMma = TiledMma_;
+  using ElementAccumulator = typename TiledMma::ValTypeC;
+  using ElementBlockScale = ElementAccumulator;
+  using GmemTiledCopyA = GmemTiledCopyA_;
+  using GmemTiledCopyB = GmemTiledCopyB_;
+  using SmemLayoutAtomA = SmemLayoutAtomA_;
+  using SmemLayoutAtomB = SmemLayoutAtomB_;
+  using SmemCopyAtomA = SmemCopyAtomA_;
+  using SmemCopyAtomB = SmemCopyAtomB_;
+  using TransformA = TransformA_;
+  using TransformB = TransformB_;
+  using ArchTag = typename DispatchPolicy::ArchTag;
+
+  using MainloopPipeline = mutlass::Mp31PipelineTmeAsync<DispatchPolicy::Stages>;
+
+  using PipelineParams = typename MainloopPipeline::Params;
+  using PipelineState  = typename mutlass::PipelineState<DispatchPolicy::Stages>;
+
+  static constexpr int ScaleGranularityM = ScaleGranularityM_ == 0 ? size<0>(TileShape{}) : ScaleGranularityM_;
+  static constexpr int ScaleGranularityN = ScaleGranularityN_ == 0 ? size<1>(TileShape{}) : ScaleGranularityN_;
+  static constexpr int ScaleGranularityK = ScaleGranularityK_ == 0 ? size<2>(TileShape{}) : ScaleGranularityK_;
+
+  static constexpr int ScaleMsPerTile = size<0>(TileShape{}) / ScaleGranularityM;
+  static constexpr int ScaleNsPerTile = size<1>(TileShape{}) / ScaleGranularityN;
+  static constexpr int ScalePromotionInterleave = ScaleGranularityK / size<2>(TileShape{});
+
+  static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<0>(TileShape{}) % size<0>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomA{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2 (M/N, K)");
+  static_assert((size<1>(TileShape{}) % size<0>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+  static_assert((size<2>(TileShape{}) % size<1>(SmemLayoutAtomB{})) == 0, "SmemLayoutAtom must evenly divide tile shape.");
+
+  static_assert((size<0>(TileShape{}) % ScaleGranularityM) == 0, "Scaling granularity M must evenly divide tile shape along M.");
+  static_assert((size<1>(TileShape{}) % ScaleGranularityN) == 0, "Scaling granularity N must evenly divide tile shape along N.");
+  static_assert((ScaleGranularityK % size<2>(TileShape{})) == 0, "Tile shape along K must evenly divide Scaling granularity K.");
+
+  static_assert(DispatchPolicy::Stages >= 2, "Stages should set to value 2 or more.");
+
+  static_assert(mute::is_base_of<mute::MP31::SQMMA::DescriptorIterator, typename TiledMma::FrgTypeA>::value &&
+                mute::is_base_of<mute::MP31::SQMMA::DescriptorIterator, typename TiledMma::FrgTypeB>::value,
+                "MMA atom must source both A and B operand from smem_desc for this mainloop.");
+
+  static_assert(mute::is_same_v<GmemTiledCopyA, MP31_TME_LOAD>, "GmemTiledCopy - invalid MP31 TME copy atom specified.");
+  static_assert(mute::is_same_v<GmemTiledCopyB, MP31_TME_LOAD>, "GmemTiledCopy - invalid MP31 TME copy atom specified.");
+
+  using SmemLayoutA = decltype(tile_to_shape(
+      SmemLayoutAtomA{},
+      make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
+  using SmemLayoutB = decltype(tile_to_shape(
+      SmemLayoutAtomB{},
+      make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}), Int<DispatchPolicy::Stages>{})));
+
+  static constexpr int NumWarps = size(TiledMma{}) / NumThreadsPerWarp;
+
+  static constexpr int NumBarriers = MainloopPipeline::NumBarriers;
+
+  static constexpr int SmemAlignmentBytes = 256;
+
+  struct SharedStorage
+  {
+    mute::array_aligned<typename TiledMma::ValTypeA, mute::cosize_v<SmemLayoutA>, SmemAlignmentBytes> smem_A;
+    mute::array_aligned<typename TiledMma::ValTypeB, mute::cosize_v<SmemLayoutB>, SmemAlignmentBytes> smem_B;
+  };
+
+  // Host side kernel arguments
+  struct Arguments {
+    ElementA const* ptr_A;
+    StrideA dA;
+    ElementB const* ptr_B;
+    StrideB dB;
+    ElementBlockScale const* ptr_scale_A;
+    ElementBlockScale const* ptr_scale_B;
+  };
+
+  // Device side kernel params
+  struct Params {
+    using TME_A = decltype(make_tme_copy(
+        GmemTiledCopyA{},
+        make_tensor(static_cast<ElementA const*>(nullptr), repeat_like(StrideA{}, int32_t(0)), StrideA{}),
+        take<0,2>(SmemLayoutA{}),
+        make_shape(shape<0>(TileShape{}), shape<2>(TileShape{}))));
+
+    using TME_B = decltype(make_tme_copy(
+        GmemTiledCopyB{},
+        make_tensor(static_cast<ElementB const*>(nullptr), repeat_like(StrideB{}, int32_t(0)), StrideB{}),
+        take<0,2>(SmemLayoutB{}),
+        make_shape(shape<1>(TileShape{}), shape<2>(TileShape{}))));
+    TME_A tme_load_a;
+    TME_B tme_load_b;
+
+    ElementBlockScale const* ptr_scale_A;
+    ElementBlockScale const* ptr_scale_B;
+  };
+
+  //
+  // Methods
+  //
+
+  template <class ProblemShape>
+  static constexpr Params
+  to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
+
+    // Optionally append 1s until problem shape is rank-4 (MNKL), in case it is only rank-3 (MNK)
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    Tensor tensor_a = make_tensor(args.ptr_A, make_layout(make_shape(M,K,L), args.dA));
+    Tensor tensor_b = make_tensor(args.ptr_B, make_layout(make_shape(N,K,L), args.dB));
+
+    typename Params::TME_A tme_load_a = make_tme_copy(
+        GmemTiledCopyA{},
+        tensor_a,
+        take<0,2>(SmemLayoutA{}),
+        make_shape(shape<0>(TileShape{}), shape<2>(TileShape{})));
+
+    typename Params::TME_B tme_load_b = make_tme_copy(
+        GmemTiledCopyB{},
+        tensor_b,
+        take<0,2>(SmemLayoutB{}),
+        make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})));
+    return {
+      tme_load_a,
+      tme_load_b,
+      args.ptr_scale_A,
+      args.ptr_scale_B
+    };
+  }
+
+  template <class ProblemShape>
+  MUTLASS_HOST_DEVICE static bool
+  can_implement(
+      ProblemShape const& problem_shape,
+      [[maybe_unused]] Arguments const& args) {
+    constexpr int tme_alignment_bits = 32;
+    auto problem_shape_MNKL = append<4>(problem_shape, 1);
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    bool implementable = true;
+    constexpr int min_tme_aligned_elements_A = tme_alignment_bits / mutlass::sizeof_bits<ElementA>::value;
+    implementable = implementable && mutlass::detail::check_alignment<min_tme_aligned_elements_A>(mute::make_shape(M,K,L), StrideA{});
+    constexpr int min_tme_aligned_elements_B = tme_alignment_bits / mutlass::sizeof_bits<ElementB>::value;
+    implementable = implementable && mutlass::detail::check_alignment<min_tme_aligned_elements_B>(mute::make_shape(N,K,L), StrideB{});
+
+    if (!implementable) {
+      MUTLASS_TRACE_HOST("  CAN IMPLEMENT: Problem Size doesn't meet the minimum alignment requirements for TME.\n");
+    }
+    return implementable;
+  }
+
+  template <class ProblemShape_MNKL>
+  MUTLASS_DEVICE auto
+  load_init(ProblemShape_MNKL const& problem_shape_MNKL, Params const& mainloop_params) const {
+    using X = Underscore;
+    // Separate out problem shape for convenience
+    auto [M,N,K,L] = problem_shape_MNKL;
+
+    // TME requires special handling of strides to deal with coord codomain mapping
+    // Represent the full tensors -- get these from TME
+    Tensor mA_mkl = mainloop_params.tme_load_a.get_tme_tensor(make_shape(M,K,L));                            // (m,k,l)
+    Tensor mB_nkl = mainloop_params.tme_load_b.get_tme_tensor(make_shape(N,K,L));                            // (n,k,l)
+
+    auto blk_shape = TileShape{};                                                                // (BLK_M,BLK_N,BLK_K)
+    auto blk_coord = make_coord(_,_,_);                                                   // (m,n,k) -- defer the slice
+
+    // Makke tiled views
+    Tensor gA_mkl = local_tile(mA_mkl, blk_shape, blk_coord, Step<_1, X,_1>{});                  // (BLK_M,BLK_K,m,k,l)
+    Tensor gB_nkl = local_tile(mB_nkl, blk_shape, blk_coord, Step< X,_1,_1>{});                  // (BLK_N,BLK_K,n,k,l)
+
+    // Scaling Tensor
+    auto scaleA_shape = make_shape(M / ScaleGranularityM, K / ScaleGranularityK, L);
+    auto scaleA_layout = make_ordered_layout(scaleA_shape, Step<_0, _1, _2>{});
+
+    auto scaleB_shape = make_shape(N / ScaleGranularityN, K / ScaleGranularityK, L);
+    auto scaleB_layout = make_ordered_layout(scaleB_shape, Step<_0, _1, _2>{});
+
+    Tensor mScaleA_mkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_A), scaleA_layout);
+    Tensor mScaleB_nkl = make_tensor(make_gmem_ptr(mainloop_params.ptr_scale_B), scaleB_layout);
+
+    return mute::make_tuple(gA_mkl, gB_nkl, mScaleA_mkl, mScaleB_nkl);
+  }
+
+  /// Perform a collective-scoped matrix multiply-accumulate
+  template <
+    class TensorA, class TensorB,
+    class TensorScaleA, class TensorScaleB,
+    class BlockCoord,
+    class FrgTensorC,
+    class KTileIterator
+  >
+  MUTLASS_DEVICE void
+  operator() (
+      mute::tuple<TensorA, TensorB, TensorScaleA, TensorScaleB> const& load_inputs,
+      BlockCoord const& blk_coord,
+      FrgTensorC& accum,
+      KTileIterator k_tile_iter, int k_tile_count,
+      int thread_idx,
+      char* shared_memory,
+      Params const& mainloop_params)
+  {
+    using namespace mute;
+
+    static_assert(is_rmem<FrgTensorC>::value, "C tensor must be rmem resident.");
+    static_assert(rank(SmemLayoutAtomA{}) == 2, "SmemLayoutAtom must be rank 2.");
+    static_assert(rank(SmemLayoutAtomB{}) == 2, "SmemLayoutAtom must be rank 2.");
+    static_assert(rank(SmemLayoutA{}) == 3, "Smem layout must be rank 3.");
+    static_assert(rank(SmemLayoutB{}) == 3, "Smem layout must be rank 3.");
+    static_assert(mute::is_void_v<SmemCopyAtomA>,
+      "MP31 SQMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
+    static_assert(mute::is_void_v<SmemCopyAtomB>,
+      "MP31 SQMMA mainloops cannot have a non-void copy atom for smem sourced instructions.");
+
+    SharedStorage& storage = *reinterpret_cast<SharedStorage*>(shared_memory);
+    Tensor sA = make_tensor(make_smem_ptr(storage.smem_A.data()), SmemLayoutA{}); // (BLK_M,BLK_K,PIPE)
+    Tensor sB = make_tensor(make_smem_ptr(storage.smem_B.data()), SmemLayoutB{}); // (BLK_N,BLK_K,PIPE)
+
+    auto cta_tme_a = mainloop_params.tme_load_a.get_slice(0);
+    auto cta_tme_b = mainloop_params.tme_load_b.get_slice(0);
+
+    Tensor gA_mkl = get<0>(load_inputs);
+    Tensor gB_nkl = get<1>(load_inputs);
+
+    auto [m_coord, n_coord, k_coord, l_coord] = blk_coord;
+
+    Tensor gA = gA_mkl(_,_,m_coord,_,l_coord);
+    Tensor gB = gB_nkl(_,_,n_coord,_,l_coord);
+
+    // Applies the mapping frome tme
+    Tensor tAgA = cta_tme_a.partition_S(gA);
+    Tensor tAsA = cta_tme_a.partition_D(sA);
+
+    Tensor tBgB = cta_tme_b.partition_S(gB);
+    Tensor tBsB = cta_tme_b.partition_D(sB);
+
+    // Scale tensor
+    Tensor mScaleA_mkl = get<2>(load_inputs); // (M/ScaleGranularityM,K/ScaleGranularityK,L)
+    Tensor mScaleB_nkl = get<3>(load_inputs); // (N/ScaleGranularityN,K/ScaleGranularityK,L)
+
+    Tensor gScaleA = local_tile(mScaleA_mkl, make_tile(Int<ScaleMsPerTile>{}),
+                make_coord(m_coord, _, l_coord)); // (ScaleMsPerTile, K/ScaleGranularityK, 1)
+
+    Tensor gScaleB = local_tile(mScaleB_nkl, make_tile(Int<ScaleNsPerTile>{}),
+                make_coord(n_coord, _, l_coord)); // (ScaleNsPerTile, K/ScaleGranularityK, 1)
+
+    auto scale_k_iter = mute::make_coord_iterator(shape<1>(gScaleA));
+
+    TiledMma tiled_mma;
+    auto thread_mma = tiled_mma.get_thread_slice(thread_idx);
+
+    auto ScaleAViewAsCLayout = Layout<
+      Shape<Shape<Int<ScaleGranularityM>, Int<ScaleMsPerTile>>, mute::tuple_element_t<1, TileShape>>,
+      Stride<Stride<_0, _1>, _0>>{};
+    auto ScaleBViewAsCLayout = Layout<
+      Shape<mute::tuple_element_t<0, TileShape>, Shape<Int<ScaleGranularityN>, Int<ScaleNsPerTile>>>,
+      Stride<_0, Stride<_0, _1>>>{};
+
+
+    Tensor tCgScaleAViewAsC_ = thread_mma.partition_C(gScaleA(_, 0).compose(ScaleAViewAsCLayout)); // (MMA, MMA_M, MMA_N)
+    Tensor tCgScaleBViewAsC_ = thread_mma.partition_C(gScaleB(_, 0).compose(ScaleBViewAsCLayout)); // (MMA, MMA_M, MMA_N)
+
+    Tensor tCrScaleAViewAsC = make_tensor_like<ElementBlockScale>(tCgScaleAViewAsC_);
+    Tensor tCrScaleBViewAsC = make_tensor_like<ElementBlockScale>(tCgScaleBViewAsC_);
+
+
+    // Set the bytes transferred in this TME transaction (may involve multiple issues)
+    constexpr uint32_t TmeTransactionBytes = static_cast<uint32_t>(
+        (size<0>(sA) * size<1>(sA) * sizeof(ElementA)) +
+        (size<0>(sB) * size<1>(sB) * sizeof(ElementB)));
+
+    // Obtain warp index
+    int warp_idx = canonical_warp_idx();
+
+    PipelineParams params;
+    params.transaction_bytes = TmeTransactionBytes;
+    params.num_consumers = NumWarps;
+
+
+    // Init pipeline
+    MainloopPipeline pipeline(params);
+
+    // State variables used for iterating the circular buffer
+    // smem_pipe_read / release is used by the consumer of SMEM data - i.e MMA
+    // smem_pipe_write is used by the producer of SMEM data - i.e. TME
+    PipelineState smem_pipe_read;
+    PipelineState smem_pipe_release;
+    PipelineState smem_pipe_write = mutlass::make_producer_start_state<MainloopPipeline>();
+
+    // We need this to gurantee that the Pipeline init is visible
+    // to all warps in the CTA
+    __syncthreads();
+
+
+
+    //
+    // Prologue TME
+    //
+
+    // Keep a copy to know when to stop issuing loads
+    int k_tile_count_tme = k_tile_count;
+
+    // Start async tme loads for all pipes but the last
+    MUTLASS_PRAGMA_UNROLL
+    for (int k_pipe = 0; k_pipe < DispatchPolicy::Stages-1; ++k_pipe) {
+      if (warp_idx == 0 && k_tile_count_tme > 0) {
+        pipeline.producer_acquire(smem_pipe_write);
+        uint32_t bar_id = pipeline.producer_get_barrier_id(smem_pipe_write);
+        copy(mainloop_params.tme_load_a.with(bar_id), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,k_pipe));
+        copy(mainloop_params.tme_load_b.with(bar_id), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,k_pipe));
+      }
+
+      // advance producer pipeline
+      ++smem_pipe_write;
+      ++k_tile_iter;
+      --k_tile_count_tme;
+    }
+
+    //
+    // Define C accumulators and A/B partitioning
+    //
+
+    Tensor tCsA = thread_mma.partition_A(sA);                                  // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCsB = thread_mma.partition_B(sB);                                  // (MMA,MMA_N,MMA_K,PIPE)
+    // Allocate "fragments/descriptors"
+    Tensor tCrA = thread_mma.make_fragment_A(tCsA);                            // (MMA,MMA_M,MMA_K,PIPE)
+    Tensor tCrB = thread_mma.make_fragment_B(tCsB);                            // (MMA,MMA_N,MMA_K,PIPE)
+
+    clear(accum);
+    ScalingAccumulation accumulation(accum, ScalePromotionInterleave, size<2>(tCrA));
+
+    MUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(accum));                     // M
+    MUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<2>(accum));                     // N
+    MUTE_STATIC_ASSERT_V(size<2>(tCsA) == size<2>(tCsB));                      // K
+    MUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tCsB));                      // PIPE
+    MUTE_STATIC_ASSERT_V(size<3>(tCsA) == size<3>(tAsA));                      // PIPE
+    MUTE_STATIC_ASSERT_V(size<3>(tCsB) == size<3>(tBsB));                      // PIPE
+    MUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sA));        // PIPE
+    MUTE_STATIC_ASSERT_V(Int<DispatchPolicy::Stages>{} == size<2>(sB));        // PIPE
+    // Currently we prefer one instruction per squad atom
+    MUTE_STATIC_ASSERT_V(size<1>(accum) == Int<1>{});
+    MUTE_STATIC_ASSERT_V(size<2>(accum) == Int<1>{});
+    MUTE_STATIC_ASSERT_V(size<2>(tCsA)  == Int<1>{});
+
+
+    // Prologue MMA
+    {
+      pipeline.consumer_wait(smem_pipe_read);
+
+      // load fisrt scale tensor
+      Tensor tCgScaleAViewAsC = thread_mma.partition_C(gScaleA(_, *scale_k_iter).compose(ScaleAViewAsCLayout)); // (MMA, MMA_M, MMA_N)
+      Tensor tCgScaleBViewAsC = thread_mma.partition_C(gScaleB(_, *scale_k_iter).compose(ScaleBViewAsCLayout)); // (MMA, MMA_M, MMA_N)
+
+      // Load per block scale values from global memory to registers
+      copy(tCgScaleAViewAsC, tCrScaleAViewAsC);
+      copy(tCgScaleBViewAsC, tCrScaleBViewAsC);
+
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
+        tCrScaleAViewAsC(0) = tCrScaleAViewAsC(0) * tCrScaleBViewAsC(0);
+      }
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
+        ElementBlockScale scale_a = tCrScaleAViewAsC(0);
+        MUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < size(tCrScaleBViewAsC); ++i) {
+          tCrScaleBViewAsC(i) = scale_a * tCrScaleBViewAsC(i);
+        }
+      }
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
+        ElementBlockScale scale_b = tCrScaleBViewAsC(0);
+        MUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < size(tCrScaleAViewAsC); ++i) {
+          tCrScaleAViewAsC(i) = scale_b * tCrScaleAViewAsC(i);
+        }
+      }
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+        mute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accumulation());
+      }
+
+      // Last stage tme
+      if (warp_idx == 0 && k_tile_count_tme > 0) {
+        pipeline.producer_acquire(smem_pipe_write);
+        uint32_t bar_id = pipeline.producer_get_barrier_id(smem_pipe_write);
+        copy(mainloop_params.tme_load_a.with(bar_id), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,DispatchPolicy::Stages-1));
+        copy(mainloop_params.tme_load_b.with(bar_id), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,DispatchPolicy::Stages-1));
+      }
+      // advance producer pipeline
+      ++smem_pipe_write;
+      ++k_tile_iter;
+      --k_tile_count_tme;
+
+      warpsquad_wait();
+      pipeline.consumer_release(smem_pipe_release);
+
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
+        ElementBlockScale scale_ab = tCrScaleAViewAsC(0);
+        accumulation.scale_if_needed(scale_ab);
+      }
+
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
+        accumulation.scale_if_needed(tCrScaleBViewAsC);
+      }
+
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
+        accumulation.scale_if_needed(tCrScaleAViewAsC);
+      }
+
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
+        accumulation.scale_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
+      }
+
+      // advance consumer pipeline
+      ++smem_pipe_read;
+      ++smem_pipe_release;
+      --k_tile_count;
+    }
+
+    if (k_tile_count == 0) {
+      // Scale the unscaled accum before return
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
+        ElementBlockScale scale_ab = tCrScaleAViewAsC(0);
+        accumulation.scale_residue_if_needed(scale_ab);
+      }
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
+        accumulation.scale_residue_if_needed(tCrScaleBViewAsC);
+      }
+
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
+        accumulation.scale_residue_if_needed(tCrScaleAViewAsC);
+      }
+
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
+        accumulation.scale_residue_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
+      }
+      return;
+    }
+
+    // PIPELINED MAIN LOOP
+    MUTLASS_PRAGMA_NO_UNROLL
+    for ( ; k_tile_count > 1; --k_tile_count)
+    {
+      pipeline.consumer_wait(smem_pipe_read);
+
+      if (accumulation.prepare_if_needed()) {
+        tiled_mma.accumulate_ = MP31::SQMMA::ScaleOut::Zero;
+        ++scale_k_iter;
+        Tensor tCgScaleAViewAsC = thread_mma.partition_C(gScaleA(_, *scale_k_iter).compose(ScaleAViewAsCLayout)); // (MMA, MMA_M, MMA_N)
+        Tensor tCgScaleBViewAsC = thread_mma.partition_C(gScaleB(_, *scale_k_iter).compose(ScaleBViewAsCLayout)); // (MMA, MMA_M, MMA_N)
+
+        copy(tCgScaleAViewAsC, tCrScaleAViewAsC);
+        copy(tCgScaleBViewAsC, tCrScaleBViewAsC);
+        if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
+          tCrScaleAViewAsC(0) = tCrScaleAViewAsC(0) * tCrScaleBViewAsC(0);
+        }
+        if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
+          ElementBlockScale scale_a = tCrScaleAViewAsC(0);
+          MUTLASS_PRAGMA_UNROLL
+          for (int i = 0; i < size(tCrScaleBViewAsC); ++i) {
+            tCrScaleBViewAsC(i) = scale_a * tCrScaleBViewAsC(i);
+          }
+        }
+        if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
+          ElementBlockScale scale_b = tCrScaleBViewAsC(0);
+          MUTLASS_PRAGMA_UNROLL
+          for (int i = 0; i < size(tCrScaleAViewAsC); ++i) {
+            tCrScaleAViewAsC(i) = scale_b * tCrScaleAViewAsC(i);
+          }
+        }
+      }
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+        mute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accumulation());
+        tiled_mma.accumulate_ = MP31::SQMMA::ScaleOut::One;
+      }
+
+      if (warp_idx == 0 && k_tile_count_tme > 0) {
+        pipeline.producer_acquire(smem_pipe_write);
+        uint32_t bar_id = pipeline.producer_get_barrier_id(smem_pipe_write);
+        copy(mainloop_params.tme_load_a.with(bar_id), tAgA(_,_,_,*k_tile_iter), tAsA(_,_,_,smem_pipe_write.index()));
+        copy(mainloop_params.tme_load_b.with(bar_id), tBgB(_,_,_,*k_tile_iter), tBsB(_,_,_,smem_pipe_write.index()));
+      }
+      // advance producer pipeline
+      ++smem_pipe_write;
+      ++k_tile_iter;
+      --k_tile_count_tme;
+
+      warpsquad_wait();
+      pipeline.consumer_release(smem_pipe_release);
+
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
+        ElementBlockScale scale_ab = tCrScaleAViewAsC(0);
+        accumulation.scale_if_needed(scale_ab);
+      }
+
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
+        accumulation.scale_if_needed(tCrScaleBViewAsC);
+      }
+
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
+        accumulation.scale_if_needed(tCrScaleAViewAsC);
+      }
+
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
+        accumulation.scale_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
+      }
+
+      // advance consumer pipeline
+      ++smem_pipe_read;
+      ++smem_pipe_release;
+    }
+
+    {
+      pipeline.consumer_wait(smem_pipe_read);
+
+      if (accumulation.prepare_if_needed()) {
+        tiled_mma.accumulate_ = MP31::SQMMA::ScaleOut::Zero;
+        ++scale_k_iter;
+        Tensor tCgScaleAViewAsC = thread_mma.partition_C(gScaleA(_, *scale_k_iter).compose(ScaleAViewAsCLayout)); // (MMA, MMA_M, MMA_N)
+        Tensor tCgScaleBViewAsC = thread_mma.partition_C(gScaleB(_, *scale_k_iter).compose(ScaleBViewAsCLayout)); // (MMA, MMA_M, MMA_N)
+
+        copy(tCgScaleAViewAsC, tCrScaleAViewAsC);
+        copy(tCgScaleBViewAsC, tCrScaleBViewAsC);
+        if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
+          tCrScaleAViewAsC(0) = tCrScaleAViewAsC(0) * tCrScaleBViewAsC(0);
+        }
+        if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
+          ElementBlockScale scale_a = tCrScaleAViewAsC(0);
+          MUTLASS_PRAGMA_UNROLL
+          for (int i = 0; i < size(tCrScaleBViewAsC); ++i) {
+            tCrScaleBViewAsC(i) = scale_a * tCrScaleBViewAsC(i);
+          }
+        }
+        if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
+          ElementBlockScale scale_b = tCrScaleBViewAsC(0);
+          MUTLASS_PRAGMA_UNROLL
+          for (int i = 0; i < size(tCrScaleAViewAsC); ++i) {
+            tCrScaleAViewAsC(i) = scale_b * tCrScaleAViewAsC(i);
+          }
+        }
+      }
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int k_block = 0; k_block < size<2>(tCrA); ++k_block) {
+        mute::gemm(tiled_mma, tCrA(_,_,k_block,smem_pipe_read.index()), tCrB(_,_,k_block,smem_pipe_read.index()), accumulation());
+        tiled_mma.accumulate_ = MP31::SQMMA::ScaleOut::One;
+      }
+
+      warpsquad_wait();
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
+        ElementBlockScale scale_ab = tCrScaleAViewAsC(0);
+        accumulation.scale_if_needed(scale_ab);
+      }
+
+      if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
+        accumulation.scale_if_needed(tCrScaleBViewAsC);
+      }
+
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
+        accumulation.scale_if_needed(tCrScaleAViewAsC);
+      }
+
+      if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
+        accumulation.scale_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
+      }
+
+    }
+
+    // Scale the unscaled accum before return
+    if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile == 1) {
+      ElementBlockScale scale_ab = tCrScaleAViewAsC(0);
+      accumulation.scale_residue_if_needed(scale_ab);
+    }
+    if constexpr (ScaleMsPerTile == 1 && ScaleNsPerTile > 1) {
+      accumulation.scale_residue_if_needed(tCrScaleBViewAsC);
+    }
+
+    if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile == 1) {
+      accumulation.scale_residue_if_needed(tCrScaleAViewAsC);
+    }
+
+    if constexpr (ScaleMsPerTile > 1 && ScaleNsPerTile > 1) {
+      accumulation.scale_residue_if_needed(tCrScaleAViewAsC, tCrScaleBViewAsC);
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::collective
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/scaling_accumulation.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/scaling_accumulation.hpp
new file mode 100644
index 000000000..f4d7a17c1
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/collective/scaling_accumulation.hpp
@@ -0,0 +1,211 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+
+#pragma once
+
+#include "mute/algorithm/clear.hpp"
+#include "mute/tensor.hpp"
+
+namespace mutlass::gemm::collective {
+
+template <
+  class EngineAccum,
+  class LayoutAccum>
+struct ScalingAccumulation {
+  using TensorAccum = mute::Tensor<EngineAccum, LayoutAccum>;
+  using ElementAccumulator = typename EngineAccum::value_type;
+
+  static_assert(is_static<LayoutAccum>::value, "Accumulator Layout should be static");
+  static_assert(is_rmem<TensorAccum>::value, "Accumulator tensor must be rmem resident.");
+
+private:
+  TensorAccum& accum_;
+  TensorAccum accum_temp_;
+
+  uint32_t accum_promotion_interval_;         // defines the max num of executed MMAs after which accum should be promoted.
+  uint32_t mma_count_per_mainloop_iteration_; // num of MMAs per k_tile of mainloop
+  uint32_t mma_count_;                        // current executed MMAs
+  uint32_t reset_accum_flag_;                 // accum needs to be zeroed or not.
+
+  template <
+    class EngineScale,
+    class LayoutScale>
+  MUTLASS_DEVICE
+  void scale_core(const mute::Tensor<EngineScale, LayoutScale> &scale) {
+    using TensorScale = mute::Tensor<EngineScale, LayoutScale>;
+
+    static_assert(is_static<LayoutScale>::value, "Scale Layout should be static");
+    static_assert(is_rmem<TensorScale>::value , "Scale tensor must be rmem resident.");
+
+    static_assert(LayoutAccum{}.shape() == LayoutScale{}.shape(), "Accumulator and scale must have same shape.");
+
+    warpsquad_wait();
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accum_); ++i) {
+      accum_(i) += accum_temp_(i) * scale(i);
+    }
+  }
+
+  template <
+    class EngineScaleA,
+    class LayoutScaleA,
+    class EngineScaleB,
+    class LayoutScaleB>
+  MUTLASS_DEVICE
+  void scale_core(const mute::Tensor<EngineScaleA, LayoutScaleA> &scaleA,
+                  const mute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
+    using TensorScaleA = mute::Tensor<EngineScaleA, LayoutScaleA>;
+    using TensorScaleB = mute::Tensor<EngineScaleB, LayoutScaleB>;
+
+    static_assert(is_static<LayoutScaleA>::value, "ScaleA Layout should be static");
+    static_assert(is_static<LayoutScaleB>::value, "ScaleB Layout should be static");
+    static_assert(is_rmem<TensorScaleA>::value , "ScaleA tensor must be rmem resident.");
+    static_assert(is_rmem<TensorScaleB>::value , "ScaleB tensor must be rmem resident.");
+
+    static_assert(LayoutAccum{}.shape() == LayoutScaleA{}.shape(), "Accumulator and scaleA must have same shape.");
+    static_assert(LayoutAccum{}.shape() == LayoutScaleB{}.shape(), "Accumulator and scaleB must have same shape.");
+
+    warpsquad_wait();
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size(accum_); ++i) {
+      accum_(i) += accum_temp_(i) * scaleA(i) * scaleB(i);
+    }
+  }
+
+
+public:
+  MUTLASS_DEVICE
+  ScalingAccumulation(
+      TensorAccum &accum,
+      uint32_t accum_promotion_interval,
+      uint32_t mma_count_per_mainloop_iteration)
+      : accum_(accum),
+        accum_promotion_interval_(accum_promotion_interval),
+        mma_count_per_mainloop_iteration_(mma_count_per_mainloop_iteration),
+        mma_count_(0),
+        reset_accum_flag_(0)
+  {
+    accum_temp_ = mute::make_fragment_like(accum);
+  }
+
+  //
+  // Methods (Common)
+  //
+  MUTLASS_DEVICE
+  TensorAccum& operator()() {
+    return accum_temp_;
+  }
+
+  MUTLASS_DEVICE
+  bool prepare_if_needed() {
+    return reset_accum_flag_;
+  }
+
+  /// scale (multiply_add) the results from the MMA accumulators to main accumulator if needed.
+
+  template <class ElementScale>
+  MUTLASS_DEVICE
+  void scale_if_needed(ElementScale const &scale) {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = mma_count_ == accum_promotion_interval_;
+    if (reset_accum_flag_) {
+      scale_core(scale);
+      mma_count_ = 0;
+    }
+  }
+
+  template <
+    class EngineScale,
+    class LayoutScale>
+  MUTLASS_DEVICE
+  void scale_if_needed(const mute::Tensor<EngineScale, LayoutScale> &scale) {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = mma_count_ == accum_promotion_interval_;
+    if (reset_accum_flag_) {
+      scale_core(scale);
+      mma_count_ = 0;
+    }
+  }
+
+  template <
+    class EngineScaleA,
+    class LayoutScaleA,
+    class EngineScaleB,
+    class LayoutScaleB>
+  MUTLASS_DEVICE
+  void scale_if_needed(const mute::Tensor<EngineScaleA, LayoutScaleA> &scaleA,
+                       const mute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
+    mma_count_ += mma_count_per_mainloop_iteration_;
+    reset_accum_flag_ = mma_count_ == accum_promotion_interval_;
+    if (reset_accum_flag_) {
+      scale_core(scaleA, scaleB);
+      mma_count_ = 0;
+    }
+  }
+
+
+
+  /// scale (multiply_add) the residue results from the MMA accumulators to main accumulator if needed.
+  template <class ElementScale>
+  MUTLASS_DEVICE
+  void scale_residue_if_needed(ElementScale const &scale) {
+    if (mma_count_ > 0) {
+      scale_core(scale);
+    }
+  }
+
+  template <
+    class EngineScale,
+    class LayoutScale>
+  MUTLASS_DEVICE
+  void scale_residue_if_needed(const mute::Tensor<EngineScale, LayoutScale> &scale) {
+    if (mma_count_ > 0) {
+      scale_core(scale);
+    }
+  }
+
+  template <
+    class EngineScaleA,
+    class LayoutScaleA,
+    class EngineScaleB,
+    class LayoutScaleB>
+  MUTLASS_DEVICE
+  void scale_residue_if_needed(const mute::Tensor<EngineScaleA, LayoutScaleA> &scaleA,
+                               const mute::Tensor<EngineScaleB, LayoutScaleB> &scaleB) {
+    if (mma_count_ > 0) {
+      scale_core(scaleA, scaleB);
+    }
+  }
+};
+
+} // namespace mutlass::gemm::collective
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/device/gemm_universal_adapter.h b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/device/gemm_universal_adapter.h
new file mode 100644
index 000000000..37cc5b226
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/device/gemm_universal_adapter.h
@@ -0,0 +1,426 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 The universal GEMM accommodates serial reductions, parallel reductions, batched strided, and
+    batched array variants.
+*/
+
+#pragma once
+
+// common
+#include "mutlass/mutlass.h"
+#include "mutlass/device_kernel.h"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/detail/layout.hpp"
+#include "mutlass/detail/mma.hpp"
+#include "mutlass/musa_host_adapter.hpp"
+
+#if !defined(__MUSACC_RTC__)
+//#include "mutlass/cluster_launch.hpp"
+#include "mutlass/trace.h"
+#endif // !defined(__MUSACC_RTC__)
+
+#include "mutlass/gemm/threadblock/threadblock_swizzle.h" // placeholder for thread block swizzle
+
+// 3.x
+#include "mutlass/arch/mma.h" // mutlass::arch::OpMultiplyAdd
+#include "mutlass/gemm/kernel/gemm_universal.hpp"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::device {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/*!
+  GemmUniversalAdapter is a stateful, reusable GEMM handle built around a kernel
+  of type mutlass::gemm::kernel::Gemm or mutlass::gemm::kernel::GemmUniversal.
+
+  It manages the lifetime of the underlying `kernel::Params` struct, and exposes APIs
+  to create it from the host facing arguments. For power users, new static methods
+  are exposed in 3.x APIs that bypass the stateful methods or args->params lowering.
+
+  It supports kernel types that implement both the 2.x and 3.0 APIs,
+  however, this is done by specializing the implementation of GemmUniversalAdapter
+  on the two kernel API types, and thus, GemmUniversalAdapter's behaviour might
+  differ between the two specializations.
+*/
+template <class GemmKernel_, class Enable = void>
+class GemmUniversalAdapter;
+
+////////////////////////////////////////////////////////////////////////////////
+////////////////////////////// MUTLASS 3.x API /////////////////////////////////
+////////////////////////////////////////////////////////////////////////////////
+
+template <class GemmKernel_>
+class GemmUniversalAdapter<
+  GemmKernel_,
+  mute::enable_if_t<gemm::detail::IsMutlass3GemmKernel<GemmKernel_>::value>>
+{
+public:
+  using GemmKernel = GemmKernel_;
+  using TileShape = typename GemmKernel::TileShape;
+  using ElementA = typename GemmKernel::ElementA;
+  using ElementB = typename GemmKernel::ElementB;
+  // using ElementC = typename GemmKernel::ElementC;
+  // using ElementC = mute::conditional_t<mute::is_void_v<GemmKernel::ElementC>, ElementD, GemmKernel::ElementC>;
+  using ElementD = typename GemmKernel::ElementD;
+  using ElementC = mute::conditional_t<mute::is_void_v<typename GemmKernel::ElementC>, ElementD, typename GemmKernel::ElementC>;
+  using ElementAccumulator = typename GemmKernel::ElementAccumulator;
+  using DispatchPolicy = typename GemmKernel::DispatchPolicy;
+  using CollectiveMainloop = typename GemmKernel::CollectiveMainloop;
+  using CollectiveEpilogue = typename GemmKernel::CollectiveEpilogue;
+
+  // Map back to 2.x type as best as possible
+  using LayoutA = gemm::detail::StrideToLayoutTagA_t<typename GemmKernel::StrideA>;
+  using LayoutB = gemm::detail::StrideToLayoutTagB_t<typename GemmKernel::StrideB>;
+  using LayoutC = gemm::detail::StrideToLayoutTagC_t<typename GemmKernel::StrideC>;
+  using LayoutD = gemm::detail::StrideToLayoutTagC_t<typename GemmKernel::StrideD>;
+
+  static bool const kEnableMusaHostAdapter = MUTLASS_ENABLE_MUSA_HOST_ADAPTER;
+
+  static ComplexTransform const kTransformA = mute::is_same_v<typename GemmKernel::CollectiveMainloop::TransformA, mute::conjugate> ?
+                                              ComplexTransform::kConjugate : ComplexTransform::kNone;
+  static ComplexTransform const kTransformB = mute::is_same_v<typename GemmKernel::CollectiveMainloop::TransformB, mute::conjugate> ?
+                                              ComplexTransform::kConjugate : ComplexTransform::kNone;
+
+  // Legacy: Assume MultiplyAdd only since we do not use this tag type in 3.0
+  using MathOperator = mutlass::arch::OpMultiplyAdd;
+
+  using OperatorClass = mutlass::detail::get_operator_class_t<typename CollectiveMainloop::TiledMma>;
+
+  using ArchTag = typename GemmKernel::ArchTag;
+
+  // NOTE: Assume identity swizzle for now
+  using ThreadblockSwizzle = mutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
+
+  // Assume TiledMma's ShapeMNK is the same as 2.x's ThreadblockShape
+  using ThreadblockShape = mutlass::gemm::GemmShape<
+      mute::size<0>(TileShape{}),
+      mute::size<1>(TileShape{}),
+      mute::size<2>(TileShape{})>;
+
+  using ClusterShape = mutlass::gemm::GemmShape<
+      mute::size<0>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
+      mute::size<1>(typename GemmKernel::DispatchPolicy::ClusterShape{}),
+      mute::size<2>(typename GemmKernel::DispatchPolicy::ClusterShape{})>;
+
+  // Instruction shape is easy too, since we get that directly from our TiledMma's atom shape
+  using InstructionShape = mutlass::gemm::GemmShape<
+      mute::size<0>(typename CollectiveMainloop::TiledMma::AtomShape_MNK{}),
+      mute::size<1>(typename CollectiveMainloop::TiledMma::AtomShape_MNK{}),
+      mute::size<2>(typename CollectiveMainloop::TiledMma::AtomShape_MNK{})>;
+
+  // Legacy: provide a correct warp count, but no reliable warp shape
+  static int const kThreadCount = GemmKernel::MaxThreadsPerBlock;
+
+  // Warp shape is not a primary API type in 3.x
+  // But we can best approximate it by inspecting the TiledMma
+  // For this, we make the assumption that we always have 4 warps along M, and rest along N, none along K
+  // We also always round up the warp count to 4 if the tiled mma is smaller than 128 threads
+  static constexpr int WarpsInMma = mute::max(4, MUTE_STATIC_V(mute::size(typename GemmKernel::TiledMma{})) / 32);
+  static constexpr int WarpsInMmaM = 4;
+  static constexpr int WarpsInMmaN = mute::ceil_div(WarpsInMma, WarpsInMmaM);
+  using WarpCount = mutlass::gemm::GemmShape<WarpsInMmaM, WarpsInMmaN, 1>;
+  using WarpShape = mutlass::gemm::GemmShape<
+      MUTE_STATIC_V(mute::tile_size<0>(typename CollectiveMainloop::TiledMma{})) / WarpsInMmaM,
+      MUTE_STATIC_V(mute::tile_size<1>(typename CollectiveMainloop::TiledMma{})) / WarpsInMmaN,
+      MUTE_STATIC_V(mute::tile_size<2>(typename CollectiveMainloop::TiledMma{}))>;
+
+  static int constexpr kStages = CollectiveMainloop::DispatchPolicy::Stages;
+
+  // Inspect TiledCopy for A and B to compute the alignment size
+  static int constexpr kAlignmentA = mutlass::detail::get_alignment_count_from_gmem_tiled_copy<
+      typename CollectiveMainloop::GmemTiledCopyA, ElementA, typename CollectiveMainloop::TiledMma::ValTypeA>();
+  static int constexpr kAlignmentB = mutlass::detail::get_alignment_count_from_gmem_tiled_copy<
+      typename CollectiveMainloop::GmemTiledCopyB, ElementB, typename CollectiveMainloop::TiledMma::ValTypeB>();
+  static int constexpr kAlignmentC = mutlass::detail::get_alignment_count_from_gmem_tiled_copy<
+      typename CollectiveEpilogue::GmemTiledCopyC, ElementC>();
+  static int constexpr kAlignmentD = mutlass::detail::get_alignment_count_from_gmem_tiled_copy<
+      typename CollectiveEpilogue::GmemTiledCopyD, ElementD>();
+
+  using EpilogueOutputOp = typename CollectiveEpilogue::ThreadEpilogueOp;
+
+  // Split-K preserves splits that are 128b aligned
+  static int constexpr kSplitKAlignment = mute::max(
+      128 / sizeof_bits<ElementA>::value, 128 / sizeof_bits<ElementB>::value);
+
+  /// Argument structure: User API
+  using Arguments = typename GemmKernel::Arguments;
+  /// Argument structure: Kernel API
+  using Params = typename GemmKernel::Params;
+
+private:
+
+  /// Kernel API parameters object
+  Params params_;
+
+public:
+
+  /// Access the Params structure
+  Params const& params() const {
+    return params_;
+  }
+
+  /// Determines whether the GEMM can execute the given problem.
+  static Status
+  can_implement(Arguments const& args) {
+    if (GemmKernel::can_implement(args)) {
+      return Status::kSuccess;
+    }
+    else {
+      return Status::kInvalid;
+    }
+  }
+
+  /// Gets the workspace size
+  static size_t
+  get_workspace_size(Arguments const& args) {
+    size_t workspace_bytes = 0;
+    if (args.mode == GemmUniversalMode::kGemmSplitKParallel) {
+      workspace_bytes += sizeof(int) * size_t(mute::size<0>(TileShape{})) * size_t(mute::size<1>(TileShape{}));
+    }
+
+    MUTLASS_TRACE_HOST("  workspace_bytes: " << workspace_bytes);
+
+    workspace_bytes += GemmKernel::get_workspace_size(args);
+    return workspace_bytes;
+  }
+
+  /// Computes the grid shape
+  static dim3
+  get_grid_shape(Arguments const& args, void* workspace = nullptr) {
+    auto tmp_params = GemmKernel::to_underlying_arguments(args, workspace);
+    return GemmKernel::get_grid_shape(tmp_params);
+  }
+
+  /// Computes the grid shape
+  static dim3
+  get_grid_shape(Params const& params) {
+    return GemmKernel::get_grid_shape(params);
+  }
+
+  /// Computes the maximum number of active blocks per multiprocessor
+  static int maximum_active_blocks(int /* smem_capacity */ = -1) {
+    MUTLASS_TRACE_HOST("GemmUniversal::maximum_active_blocks()");
+    int max_active_blocks = -1;
+    int smem_size = GemmKernel::SharedStorageSize;
+
+    // first, account for dynamic smem capacity if needed
+    musaError_t result;
+    if (smem_size >= (48 << 10)) {
+      MUTLASS_TRACE_HOST("  Setting smem size to " << smem_size);
+      result = musaFuncSetAttribute(
+          device_kernel<GemmKernel>,
+          musaFuncAttributeMaxDynamicSharedMemorySize,
+          smem_size);
+      if (musaSuccess != result) {
+        result = musaGetLastError(); // to clear the error bit
+        MUTLASS_TRACE_HOST(
+          "  musaFuncSetAttribute() returned error: "
+          << musaGetErrorString(result));
+        return -1;
+      }
+    }
+
+    // query occupancy after setting smem size
+    result = musaOccupancyMaxActiveBlocksPerMultiprocessor(
+        &max_active_blocks,
+        device_kernel<GemmKernel>,
+        GemmKernel::MaxThreadsPerBlock,
+        smem_size);
+
+    if (musaSuccess != result) {
+      result = musaGetLastError(); // to clear the error bit
+      MUTLASS_TRACE_HOST(
+        "  musaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: "
+        << musaGetErrorString(result));
+      return -1;
+    }
+
+    MUTLASS_TRACE_HOST("  max_active_blocks: " << max_active_blocks);
+    return max_active_blocks;
+  }
+
+  /// Initializes GEMM state from arguments.
+  Status
+  initialize(
+    Arguments const& args,
+    void* workspace = nullptr,
+    musaStream_t stream = nullptr,
+    MusaHostAdapter* musa_adapter = nullptr) {
+
+    MUTLASS_TRACE_HOST("GemmUniversal::initialize() - workspace "
+      << workspace << ", stream: " << (stream ? "non-null" : "null"));
+
+    // Initialize the workspace
+    Status status = GemmKernel::initialize_workspace(args, workspace, stream, musa_adapter);
+    if (status != Status::kSuccess) {
+      return status;
+    }
+    // Initialize the Params structure
+    params_ = GemmKernel::to_underlying_arguments(args, workspace);
+    // Don't set the function attributes - require the MusaHostAdapter to set it.
+    if constexpr (kEnableMusaHostAdapter) {
+      MUTLASS_ASSERT(musa_adapter);
+      return Status::kSuccess;
+    }
+    else {
+      //
+      // Account for dynamic smem capacity if needed
+      //
+      int smem_size = GemmKernel::SharedStorageSize;
+
+      MUTLASS_ASSERT(musa_adapter == nullptr);
+
+      if (smem_size >= (48 << 10)) {
+        MUTLASS_TRACE_HOST("  Setting smem size to " << smem_size);
+        musaError_t result = musaFuncSetAttribute(
+            device_kernel<GemmKernel>,
+            musaFuncAttributeMaxDynamicSharedMemorySize,
+            smem_size);
+        if (musaSuccess != result) {
+          result = musaGetLastError(); // to clear the error bit
+          MUTLASS_TRACE_HOST("  musaFuncSetAttribute() returned error: " << musaGetErrorString(result));
+          return Status::kErrorInternal;
+        }
+      }
+    }
+    return Status::kSuccess;
+  }
+
+  /// Update API is preserved in 3.0, but does not guarantee a lightweight update of params.
+  Status
+  update(Arguments const& args, void* workspace = nullptr) {
+    MUTLASS_TRACE_HOST("GemmUniversal()::update() - workspace: " << workspace);
+
+    size_t workspace_bytes = get_workspace_size(args);
+    if (workspace_bytes > 0 && nullptr == workspace) {
+      return Status::kErrorWorkspaceNull;
+    }
+
+    params_ = GemmKernel::to_underlying_arguments(args, workspace);
+    return Status::kSuccess;
+  }
+
+  /// Primary run() entry point API that is static allowing users to create and manage their own params.
+  /// Supplied params struct must be construct by calling GemmKernel::to_underling_arguments()
+  static Status
+  run(Params& params,
+      musaStream_t stream = nullptr,
+      MusaHostAdapter *musa_adapter = nullptr) {
+    MUTLASS_TRACE_HOST("GemmUniversal::run()");
+    dim3 const block = GemmKernel::get_block_shape();
+    dim3 const grid = get_grid_shape(params);
+
+    // configure smem size and carveout
+    int smem_size = GemmKernel::SharedStorageSize;
+
+    Status launch_result{ Status::kSuccess };
+    // Use extended launch API only for mainloops that use it
+    launch_result = Status::kSuccess;
+    if constexpr (kEnableMusaHostAdapter) {
+      MUTLASS_ASSERT(musa_adapter);
+      if (musa_adapter) {
+        void* kernel_params[] = {&params};
+
+        launch_result = musa_adapter->launch(
+          grid, block, smem_size, stream, kernel_params, 0
+        );
+
+      }
+      else {
+        return Status::kErrorInternal;
+      }
+    }
+    else {
+      MUTLASS_ASSERT(musa_adapter == nullptr);
+      device_kernel<GemmKernel><<<grid, block, smem_size, stream>>>(params);
+    }
+
+    musaError_t result = musaGetLastError();
+    if (musaSuccess == result && Status::kSuccess == launch_result) {
+      return Status::kSuccess;
+    }
+    else {
+      MUTLASS_TRACE_HOST("  Kernel launch failed. Reason: " << result);
+      return Status::kErrorInternal;
+    }
+  }
+
+  //
+  // Non-static launch overloads that first create and set the internal params struct of this kernel handle.
+  //
+
+  /// Launches the kernel after first constructing Params internal state from supplied arguments.
+  Status
+  run(
+    Arguments const& args,
+    void* workspace = nullptr,
+    musaStream_t stream = nullptr,
+    MusaHostAdapter *musa_adapter = nullptr
+  ) {
+    Status status = initialize(args, workspace, stream, musa_adapter);
+
+    if (Status::kSuccess == status) {
+      status = run(params_, stream, musa_adapter);
+    }
+    return status;
+  }
+
+  /// Launches the kernel after first constructing Params internal state from supplied arguments.
+  Status
+  operator()(
+    Arguments const& args,
+    void* workspace = nullptr,
+    musaStream_t stream = nullptr,
+    MusaHostAdapter *musa_adapter = nullptr) {
+    return run(args, workspace, stream, musa_adapter);
+  }
+
+  /// Overload that allows a user to re-launch the same kernel without updating internal params struct.
+  Status
+  run(musaStream_t stream = nullptr, MusaHostAdapter *musa_adapter = nullptr) {
+    return run(params_, stream, musa_adapter);
+  }
+
+  /// Overload that allows a user to re-launch the same kernel without updating internal params struct.
+  Status
+  operator()(musaStream_t stream = nullptr, MusaHostAdapter *musa_adapter = nullptr) {
+    return run(params_, stream, musa_adapter);
+  }
+};
+
+} // namespace mutlass::gemm::device
+
+////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/dispatch_policy.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/dispatch_policy.hpp
new file mode 100644
index 000000000..cc3b65343
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/dispatch_policy.hpp
@@ -0,0 +1,110 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/arch/arch.h"
+#include "mutlass/gemm/gemm.h"
+
+#include "mute/layout.hpp"
+#include "mute/numeric/integral_constant.hpp"
+//////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm {
+using namespace mute;
+
+//////////////////////////////////////////////////////////////////////////////
+
+//
+// Policies for categorical dispatch of mainloop against kernel grid schedules
+//
+struct KernelMultistage { };
+struct KernelTme { };
+
+template <
+  int ScaleGranularityM = 0,
+  int ScaleGranularityN = 0,
+  int ScaleGranularityK = 0
+>
+struct KernelTmeGroupScaledAccum : KernelTme { };
+
+// Policies for dispatch of epilogue
+struct EpilogueDefault { };
+struct EpilogueTransposed { };
+
+//
+// Collective Mainloop Policies
+//
+
+// 2 stage pipeline through 1 stage in smem, 1 in rmem, WITHOUT predicated gmem loads
+struct MainloopMp22TwoStageUnpredicated {
+  constexpr static int Stages = 2;
+  using ArchTag = arch::Mp22;
+  using Schedule = KernelMultistage;
+  using ClusterShape = Shape<_1,_1,_1>;
+};
+
+// 2 stage pipeline through 1 stage in smem, 1 in rmem, with predicated gmem loads
+struct MainloopMp22TwoStage {
+  constexpr static int Stages = 2;
+  using ArchTag = arch::Mp22;
+  using Schedule = KernelMultistage;
+  using ClusterShape = Shape<_1,_1,_1>;
+};
+
+// n-buffer in smem(TME), pipelined with SQMMA and TME, static schedule between TME and SQMMA
+template <
+  int Stages_,
+  class KernelSchedule = KernelTme>
+struct MainloopMp31TmeSqmma {
+  constexpr static int Stages = Stages_;
+  using ArchTag = arch::Mp31;
+  using Schedule = KernelSchedule;
+  using ClusterShape = Shape<_1,_1,_1>;
+};
+
+template <
+  int Stages_,
+  class KernelSchedule = KernelTme,
+  int ScaleGranularityM = 0, // `ScaleGranularityM` specifies scaling granularity along M, while zero-value `ScaleGranularityM` indicates that scaling granularity is `size<0>(TileShape_MNK{})` along M
+  int ScaleGranularityN = 0, // `ScaleGranularityN` specifies scaling granularity along N, while zero-value `ScaleGranularityN` indicates that scaling granularity is `size<1>(TileShape_MNK{})` along N
+  int ScaleGranularityK = 0
+>
+struct MainloopMp31TmeSqmmaBlockScalingFP8
+  : MainloopMp31TmeSqmma<Stages_, KernelSchedule> {
+  static_assert(
+    mute::is_same_v<KernelSchedule, KernelTmeGroupScaledAccum<ScaleGranularityM, ScaleGranularityN, ScaleGranularityK>>,
+    "Kernel Schedule must be the scaled accum policy.");
+};
+
+//////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/gemm.h b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/gemm.h
new file mode 100644
index 000000000..718b2d8e2
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/gemm.h
@@ -0,0 +1,134 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 common types used for all GEMM-like operators.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+#include "mutlass/gemm_coord.h"
+#include "mutlass/layout/matrix.h"
+#include "mutlass/gemm/gemm_enumerated_types.h"
+#include "mute/layout.hpp"
+#include "mutlass/detail/layout.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace gemm {
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+using mutlass::detail::TagToStrideA;
+using mutlass::detail::TagToStrideB;
+using mutlass::detail::TagToStrideC;
+using mutlass::detail::TagToStrideA_t;
+using mutlass::detail::TagToStrideB_t;
+using mutlass::detail::TagToStrideC_t;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+using mutlass::detail::StrideToLayoutTagA;
+using mutlass::detail::StrideToLayoutTagB;
+using mutlass::detail::StrideToLayoutTagC;
+using mutlass::detail::StrideToLayoutTagA_t;
+using mutlass::detail::StrideToLayoutTagB_t;
+using mutlass::detail::StrideToLayoutTagC_t;
+
+template<int ModeIndex, class Stride>
+constexpr bool
+is_major(Stride = {}) {
+  return ::mutlass::detail::is_major<ModeIndex, Stride>();
+}
+
+template<class Stride>
+constexpr bool
+is_mn_major() {
+  return is_major<0,Stride>();
+}
+
+template<class Stride>
+constexpr
+bool
+is_k_major() {
+  return is_major<1,Stride>();
+}
+
+template<class LayoutA>
+constexpr bool
+is_mn_major_A() {
+  return is_mn_major<TagToStrideA_t<LayoutA>>();
+}
+
+template<class LayoutB>
+constexpr bool
+is_mn_major_B() {
+  return is_mn_major<TagToStrideB_t<LayoutB>>();
+}
+
+template<class LayoutA>
+constexpr bool
+is_k_major_A() {
+  return is_k_major<TagToStrideA_t<LayoutA>>();
+}
+
+template<class LayoutB>
+constexpr bool
+is_k_major_B() {
+  return is_k_major<TagToStrideB_t<LayoutB>>();
+}
+
+///////////////////////////////////////////////////////////////////////////////
+
+// The following two metafunctions are used to detect whether a `kernel::Gemm` or `kernel::GemmUniversal`
+// is implementing the MUTLASS 3.x API or not, by checking if the problem shape type is aliased within or not.
+template <class GemmKernel, class = void>
+struct IsMutlass3GemmKernel : mute::false_type { };
+
+template <typename GemmKernel>
+struct IsMutlass3GemmKernel<GemmKernel, mute::void_t<typename GemmKernel::ProblemShape>>
+    : mute::true_type { };
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace gemm
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/gemm_enumerated_types.h b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/gemm_enumerated_types.h
new file mode 100644
index 000000000..1c72252dd
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/gemm_enumerated_types.h
@@ -0,0 +1,90 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 common types used for all GEMM-like operators.
+*/
+
+/*
+  Note:  MUTLASS 3x increases the host compiler requirements to C++17. However, certain
+         existing integrations of MUTLASS require C++11 host compilers.
+
+         Until this requirement can be lifted, certain headers with this annotation are required
+         to be remain consistent with C++11 syntax.
+
+         C++11 compatibility is enforced by `mutlass_test_unit_core_cpp11`.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+#include "mutlass/gemm_coord.h"
+#include "mutlass/layout/matrix.h"
+
+namespace mutlass {
+namespace gemm {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// GEMM operand enumeration: D = A * B + C
+enum class Operand {
+  kA, /// A multiplicand
+  kB, /// B multiplicand
+  kC, /// Source accumulator
+  kD  /// Destination accumulator
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+enum class GemmUniversalMode {
+  kGemm,
+  kGemmSplitKParallel,
+  kBatched,
+  kArray,
+  kInvalid
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+/// Some options for clearing shared memory
+enum class SharedMemoryClearOption {
+  kNone,            ///< SMEM is in don't-care state
+  kZfill,           ///< Kernels fill out of bounds accesses with zeros
+  kClearLastStage   ///< Last SMEM stage is explicitly cleared. Mainloop uses 'kNone'
+};
+
+/////////////////////////////////////////////////////////////////////////
+
+} // namespace gemm
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/gemm_universal.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/gemm_universal.hpp
new file mode 100644
index 000000000..e157ceef1
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/gemm_universal.hpp
@@ -0,0 +1,73 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/gemm/kernel/tile_scheduler.hpp"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::kernel {
+
+////////////////////////////////////////////////////////////////////////////////
+
+/*
+ * Stateless universal device GEMM kernel type that treats GEMM as
+ * a composition of a collective mainloop and a collective epilogue.
+ *
+ * Supports both the 2.x and 3.x APIs based on whether the first type is
+ * a mute::tuple<> or not.
+ * 2.x API implementation: mutlass/gemm/kernel/gemm_universal.h
+ * 3.x API implementation: mutlass/gemm/kernel/gemm_*.hpp
+ *
+ * In the following declaration, the name preceding the 'Or' refers to
+ * 3.x API type argument order, and the name succeeding the 'Or' refers to
+ * 2.x API type argument order. Template arguments without two names
+ * belong to the 3.x API only.
+**/
+template <
+  class ProblemShapeOrThreadblockMma_, // (m, n, k) or (m, n, k, l)
+  class CollectiveMainloopOrEpilogue_,
+  class CollectiveEpilogueOrThreadblockSwizzle_,
+  class TileScheduler_ = void,
+  class Enable = void
+>
+class GemmUniversal;
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::kernel
+
+////////////////////////////////////////////////////////////////////////////////
+
+#include "mutlass/gemm/kernel/mp22_gemm.hpp"
+#include "mutlass/gemm/kernel/mp31_gemm_tme.hpp"
+////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/mp22_gemm.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/mp22_gemm.hpp
new file mode 100644
index 000000000..a7e6eacd5
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/mp22_gemm.hpp
@@ -0,0 +1,275 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/kernel_hardware_info.hpp"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+
+#include "mute/tensor.hpp"
+
+namespace mutlass::gemm::kernel {
+
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+  class ProblemShape_,
+  class CollectiveMainloop_,
+  class CollectiveEpilogue_,
+  class TileScheduler_
+>
+class GemmUniversal<
+  ProblemShape_,
+  CollectiveMainloop_,
+  CollectiveEpilogue_,
+  TileScheduler_,
+  mute::enable_if_t<mute::is_base_of_v<KernelMultistage, typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
+{
+public:
+  //
+  // Type Aliases
+  //
+  using ProblemShape = ProblemShape_;
+  static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
+    "ProblemShape{} should be <M,N,K> or <M,N,K,L>");
+
+  // Mainloop derived types
+  using CollectiveMainloop = CollectiveMainloop_;
+  using TileShape = typename CollectiveMainloop::TileShape;
+  using TiledMma  = typename CollectiveMainloop::TiledMma;
+  using ArchTag   = typename CollectiveMainloop::ArchTag;
+  using ElementA  = typename CollectiveMainloop::ElementA;
+  using StrideA   = typename CollectiveMainloop::StrideA;
+  using ElementB  = typename CollectiveMainloop::ElementB;
+  using StrideB   = typename CollectiveMainloop::StrideB;
+  using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
+  using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
+  using MainloopArguments = typename CollectiveMainloop::Arguments;
+  using MainloopParams = typename CollectiveMainloop::Params;
+
+  static_assert(mute::is_void_v<TileScheduler_>, "MP22 kernel does not support specializing the tile scheduler.");
+
+  // Epilogue derived types
+  using CollectiveEpilogue = CollectiveEpilogue_;
+  using ElementC = typename CollectiveEpilogue::ElementC;
+  using StrideC  = typename CollectiveEpilogue::StrideC;
+  using ElementD = typename CollectiveEpilogue::ElementD;
+  using StrideD  = typename CollectiveEpilogue::StrideD;
+  using EpilogueArguments = typename CollectiveEpilogue::Arguments;
+  using EpilogueParams = typename CollectiveEpilogue::Params;
+  static_assert(mute::is_same_v<ElementAccumulator, typename CollectiveEpilogue::ElementAccumulator>,
+    "Mainloop and epilogue do not agree on accumulator value type.");
+
+  // MSVC requires the cast to fix a warning-as-error.
+  static constexpr int SharedStorageSize = static_cast<int>(mute::max(
+      sizeof(typename CollectiveMainloop::SharedStorage),
+      sizeof(typename CollectiveEpilogue::SharedStorage)));
+
+  static constexpr uint32_t MaxThreadsPerBlock = MUTE_STATIC_V(mute::size(TiledMma{}));
+  static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
+
+  static constexpr int SmemAlignmentBytes = CollectiveMainloop::SmemAlignmentBytes;
+
+  // Device side arguments
+  struct Arguments {
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopArguments mainloop{};
+    EpilogueArguments epilogue{};
+    KernelHardwareInfo hw_info{};
+  };
+
+  // Kernel entry point API
+  struct Params {
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop{};
+    EpilogueParams epilogue{};
+  };
+
+  //
+  // Methods
+  //
+
+  // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+  static
+  Params
+  to_underlying_arguments(Arguments const& args, void* workspace) {
+    (void) workspace;
+
+    KernelHardwareInfo hw_info{args.hw_info.device_id, args.hw_info.sm_count};
+    auto problem_shape_MNKL = append<4>(args.problem_shape, Int<1>{});
+
+    return {
+      args.mode,
+      args.problem_shape,
+      CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
+      CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace)
+    };
+  }
+
+  static bool
+  can_implement(Arguments const& args) {
+    bool implementable =  (args.mode == GemmUniversalMode::kGemm) or
+         (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
+    if (!implementable) {
+      MUTLASS_TRACE_HOST("  CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
+      return implementable;
+    }
+
+    implementable = implementable && CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
+    implementable = implementable && CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+
+    return implementable;
+
+  }
+
+  static size_t
+  get_workspace_size(Arguments const& args) {
+    size_t workspace_size = 0;
+    return workspace_size;
+  }
+
+  static
+  mutlass::Status
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, musaStream_t stream = nullptr,
+    MusaHostAdapter* musa_adapter = nullptr) {
+    mutlass::Status status = Status::kSuccess;
+
+    return status;
+  }
+
+  static dim3
+  get_grid_shape(Params const& params) {
+    int batch_count = 1;
+    if constexpr (mute::rank(ProblemShape{}) == 4) {
+      batch_count = mute::size<3>(params.problem_shape);
+    }
+
+    return dim3(
+      mute::size(mute::ceil_div(mute::shape<0>(params.problem_shape), mute::shape<0>(TileShape{}))),
+      mute::size(mute::ceil_div(mute::shape<1>(params.problem_shape), mute::shape<1>(TileShape{}))),
+      batch_count
+    );
+  }
+
+  static dim3
+  get_block_shape() {
+    return dim3(MaxThreadsPerBlock, 1, 1);
+  }
+
+  MUTLASS_DEVICE
+  void
+  operator()(Params const& params, char* smem_buf) {
+    using namespace mute;
+    using X = Underscore;
+
+    // Preconditions
+    MUTE_STATIC_ASSERT(is_static<TileShape>::value);
+
+    // Separate out problem shape for convenience
+    // Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
+    auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
+    auto M = get<0>(problem_shape_MNKL);
+    auto N = get<1>(problem_shape_MNKL);
+    auto K = get<2>(problem_shape_MNKL);
+    auto L = get<3>(problem_shape_MNKL);
+
+    // Preconditions
+    static_assert(mute::rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(mute::rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(mute::rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(mute::rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
+
+    // Get the appropriate blocks for this thread block -- potential for thread block locality
+    int thread_idx = int(threadIdx.x);
+    auto blk_shape = TileShape{};                                                                // (BLK_M,BLK_N,BLK_K)
+    auto [m_coord, n_coord, l_coord] = static_cast<uint3>(blockIdx);
+    auto blk_coord_mnkl = make_coord(m_coord, n_coord, _, l_coord);                                        // (m,n,k,l)
+
+    // Represent the full tensors
+    Tensor mA_mkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_A), make_shape(M,K,L), params.mainloop.dA); //(m,k,l)
+    Tensor mB_nkl = make_tensor(make_gmem_ptr(params.mainloop.ptr_B), make_shape(N,K,L), params.mainloop.dB); //(n,k,l)
+
+    // Get batch slice
+    Tensor mA_mk = mA_mkl(_,_,l_coord);                                                                        // (m,k)
+    Tensor mB_nk = mB_nkl(_,_,l_coord);                                                                        // (n,k)
+
+    // Slice to get the tiles this thread block is responsible for
+    Tensor gA = local_tile(mA_mk, blk_shape, take<0,3>(blk_coord_mnkl), Step<_1, X,_1>{});           // (BLK_M,BLK_K,k)
+    Tensor gB = local_tile(mB_nk, blk_shape, take<0,3>(blk_coord_mnkl), Step< X,_1,_1>{});           // (BLK_N,BLK_K,k)
+
+    // Compute tile residues for predication
+    auto m_max_coord = M - size<0>(gA) * get<0>(blk_coord_mnkl);                             // M - BLK_M * m_coord
+    auto n_max_coord = N - size<0>(gB) * get<1>(blk_coord_mnkl);                             // N - BLK_N * n_coord
+    auto k_residue   = K - size<1>(gA) * size<2>(gA);                                        // K - BLK_K * k_coord_max
+    auto residue_mnk = make_tuple(m_max_coord, n_max_coord, k_residue);
+
+    // Allocate the tiled_mma and the accumulators for the (M,N) blk_shape
+    TiledMma tiled_mma;
+    Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N)
+    clear(accumulators);
+
+    auto k_tile_iter  = mute::make_coord_iterator(shape<2>(gA));
+    int  k_tile_count = size<2>(gA);
+
+    // Perform the collective scoped MMA
+    CollectiveMainloop collective_mma;
+    collective_mma(
+      accumulators,
+      gA,
+      gB,
+      accumulators,
+      k_tile_iter, k_tile_count,
+      residue_mnk,
+      thread_idx,
+      smem_buf
+    );
+    // Epilogue and write to gD
+    CollectiveEpilogue epilogue{params.epilogue};
+    epilogue(
+      problem_shape_MNKL,
+      blk_shape,
+      blk_coord_mnkl,
+      accumulators,
+      tiled_mma,
+      residue_mnk,
+      thread_idx,
+      smem_buf
+    );
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::kernel
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/mp31_gemm_tme.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/mp31_gemm_tme.hpp
new file mode 100644
index 000000000..eecd06d57
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/mp31_gemm_tme.hpp
@@ -0,0 +1,310 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/kernel_hardware_info.hpp"
+#include "mutlass/epilogue/collective/detail.hpp"
+#include "mutlass/gemm/kernel/tile_scheduler.hpp"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+#include "mutlass/trace.h"
+#include "mutlass/arch/barrier.hpp"
+
+#include "mute/tensor.hpp"
+
+///////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::kernel {
+
+///////////////////////////////////////////////////////////////////////////////
+
+template <
+  class ProblemShape_,
+  class CollectiveMainloop_,
+  class CollectiveEpilogue_,
+  class TileScheduler_
+>
+class GemmUniversal<
+  ProblemShape_,
+  CollectiveMainloop_,
+  CollectiveEpilogue_,
+  TileScheduler_,
+  mute::enable_if_t<mute::is_base_of_v<KernelTme, typename CollectiveMainloop_::DispatchPolicy::Schedule>>>
+{
+public:
+  //
+  // Type Aliases
+  //
+
+  using ProblemShape = ProblemShape_;
+
+  static_assert(rank(ProblemShape{}) == 3 or rank(ProblemShape{}) == 4,
+    "ProblemShape{} should be <M,N,K> or <M,N,K,L>");
+
+  // Mainloop derived types
+  using CollectiveMainloop = CollectiveMainloop_;
+  using TileShape = typename CollectiveMainloop::TileShape;
+  using TiledMma  = typename CollectiveMainloop::TiledMma;
+  using ArchTag   = typename CollectiveMainloop::ArchTag;
+  using ElementA  = typename CollectiveMainloop::ElementA;
+  using StrideA   = typename CollectiveMainloop::StrideA;
+  using ElementB  = typename CollectiveMainloop::ElementB;
+  using StrideB   = typename CollectiveMainloop::StrideB;
+  using DispatchPolicy = typename CollectiveMainloop::DispatchPolicy;
+  using ElementAccumulator = typename CollectiveMainloop::ElementAccumulator;
+  using ClusterShape = typename DispatchPolicy::ClusterShape;
+  using MainloopArguments = typename CollectiveMainloop::Arguments;
+  using MainloopParams = typename CollectiveMainloop::Params;
+
+  static_assert(ArchTag::kMinComputeCapability == 31);
+
+  // Epilogue derived types
+  using CollectiveEpilogue = CollectiveEpilogue_;
+  using ElementC = typename CollectiveEpilogue::ElementC;
+  using StrideC  = typename CollectiveEpilogue::StrideC;
+  using ElementD = typename CollectiveEpilogue::ElementD;
+  using StrideD  = typename CollectiveEpilogue::StrideD;
+  using EpilogueArguments = typename CollectiveEpilogue::Arguments;
+  using EpilogueParams = typename CollectiveEpilogue::Params;
+  //static_assert(mute::is_same_v<ElementAccumulator, typename CollectiveEpilogue::ElementAccumulator>,
+  //  "Mainloop and epilogue do not agree on accumulator value type.");
+  static_assert(mute::is_void_v<TileScheduler_> || mute::is_same_v<TileScheduler_, DefaultScheduler>,
+    "TME kernel only support default tile scheduler");
+
+  using TileSchedulerTag = TileScheduler_;
+  using TileScheduler = typename detail::TileSchedulerSelector<
+        TileScheduler_, ArchTag, TileShape, ClusterShape>::Scheduler;
+  using TileSchedulerArguments = typename TileScheduler::Arguments;
+  using TileSchedulerParams = typename TileScheduler::Params;
+
+  static constexpr int SharedStorageSize = static_cast<int>(mute::max(
+        sizeof(typename CollectiveMainloop::SharedStorage),
+        sizeof(typename CollectiveEpilogue::SharedStorage)));
+
+  static constexpr uint32_t MaxThreadsPerBlock = size(TiledMma{});
+  static constexpr uint32_t MinBlocksPerMultiprocessor = 1;
+
+  static constexpr int SmemAlignmentBytes = CollectiveMainloop::SmemAlignmentBytes;
+
+  static constexpr int NumBarriers = mute::max(CollectiveEpilogue::NumBarriers, CollectiveMainloop::NumBarriers);
+  // Device side arguments
+  struct Arguments {
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopArguments mainloop{};
+    EpilogueArguments epilogue{};
+    KernelHardwareInfo hw_info{};
+    TileSchedulerArguments scheduler{};
+  };
+
+
+  // Kernel entry point API
+  struct Params {
+    GemmUniversalMode mode{};
+    ProblemShape problem_shape{};
+    MainloopParams mainloop;
+    EpilogueParams epilogue;
+    KernelHardwareInfo hw_info{};
+    TileSchedulerParams scheduler{};
+  };
+
+  //
+  // Methods
+  //
+
+  // Convert to underlying arguments. In this case, a simple copy for the aliased type.
+  static
+  Params
+  to_underlying_arguments(Arguments const& args, void* workspace) {
+    (void) workspace;
+    KernelHardwareInfo hw_info{args.hw_info.device_id, args.hw_info.sm_count};
+    auto problem_shape_MNKL = append<4>(args.problem_shape, Int<1>{});
+    return {
+      args.mode,
+      args.problem_shape,
+      CollectiveMainloop::to_underlying_arguments(args.problem_shape, args.mainloop, workspace),
+      CollectiveEpilogue::to_underlying_arguments(args.problem_shape, args.epilogue, workspace),
+      hw_info,
+      TileScheduler::to_underlying_arguments(problem_shape_MNKL, TileShape{}, hw_info, args.scheduler)
+    };
+  }
+
+  MUTLASS_HOST_DEVICE static
+  bool
+  can_implement(Arguments const& args) {
+    bool implementable = (args.mode == GemmUniversalMode::kGemm) or
+        (args.mode == GemmUniversalMode::kBatched && rank(ProblemShape{}) == 4);
+    if (!implementable) {
+      MUTLASS_TRACE_HOST("  CAN IMPLEMENT: Arguments or Problem Shape don't meet the requirements.\n");
+      return implementable;
+    }
+    implementable &= CollectiveMainloop::can_implement(args.problem_shape, args.mainloop);
+    implementable &= CollectiveEpilogue::can_implement(args.problem_shape, args.epilogue);
+    implementable &= TileScheduler::can_implement(args.scheduler);
+    return implementable;
+  }
+
+  static int
+  get_workspace_size(Arguments const& args) {
+    return 0;
+  }
+
+  static mutlass::Status
+  initialize_workspace(Arguments const& args, void* workspace = nullptr, musaStream_t stream = nullptr,
+      MusaHostAdapter* musa_adapter = nullptr) {
+    return Status::kSuccess;
+  }
+
+  // Computes the kernel launch grid shape based on runtime parameters
+  static dim3
+  get_grid_shape(Params const& params) {
+    TileSchedulerArguments args{};
+    args.swizzle_size = params.scheduler.swizzle_size_;
+    args.raster_order = params.scheduler.raster_order_ == TileScheduler::RasterOrder::AlongN ? TileScheduler::RasterOrderOptions::AlongN : TileScheduler::RasterOrderOptions::AlongM;
+    return TileScheduler::get_grid_shape(params.scheduler, params.problem_shape, TileShape{}, params.hw_info, args);
+  }
+
+  static dim3
+  get_block_shape() {
+    return dim3(MaxThreadsPerBlock, 1, 1);
+  }
+
+  MUTLASS_DEVICE
+  void
+  operator()(Params const& params, char* smem_buf) {
+    using namespace mute;
+    using X = Underscore;
+
+    // Preconditions
+    static_assert(rank(StrideA{}) == 3, "StrideA must be rank-3: [M, K, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(rank(StrideB{}) == 3, "StrideB must be rank-3: [N, K, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(rank(StrideC{}) == 3, "StrideC must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
+    static_assert(rank(StrideD{}) == 3, "StrideD must be rank-3: [M, N, L]. If batch mode is not needed, set L stride to Int<0>.");
+
+    int thread_idx = int(threadIdx.x);
+    int warp_idx   = canonical_warp_idx();
+
+    arch::allocate_async_barriers(NumBarriers);
+
+    // Separate out problem shape for convenience
+    // Optionally append 1s until problem shape is rank-4 in case its is only rank-3 (MNK)
+    auto problem_shape_MNKL = append<4>(params.problem_shape, Int<1>{});
+    auto M = get<0>(problem_shape_MNKL);
+    auto N = get<1>(problem_shape_MNKL);
+    auto K = get<2>(problem_shape_MNKL);
+    auto L = get<3>(problem_shape_MNKL);
+
+    //// TME requires special handling of strides to deal with coord codomain mapping
+    //// Represent the full tensors -- get these from TMA
+    //Tensor mA_mkl = params.mainloop.tme_load_a.get_tme_tensor(make_shape(M,K,L));                            // (m,k,l)
+    //Tensor mB_nkl = params.mainloop.tme_load_b.get_tme_tensor(make_shape(N,K,L));                            // (n,k,l)
+
+    auto blk_shape = TileShape{};                                                                // (BLK_M,BLK_N,BLK_K)
+    //auto blk_coord = make_coord(_,_,_);                                                   // (m,n,k) -- defer the slice
+
+    //// Make tiled views
+    //Tensor gA_mkl = local_tile(mA_mkl, blk_shape, blk_coord, Step<_1, X,_1>{});                  // (BLK_M,BLK_K,m,k,l)
+    //Tensor gB_nkl = local_tile(mB_nkl, blk_shape, blk_coord, Step< X,_1,_1>{});                  // (BLK_N,BLK_K,n,k,l)
+
+    // Perform the collective scoped MMA
+    CollectiveMainloop collective_mma;
+
+    auto load_inputs = collective_mma.load_init(problem_shape_MNKL, params.mainloop);
+
+    static_assert(mute::tuple_size_v<decltype(load_inputs)> >= 2, "Output of load_init must have at least two elements (A, B)");
+
+    // Extract out partitioned A and B
+    Tensor gA_mkl = get<0>(load_inputs);
+    Tensor gB_nkl = get<1>(load_inputs);
+
+    TileScheduler scheduler{params.scheduler};
+    auto work_tile_info = scheduler.initial_work_tile_info();
+    auto m_coord = idx2crd(int(work_tile_info.M_idx), shape<2>(gA_mkl));
+    auto n_coord = idx2crd(int(work_tile_info.N_idx), shape<2>(gB_nkl));
+    auto l_coord = idx2crd(int(work_tile_info.L_idx), shape<4>(gB_nkl));
+    auto output_tile_coord = make_coord(m_coord, n_coord, _, l_coord);
+
+    // Slice with m_coord and n_coord
+    Tensor gA = gA_mkl(_,_,m_coord,_,l_coord);                                                       // (BLK_M,BLK_K,k)
+    Tensor gB = gB_nkl(_,_,n_coord,_,l_coord);                                                       // (BLK_N,BLK_K,k)
+
+    // Allocate the tiled_mma and the accumulators for the (M,N) blk_shape
+    TiledMma tiled_mma;
+    Tensor accumulators = partition_fragment_C(tiled_mma, take<0,2>(blk_shape));                   // (MMA,MMA_M,MMA_N)
+
+    auto k_tile_iter  = mute::make_coord_iterator(shape<2>(gA));
+    auto k_tile_count = size<2>(gA);
+
+    collective_mma(
+      load_inputs,
+      output_tile_coord,
+      accumulators,
+      k_tile_iter, k_tile_count,
+      thread_idx,
+      smem_buf,
+      params.mainloop
+    );
+
+    constexpr int BLK_M_RANK = rank<0>(blk_shape);
+    bool m_oob = int(work_tile_info.M_idx) >= size<2>(gA_mkl);
+    auto m_max_coord = unwrap(mute::transform(make_seq<BLK_M_RANK>{}, [&](auto i) {
+        return  m_oob ? 0 : get<i>(M) - get<0,i>(blk_shape) * get<i>(m_coord);
+      }));
+
+    constexpr int BLK_N_RANK = rank<1>(blk_shape);
+    bool n_oob = int(work_tile_info.N_idx) >= size<2>(gB_nkl);
+    auto n_max_coord = unwrap(mute::transform(make_seq<BLK_N_RANK>{}, [&](auto i) {
+        return  n_oob ? 0 : get<i>(N) - get<1,i>(blk_shape) * get<i>(n_coord);
+      }));
+    auto residue_mnk = mute::make_tuple(m_max_coord, n_max_coord, Int<0>{});
+
+    // Epilogue and write to gD
+    CollectiveEpilogue epilogue{params.epilogue};
+    epilogue(
+      problem_shape_MNKL,
+      blk_shape,
+      output_tile_coord,
+      accumulators,
+      tiled_mma,
+      residue_mnk,
+      thread_idx,
+      smem_buf
+    );
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::kernel
+
+///////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/static_tile_scheduler.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/static_tile_scheduler.hpp
new file mode 100644
index 000000000..bd162734a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/static_tile_scheduler.hpp
@@ -0,0 +1,189 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/fast_math.h"
+#include "mutlass/gemm_coord.hpp"
+#include "mutlass/kernel_hardware_info.hpp"
+#include "mutlass/gemm/kernel/tile_scheduler_params.hpp"
+#include "mute/layout.hpp"
+#include "mute/tensor.hpp"
+#include "mutlass/pipeline/pipeline.hpp"
+
+namespace mutlass::gemm::kernel::detail {
+
+using namespace mute;
+///////////////////////////////////////////////////////////////////////////////
+
+class StaticTileScheduler {
+
+public:
+  struct WorkTileInfo {
+    int32_t M_idx = 0;
+    int32_t N_idx = 0;
+    int32_t L_idx = 0;
+  };
+
+  using Params = TileSchedulerDefaultParams;
+  using RasterOrder = typename Params::RasterOrder;
+  using RasterOrderOptions = typename Params::RasterOrderOptions;
+
+public:
+  struct Arguments {
+    int swizzle_size = 1;
+    RasterOrderOptions raster_order = RasterOrderOptions::AlongN;
+  };
+
+  template <class ProblemShapeMNKL, class TileShape>
+  static Params
+  to_underlying_arguments(
+      ProblemShapeMNKL problem_shape_mnkl,
+      TileShape tile_shape,
+      [[maybe_unused]] KernelHardwareInfo const& hw_info,
+      Arguments const& arguments,
+      [[maybe_unused]] void* workspace=nullptr,
+      [[maybe_unused]] const uint32_t epilogue_subtile = 1,
+      [[maybe_unused]] uint32_t ktile_start_alignment_count = 1u) {
+
+    static_assert(mute::is_static<TileShape>::value);
+
+    dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, tile_shape);
+
+    Params params;
+    params.initialize(
+      problem_blocks,
+      hw_info,
+      arguments.swizzle_size,
+      arguments.raster_order
+    );
+    return params;
+  }
+
+  MUTLASS_HOST_DEVICE
+  static bool
+  can_implement(Arguments const& args) {
+    return args.swizzle_size >= 1;
+  }
+
+  MUTLASS_HOST_DEVICE
+  StaticTileScheduler() { }
+
+  MUTLASS_DEVICE explicit StaticTileScheduler(Params const& params_) : scheduler_params(params_) {}
+
+  // Returns the initial work tile info that will be computed over
+  MUTLASS_DEVICE
+  WorkTileInfo
+  initial_work_tile_info() {
+    return get_current_work();
+  }
+
+  MUTLASS_DEVICE
+  WorkTileInfo
+  get_current_work() const {
+    auto problem_blocks = scheduler_params.problem_blocks_;
+    int macro_tile_x = scheduler_params.swizzle_size_;
+    int macro_tile_size, res_tile_x, res_tile_x_quo, macro_bid_max, max_tile_y, mini_x, blockidx_xy, macro_x, macro_res;
+
+    if (scheduler_params.raster_order_ == RasterOrder::AlongN) {
+      scheduler_params.divmod_macro_tile_x_.fast_divmod(res_tile_x_quo, res_tile_x, problem_blocks.y);
+      macro_bid_max = res_tile_x_quo * macro_tile_x * problem_blocks.x;
+      max_tile_y = problem_blocks.x - 1;
+    } else {
+      scheduler_params.divmod_macro_tile_x_.fast_divmod(res_tile_x_quo, res_tile_x, problem_blocks.x);
+      macro_bid_max = res_tile_x_quo * macro_tile_x * problem_blocks.y;
+      max_tile_y = problem_blocks.y - 1;
+    }
+    auto [block_idx_x, block_idx_y, block_idx_z] = static_cast<uint3>(blockIdx);
+    blockidx_xy = block_idx_x;
+    scheduler_params.divmod_macro_tile_size_.fast_divmod(macro_x, macro_res, blockidx_xy);
+
+    if (blockidx_xy >= macro_bid_max) {
+      block_idx_y = macro_res / res_tile_x;
+      mini_x = macro_res % res_tile_x;
+      if (block_idx_y % 2 == 1) mini_x = res_tile_x - 1 - mini_x;
+    } else {
+      block_idx_y = macro_res / macro_tile_x;
+      mini_x = macro_res % macro_tile_x;
+      if (block_idx_y % 2 == 1) {
+        mini_x = macro_tile_x - 1 - mini_x;
+      }
+    }
+    if (macro_x % 2 == 1) {
+      block_idx_y = max_tile_y - block_idx_y;
+    }
+    block_idx_x = macro_x * macro_tile_x + mini_x;
+    if (scheduler_params.raster_order_ == RasterOrder::AlongN) {
+      return {static_cast<int32_t>(block_idx_y), static_cast<int32_t>(block_idx_x), static_cast<int32_t>(block_idx_z)};
+    } else {
+      return {static_cast<int32_t>(block_idx_x), static_cast<int32_t>(block_idx_y), static_cast<int32_t>(block_idx_z)};
+    }
+  }
+
+  // Given the inputs, computes the total number of output blocks over which this problem will compute.
+  // Note that this is only the logical size of our grid, not the physical grid we will actually launch.
+  template<class ProblemShapeMNKL, class BlockShape>
+  MUTLASS_HOST_DEVICE static
+  dim3
+  get_tiled_cta_shape_mnl(ProblemShapeMNKL problem_shape_mnkl, BlockShape cta_shape) {
+    auto cta_m = mute::size(mute::ceil_div(mute::shape<0>(problem_shape_mnkl), mute::shape<0>(cta_shape)));
+    auto cta_n = mute::size(mute::ceil_div(mute::shape<1>(problem_shape_mnkl), mute::shape<1>(cta_shape)));
+
+    return Params::get_tiled_cta_shape_mnl(
+      to_gemm_coord(problem_shape_mnkl),
+      cta_m, cta_n
+    );
+  }
+
+  // Given the inputs, computes the physical grid we should launch.
+  template<class ProblemShapeMNKL, class BlockShape>
+  MUTLASS_HOST_DEVICE static
+  dim3
+  get_grid_shape(
+      [[maybe_unused]] Params const& params,
+      ProblemShapeMNKL problem_shape_mnk,
+      BlockShape cta_shape,
+      KernelHardwareInfo hw_info,
+      Arguments arguments = Arguments{1, RasterOrderOptions::AlongN},
+      bool truncate_by_problem_size=true) {
+
+    auto problem_shape_mnkl = mute::append<4>(problem_shape_mnk, mute::Int<1>{});
+    dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape_mnkl, cta_shape);
+
+    return Params::get_grid_shape(problem_blocks, hw_info, arguments.swizzle_size, arguments.raster_order);
+  }
+
+public:
+  // Sink scheduler params as a member
+  Params scheduler_params;
+};
+
+} // namespace mutlass::gemm::kernel::detail
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/tile_scheduler.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/tile_scheduler.hpp
new file mode 100644
index 000000000..52e5b8d9d
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/tile_scheduler.hpp
@@ -0,0 +1,124 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+/*! \file
+    \brief Utilities for selecting default tile schedulers
+*/
+
+#include "mutlass/detail/dependent_false.hpp"
+#include "mutlass/gemm/kernel/static_tile_scheduler.hpp"
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm {
+
+////////////////////////////////////////////////////////////////////////////////
+
+//
+// Tags for specifying tile schedulers
+//
+
+struct PersistentScheduler { };
+
+struct StreamKScheduler { };
+
+struct GroupScheduler { }; // Only used for Grouped GEMMs
+
+struct DefaultScheduler { };
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm
+
+////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::kernel::detail {
+
+//
+// Selectors mapping tile scheduler tag and arch tag to a tile scheduler class
+//
+
+template <
+  class TileSchedulerTag,
+  class ArchTag,
+  class TileShape,
+  class ClusterShape
+  , class ProblemShapeType = void
+>
+struct TileSchedulerSelector {
+  static_assert(mutlass::detail::dependent_false<ArchTag>,
+      "Could not select a tile scheduler for given parameters.");
+};
+
+template <
+  class TileShape,
+  class ClusterShape,
+  class ProblemShapeType
+>
+struct TileSchedulerSelector <
+  void,
+  arch::Mp31,
+  TileShape,
+  ClusterShape,
+  ProblemShapeType
+> {
+  using Scheduler = typename TileSchedulerSelector<
+    DefaultScheduler,
+    arch::Mp31,
+    TileShape,
+    ClusterShape,
+    ProblemShapeType
+  >::Scheduler;
+};
+
+template <
+  class TileShape,
+  class ClusterShape,
+  class ProblemShapeType
+>
+struct TileSchedulerSelector <
+  DefaultScheduler,
+  arch::Mp31,
+  TileShape,
+  ClusterShape,
+  ProblemShapeType
+> {
+  using Scheduler = StaticTileScheduler;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass::gemm::kernel::detail
+
+////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/tile_scheduler_params.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/tile_scheduler_params.hpp
new file mode 100644
index 000000000..5c8a7c9ed
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/kernel/tile_scheduler_params.hpp
@@ -0,0 +1,192 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+/*! \file
+    \brief Parameters structures for persistent tile schedulers
+*/
+
+#include "mutlass/coord.h"
+#include "mutlass/kernel_hardware_info.h"
+#include "mutlass/workspace.h"
+#include "mutlass/platform/platform.h"
+#include "mutlass/fast_math.h"
+#include "mutlass/gemm_coord.h"
+////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::gemm::kernel::detail {
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct TileSchedulerDefaultParams {
+
+  enum class RasterOrder {
+    AlongM,
+    AlongN
+  };
+
+  enum class RasterOrderOptions {
+    Heuristic,
+    AlongM,
+    AlongN
+  };
+
+
+  FastDivmod divmod_macro_tile_x_;
+  FastDivmod divmod_macro_tile_size_;
+  int32_t swizzle_size_ = 0;
+  RasterOrder raster_order_ = RasterOrder::AlongN;
+  dim3 problem_blocks_{0, 0, 0};
+
+  // Initializes members. This variant of the method should only be used when
+  // problem_shape and tile_shape contain modes of only rank 1.
+  void
+  initialize(
+    BatchedGemmCoord problem_shape,
+    GemmCoord tile_shape,
+    KernelHardwareInfo const& hw_info,
+    int swizzle_size,
+    RasterOrderOptions raster_order_option
+  ) {
+    dim3 problem_blocks = get_tiled_cta_shape_mnl(problem_shape, tile_shape);
+    return initialize(
+      problem_blocks,
+      hw_info,
+      swizzle_size,
+      raster_order_option
+    );
+  }
+
+  // Version of initialize that takes in as input the number of CTAs in the M and N and L dimensions.
+  // This is useful for calculating the tiled shape when a mode of problem and/or CTA shape has rank > 1,
+  // for which using MuTe algebra for calculating tile shapes is easiest.
+  void
+  initialize(
+    dim3 problem_blocks,
+    KernelHardwareInfo const& hw_info,
+    int swizzle_size,
+    RasterOrderOptions raster_order_option
+  ) {
+    MUTLASS_UNUSED(hw_info);
+    RasterOrder raster_order = get_rasterization_order(
+      problem_blocks.y,
+      problem_blocks.x,
+      raster_order_option
+    );
+
+    swizzle_size_ = swizzle_size;
+    raster_order_ = raster_order;
+    problem_blocks_ = problem_blocks;
+    if (raster_order_ == RasterOrder::AlongN) {
+      divmod_macro_tile_size_ = FastDivmod(swizzle_size_ * problem_blocks.x);
+    } else {
+      divmod_macro_tile_size_ = FastDivmod(swizzle_size_ * problem_blocks.y);
+    }
+    divmod_macro_tile_x_ = FastDivmod(swizzle_size_);
+  }
+
+  // Version of get_grid_shape that takes in as input the number of CTAs in the M and N and L dimensions.
+  // This is useful for calculating the tiled shape when a mode of problem and/or CTA shape has rank > 1,
+  // for which using MuTe algebra for calculating tile shapes is easiest.
+  MUTLASS_HOST_DEVICE static
+  dim3
+  get_grid_shape(
+    dim3 problem_blocks,
+    KernelHardwareInfo const& hw_info,
+    int swizzle_size,
+    RasterOrderOptions raster_order_option
+  ) {
+
+    MUTLASS_UNUSED(hw_info);
+    RasterOrder raster_order = get_rasterization_order(
+      problem_blocks.x,
+      problem_blocks.y,
+      raster_order_option
+    );
+
+    return dim3(problem_blocks.x * problem_blocks.y, 1, problem_blocks.z);
+  }
+
+  MUTLASS_HOST_DEVICE
+  static RasterOrder
+  get_rasterization_order(
+    uint32_t tiles_m,
+    uint32_t tiles_n,
+    RasterOrderOptions raster_order_option
+  ) {
+
+    if (raster_order_option == RasterOrderOptions::Heuristic) {
+      if (tiles_n > tiles_m) {
+        return RasterOrder::AlongM;
+      }
+      else {
+        return RasterOrder::AlongN;
+      }
+    }
+    else if (raster_order_option == RasterOrderOptions::AlongN) {
+        return RasterOrder::AlongN;
+    }
+    else {
+      return RasterOrder::AlongM;
+    }
+  }
+
+  // Get the number of CTA tiles in this problem. This variant of the method should only be used when
+  // problem_shape and tile_shape contain modes of only rank 1.
+  MUTLASS_HOST_DEVICE
+  static dim3
+  get_tiled_cta_shape_mnl(BatchedGemmCoord problem_shape, GemmCoord cta_shape) {
+    auto cta_m = (problem_shape.m() + cta_shape.m() - 1) / cta_shape.m();
+    auto cta_n = (problem_shape.n() + cta_shape.n() - 1) / cta_shape.n();
+
+    return get_tiled_cta_shape_mnl(problem_shape, cta_m, cta_n);
+  }
+
+  // Version of get_tiled_cta_shape_mnl that takes in as input the number of CTAs in the M and N dimensions.
+  // This is useful for calculating the tiled shape when a mode of problem and/or CTA shape has rank > 1,
+  // for which using MuTe algebra for calculating tile shapes is easiest.
+  MUTLASS_HOST_DEVICE
+  static dim3
+  get_tiled_cta_shape_mnl(BatchedGemmCoord problem_shape, uint32_t cta_m, uint32_t cta_n) {
+    return {
+      static_cast<uint32_t>(cta_m),
+      static_cast<uint32_t>(cta_n),
+      static_cast<uint32_t>(problem_shape.batch())
+    };
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////
+} // namespace mutlass::gemm::kernel::detail
+
+////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm/threadblock/threadblock_swizzle.h b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/threadblock/threadblock_swizzle.h
new file mode 100644
index 000000000..e787a057f
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm/threadblock/threadblock_swizzle.h
@@ -0,0 +1,55 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Implements several possible threadblock-swizzling functions mapping blockIdx to
+      GEMM problems.
+*/
+
+#pragma once
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace gemm {
+namespace threadblock {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Threadblock swizzling function for GEMMs
+template <int N = 1>
+struct GemmIdentityThreadblockSwizzle { };
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace threadblock
+} // namespace gemm
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm_coord.h b/vllm_musa/csrc/flash_attn/include/mutlass/gemm_coord.h
new file mode 100644
index 000000000..63d291077
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm_coord.h
@@ -0,0 +1,395 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mutlass/coord.h"
+
+namespace mutlass {
+namespace gemm {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Shape of a matrix multiply-add operation
+template <
+  /// Rows of matrix product
+  int M = 1,
+  /// Columns of matrix product
+  int N = 1,
+  /// Inner dimension of matrix product
+  int K = 1
+>
+struct GemmShape {
+  static int const kM = M;
+  static int const kN = N;
+  static int const kK = K;
+
+  static int const kMN = M * N;
+  static int const kMK = M * K;
+  static int const kKN = N * K;
+  static int const kMNK = M * N * K;
+
+  static int const kCount = kMNK;
+
+  //
+  // Static member functions
+  //
+
+  /// Returns a Coord object
+  MUTLASS_HOST_DEVICE
+  static Coord<3> toCoord() {
+    return make_Coord(kM, kN, kK);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Type alias of the transpose of a GemmShape
+template <
+  /// concept: GemmShape
+  typename Shape
+>
+using GemmShapeTranspose = GemmShape<Shape::kN, Shape::kM, Shape::kK>;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// GemmCoord is a structure derived from Coord<3> that specifies a location within the
+/// coordinate space of a GEMM problem.
+struct GemmCoord : public Coord<3, int> {
+
+  /// Integer-valued index
+  typedef int Index;
+
+  /// Base type is a Coord of rank=3
+  typedef Coord<3, Index> Base;
+
+  /// GEMM M dimension - rows of the output C matrix
+  static int const kM = 0;
+
+  /// GEMM N dimension - columns of the output C matrix
+  static int const kN = 1;
+
+  /// GEMM K dimension - inner dimension of the GEMM problem
+  static int const kK = 2;
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  GemmCoord() { }
+
+  /// Constructs from Coord<3> and a batch
+  MUTLASS_HOST_DEVICE
+  GemmCoord(Coord<3, Index> const& coord): Base(make_Coord(coord[0], coord[1], coord[2])) { }
+
+  /// Helper to construct from a K, N, M, batch variables
+  MUTLASS_HOST_DEVICE
+  GemmCoord(Index m, Index n, Index k): Base(make_Coord(m, n, k)) { }
+
+  /// Returns the GEMM M coordinate
+  MUTLASS_HOST_DEVICE
+  Index const&  m() const { return this->at(kM); }
+
+  /// Returns reference to the GEMM M coordinate
+  MUTLASS_HOST_DEVICE
+  Index & m() { return this->at(kM); }
+
+  /// Returns the GEMM N coordinate
+  MUTLASS_HOST_DEVICE
+  Index const&  n() const { return this->at(kN); }
+
+  /// Returns reference to the GEMM N coordinate
+  MUTLASS_HOST_DEVICE
+  Index & n() { return this->at(kN); }
+
+  /// Returns the GEMM K coordinate
+  MUTLASS_HOST_DEVICE
+  Index const&  k() const { return this->at(kK); }
+
+  /// Returns reference to the GEMM K coordinate
+  MUTLASS_HOST_DEVICE
+  Index & k() { return this->at(kK); }
+
+  /// Obtains a Coord<3> from GemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<3> mnk() const {
+    return make_Coord(m(), n(), k());
+  }
+
+  /// Obtains a Coord<3> from GemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<3> knm() const {
+    return make_Coord(k(), n(), m());
+  }
+
+  /// Obtains a Coord<2> from GemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<2> nm() const {
+    return make_Coord(n(), m());
+  }
+
+  /// Obtains a Coord<2> from GemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<2> mn() const {
+    return make_Coord(m(), n());
+  }
+
+  /// Obtains a Coord<2> from GemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<2> mk() const {
+    return make_Coord(m(), k());
+  }
+
+  /// Obtains a Coord<2> from GemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<2> km() const {
+    return make_Coord(k(), m());
+  }
+
+  /// Obtains a Coord<2> from GemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<2> nk() const {
+    return make_Coord(n(), k());
+  }
+
+  /// Obtains a Coord<2> from GemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<2> kn() const {
+    return make_Coord(k(), n());
+  }
+
+  //
+  // Coord operators
+  //
+
+  /// Element-wise addition
+  MUTLASS_HOST_DEVICE
+  GemmCoord operator+(Base const& b) const {
+    return GemmCoord(Base::operator+(b));
+  }
+
+  /// Element-wise subtraction
+  MUTLASS_HOST_DEVICE
+  GemmCoord operator-(Base const& b) const {
+    return GemmCoord(Base::operator-(b));
+  }
+
+  /// Element-wise multiplication
+  MUTLASS_HOST_DEVICE
+  GemmCoord operator*(Base const& b) const {
+    return GemmCoord(Base::operator*(b));
+  }
+
+  /// Element-wise division
+  MUTLASS_HOST_DEVICE
+  GemmCoord operator/(Base const& b) const {
+    return GemmCoord(Base::operator/(b));
+  }
+
+  /// In-place addition
+  MUTLASS_HOST_DEVICE
+  GemmCoord& operator+=(Base const& b) {
+    Base::operator+=(b);
+    return *this;
+  }
+
+  /// In-place subtraction
+  MUTLASS_HOST_DEVICE
+  GemmCoord& operator-=(Base const& b) {
+    Base::operator-=(b);
+    return *this;
+  }
+
+  /// In-place multiplication
+  MUTLASS_HOST_DEVICE
+  GemmCoord& operator*=(Base const& b) {
+    Base::operator*=(b);
+    return *this;
+  }
+
+  /// In-place division
+  MUTLASS_HOST_DEVICE
+  GemmCoord& operator/=(Base const& b) {
+    Base::operator/=(b);
+    return *this;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// BatchedGemmCoord is a structure derived from Coord<4> that specifies a location within the
+/// coordinate space of a batched GEMM problem.
+struct BatchedGemmCoord : public Coord<4, int> {
+
+  /// Integer-valued index
+  typedef int Index;
+
+  /// Base type is a Coord of rank=4
+  typedef Coord<4, Index> Base;
+
+  /// GEMM M dimension - rows of the output C matrix
+  static int const kM = 0;
+
+  /// GEMM N dimension - columns of the output C matrix
+  static int const kN = 1;
+
+  /// GEMM K dimension - inner dimension of the GEMM problem
+  static int const kK = 2;
+
+  /// GEMM Batch dimension - inner dimension of the GEMM problem
+  static int const kBatch = 3;
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord() { }
+
+  /// Constructs from Coord<4>
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord(Base const& coord): Base(coord) { }
+
+  /// Helper to construct from a K, N, M, and batch variables
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord(Index m, Index n, Index k, Index b): Base(make_Coord(m, n, k, b)) { }
+
+  /// Returns the GEMM M coordinate
+  MUTLASS_HOST_DEVICE
+  Index const&  m() const { return this->at(kM); }
+
+  /// Returns reference to the GEMM M coordinate
+  MUTLASS_HOST_DEVICE
+  Index & m() { return this->at(kM); }
+
+  /// Returns the GEMM N coordinate
+  MUTLASS_HOST_DEVICE
+  Index const&  n() const { return this->at(kN); }
+
+  /// Returns reference to the GEMM N coordinate
+  MUTLASS_HOST_DEVICE
+  Index & n() { return this->at(kN); }
+
+  /// Returns the GEMM K coordinate
+  MUTLASS_HOST_DEVICE
+  Index const&  k() const { return this->at(kK); }
+
+  /// Returns reference to the GEMM K coordinate
+  MUTLASS_HOST_DEVICE
+  Index & k() { return this->at(kK); }
+
+  /// Returns the GEMM batch coordinate
+  MUTLASS_HOST_DEVICE
+  Index const&  batch() const { return this->at(kBatch); }
+
+  /// Returns reference to the GEMM batch coordinate
+  MUTLASS_HOST_DEVICE
+  Index & batch() { return this->at(kBatch); }
+
+  /// Obtains a GemmCoord from BatchedGemmCoord
+  MUTLASS_HOST_DEVICE
+  GemmCoord mnk() const {
+    return GemmCoord(m(), n(), k());
+  }
+
+  /// Obtains a Coord<4> from BatchedGemmCoord
+  MUTLASS_HOST_DEVICE
+  Coord<4> mnkb() const {
+    return make_Coord(m(), n(), k(), batch());
+  }
+
+  //
+  // Coord operators
+  //
+
+  /// Element-wise addition
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord operator+(Base const& b) const {
+    return BatchedGemmCoord(Base::operator+(b));
+  }
+
+  /// Element-wise subtraction
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord operator-(Base const& b) const {
+    return BatchedGemmCoord(Base::operator-(b));
+  }
+
+  /// Element-wise multiplication
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord operator*(Base const& b) const {
+    return BatchedGemmCoord(Base::operator*(b));
+  }
+
+  /// Element-wise division
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord operator/(Base const& b) const {
+    return BatchedGemmCoord(Base::operator/(b));
+  }
+
+  /// In-place addition
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord& operator+=(Base const& b) {
+    Base::operator+=(b);
+    return *this;
+  }
+
+  /// In-place subtraction
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord& operator-=(Base const& b) {
+    Base::operator-=(b);
+    return *this;
+  }
+
+  /// In-place multiplication
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord& operator*=(Base const& b) {
+    Base::operator*=(b);
+    return *this;
+  }
+
+  /// In-place division
+  MUTLASS_HOST_DEVICE
+  BatchedGemmCoord& operator/=(Base const& b) {
+    Base::operator/=(b);
+    return *this;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace gemm
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/gemm_coord.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/gemm_coord.hpp
new file mode 100644
index 000000000..4924eb7cc
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/gemm_coord.hpp
@@ -0,0 +1,67 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Utilities to convert a MuTe tuple to a GemmCoord or BatchedGemmCoord
+*/
+
+#pragma once
+
+#include "mute/layout.hpp"
+#include "mutlass/gemm_coord.h"
+
+namespace mutlass {
+namespace gemm {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class Tuple>
+MUTLASS_HOST_DEVICE
+auto
+to_gemm_coord(Tuple tuple) {
+  static_assert(mute::rank(tuple) <= 4, "Can only convert tuples of rank <= 4.");
+
+  if constexpr (mute::rank(tuple) <= 3) {
+    auto tuple_mnk = mute::append<3>(tuple, mute::Int<0>{});
+    return GemmCoord(mute::size<0>(tuple_mnk), mute::size<1>(tuple_mnk), mute::size<2>(tuple_mnk));
+  }
+  else {
+    return BatchedGemmCoord(mute::size<0>(tuple), mute::size<1>(tuple), mute::size<2>(tuple), mute::size<3>(tuple));
+  }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace gemm
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/half.h b/vllm_musa/csrc/flash_attn/include/mutlass/half.h
new file mode 100644
index 000000000..af885e6c1
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/half.h
@@ -0,0 +1,905 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a class for using IEEE half-precision floating-point types in host or
+      device code.
+*/
+
+#pragma once
+
+#ifndef MUTLASS_ENABLE_F16C
+#define MUTLASS_ENABLE_F16C 0
+#endif
+
+#if defined(__MUSACC_RTC__)
+
+#include "mutlass/floating_point_mtrtc.h"
+
+#undef MUTLASS_ENABLE_F16C
+#define MUTLASS_ENABLE_F16C 0
+
+#else
+#include <cmath>
+#include <limits>
+#include <cstdint>
+#include <cstring>
+#endif
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include <musa_fp16.h>
+
+#include "mutlass/mutlass.h"
+#include "mutlass/float8.h"
+#include "mutlass/platform/platform.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Optionally target F16C extentions to accelerate half-precision conversion.
+#if !defined(__MUSA_ARCH__) && (MUTLASS_ENABLE_F16C)
+#if defined(_MSC_VER)
+
+#include <immintrin.h>
+
+#if defined(__i386__) || defined(__x86_64__)
+#include <intrin.h>
+#endif
+
+#define F16C_ROUND_NEAREST 0
+
+#if !defined(__MUSA_ARCH__)
+extern __inline float _cvtsh_ss (unsigned short __S) {
+  __m128i packed;
+  std::memcpy(&packed, &__S, sizeof(__S));
+
+  __m128 result = _mm_cvtph_ps(packed);
+
+  float flt;
+  std::memcpy(&flt, &result, sizeof(flt));
+
+  return flt;
+}
+
+__inline unsigned short _cvtss_sh (float __F, const int) {
+  __m128 packed;
+  std::memcpy(&packed, &__F, sizeof(__F));
+
+  __m128i result = _mm_cvtps_ph(packed, F16C_ROUND_NEAREST);
+
+  unsigned short u;
+  std::memcpy(&u, &result, sizeof(u));
+
+  return u;
+}
+#endif
+
+#else
+
+// Linux
+#include <x86intrin.h>
+
+#if defined(__i386__) || defined(__x86_64__)
+#include <cpuid.h>
+#endif
+
+#define F16C_ROUND_NEAREST (_MM_FROUND_TO_NEAREST_INT |_MM_FROUND_NO_EXC)
+
+#endif // _MSC_VER
+
+class CpuId {
+
+  bool f16c_enabled;
+
+  CpuId() {
+  #if defined(__i386__) || defined(__x86_64__)
+    #if defined(_MSC_VER)
+      int exx[4];
+
+      __cpuid (exx, 1);
+      f16c_enabled = exx[2] & 0x20000000;
+
+    #else
+    // GCC / Clang
+       int eax, ebx, ecx, edx;
+
+      __cpuid (1 , eax, ebx, ecx, edx);
+      f16c_enabled = ecx & 0x20000000;
+    #endif
+  #else
+  // Arm / PowerPC etc.
+    f16c_enabled = false;
+  #endif
+  }
+
+public:
+
+  bool is_f16c_supported() const {
+    return f16c_enabled;
+  }
+
+  static const CpuId& instance() {
+      static CpuId cpu;
+      return cpu;
+  }
+};
+#endif // !defined(__MUSA_ARCH__) && MUTLASS_ENABLE_F16C
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// IEEE half-precision floating-point type
+struct alignas(2) half_t {
+
+  //
+  // Data members
+  //
+
+  /// Storage type
+  uint16_t storage;
+
+  //
+  // Static conversion operators
+  //
+
+  /// Constructs from an unsigned short
+  MUTLASS_HOST_DEVICE
+  static half_t bitcast(uint16_t x) {
+    half_t h;
+    h.storage = x;
+    return h;
+  }
+
+  /// FP32 -> FP16 conversion - rounds to nearest even
+  MUTLASS_HOST_DEVICE
+  static half_t convert(float const& flt) {
+  #if defined(__MUSA_ARCH__)
+    return half_t(__float2half_rn(flt));
+  #else
+
+    #if !defined(__MUSA_ARCH__) && MUTLASS_ENABLE_F16C
+      if( CpuId::instance().is_f16c_supported() ) {
+        unsigned short u = _cvtss_sh(flt, F16C_ROUND_NEAREST);
+        return bitcast(u);
+      }
+    #endif
+
+    // software implementation rounds toward nearest even
+    unsigned s;
+
+    #if defined(__MUSA_ARCH__)
+    s = reinterpret_cast<unsigned const &>(flt);
+    #else
+    std::memcpy(&s, &flt, sizeof(s));
+    #endif
+
+    uint16_t sign = uint16_t((s >> 16) & 0x8000);
+    int16_t exp = uint16_t(((s >> 23) & 0xff) - 127);
+    int mantissa = s & 0x7fffff;
+    uint16_t u = 0;
+
+    if ((s & 0x7fffffff) == 0) {
+      // sign-preserving zero
+      return bitcast(sign);
+    }
+
+    if (exp > 15) {
+      if (exp == 128 && mantissa) {
+        // not a number
+        u = 0x7fff;
+      } else {
+        // overflow to infinity
+        u = sign | 0x7c00;
+      }
+      return bitcast(u);
+    }
+
+    int sticky_bit = 0;
+
+    if (exp >= -14) {
+      // normal fp32 to normal fp16
+      exp = uint16_t(exp + uint16_t(15));
+      u = uint16_t(((exp & 0x1f) << 10));
+      u = uint16_t(u | (mantissa >> 13));
+    } else {
+      // normal single-precision to subnormal half_t-precision representation
+      int rshift = (-14 - exp);
+      if (rshift < 32) {
+        mantissa |= (1 << 23);
+
+        sticky_bit = ((mantissa & ((1 << rshift) - 1)) != 0);
+
+        mantissa = (mantissa >> rshift);
+        u = (uint16_t(mantissa >> 13) & 0x3ff);
+      } else {
+        mantissa = 0;
+        u = 0;
+      }
+    }
+
+    // round to nearest even
+    int round_bit = ((mantissa >> 12) & 1);
+    sticky_bit |= ((mantissa & ((1 << 12) - 1)) != 0);
+
+    if ((round_bit && sticky_bit) || (round_bit && (u & 1))) {
+      u = uint16_t(u + 1);
+    }
+
+    u |= sign;
+
+    return bitcast(u);
+  #endif
+  }
+
+  /// FP32 -> FP16 conversion - rounds to nearest even
+  MUTLASS_HOST_DEVICE
+  static half_t convert(int const& n) {
+  #if defined(__MUSA_ARCH__)
+    return half_t(__int2half_rn(n));
+  #else
+    return convert(float(n));
+  #endif
+  }
+
+  /// FP32 -> FP16 conversion - rounds to nearest even
+  MUTLASS_HOST_DEVICE
+  static half_t convert(unsigned const& n) {
+  #if defined(__MUSA_ARCH__)
+    return half_t(__uint2half_rn(n));
+  #else
+    return convert(float(n));
+  #endif
+  }
+
+  /// Converts a half-precision value stored as a uint16_t to a float
+  MUTLASS_HOST_DEVICE
+  static float convert(half_t const& x) {
+  #if defined(__MUSA_ARCH__)
+    return __half2float(x.to_half());
+  #else
+
+    #if !defined(__MUSA_ARCH__) && MUTLASS_ENABLE_F16C
+      if( CpuId::instance().is_f16c_supported() ) {
+        unsigned short u = x.storage;
+        return _cvtsh_ss(u);
+      }
+    #endif
+
+    uint16_t const &h = x.storage;
+    uint32_t sign = ((h >> 15) & 1);
+    uint32_t exp = ((h >> 10) & 0x1f);
+    uint32_t mantissa = (h & 0x3ff);
+    unsigned f = 0;
+
+    if (exp > 0 && exp < 31) {
+      // normal
+      exp += 112;
+      f = (sign << 31) | (exp << 23) | (mantissa << 13);
+    } else if (exp == 0) {
+      if (mantissa) {
+        // subnormal
+        exp += 113;
+        while ((mantissa & (1 << 10)) == 0) {
+          mantissa <<= 1;
+          exp--;
+        }
+        mantissa &= 0x3ff;
+        f = (sign << 31) | (exp << 23) | (mantissa << 13);
+      } else {
+        // sign-preserving zero
+        f = (sign << 31);
+      }
+    } else if (exp == 31) {
+      if (mantissa) {
+        f = 0x7fffffff;  // not a number
+      } else {
+        f = (0xff << 23) | (sign << 31);  //  inf
+      }
+    }
+    #if defined(__MUSA_ARCH__)
+    return reinterpret_cast<float const&>(f);
+    #else
+    float flt;
+    std::memcpy(&flt, &f, sizeof(flt));
+    return flt;
+    #endif
+  #endif
+  }
+
+  //
+  // Methods
+  //
+
+  /// Default constructor
+  half_t() = default;
+
+  /// Reinterpret cast from MUSA's half type
+  MUTLASS_HOST_DEVICE
+  explicit half_t(half const & x) {
+    #if defined(__MUSA_ARCH__)
+    storage = reinterpret_cast<uint16_t const &>(x);
+    #else
+    __half_raw raw(x);
+    std::memcpy(&storage, &raw.x, sizeof(storage));
+    #endif
+  }
+
+  /// Floating point conversion
+  MUTLASS_HOST_DEVICE
+  explicit half_t(float x) {
+    storage = convert(x).storage;
+  }
+
+  /// Floating point conversion
+  MUTLASS_HOST_DEVICE
+  explicit half_t(double x): half_t(float(x)) {
+
+  }
+
+  /// float_e4m3_t conversion
+  MUTLASS_HOST_DEVICE
+  explicit half_t(float_e4m3_t x): half_t(float(x)) {
+
+  }
+
+  /// float_e5m2_t conversion
+  MUTLASS_HOST_DEVICE
+  explicit half_t(float_e5m2_t x): half_t(float(x)) {
+
+  }
+
+  /// Integer conversion - round to nearest even
+  MUTLASS_HOST_DEVICE
+  explicit half_t(int x) {
+    storage = convert(x).storage;
+  }
+
+  /// Integer conversion - round toward zero
+  MUTLASS_HOST_DEVICE
+  explicit half_t(unsigned x) {
+    storage = convert(x).storage;
+  }
+
+  /// Assignment
+  MUTLASS_HOST_DEVICE
+  half_t & operator=(half const &x) {
+    #if defined(__MUSA_ARCH__)
+    storage = reinterpret_cast<uint16_t const &>(x);
+    #else
+    __half_raw raw(x);
+    std::memcpy(&storage, &raw.x, sizeof(storage));
+    #endif
+    return *this;
+  }
+
+  /// Converts to float
+  MUTLASS_HOST_DEVICE
+  operator float() const {
+    return convert(*this);
+  }
+
+  /// Converts to float
+  MUTLASS_HOST_DEVICE
+  explicit operator double() const {
+    return double(convert(*this));
+  }
+
+  /// Converts to float
+  MUTLASS_HOST_DEVICE
+  explicit operator int() const {
+    return int(convert(*this));
+  }
+
+  /// Casts to bool
+  MUTLASS_HOST_DEVICE
+  explicit operator bool() const {
+    return (convert(*this) != 0.0f);
+  }
+
+  /// Bitcasts to MUSA's half type
+  MUTLASS_HOST_DEVICE
+  half to_half() const {
+    #if defined(__MUSA_ARCH__)
+    return reinterpret_cast<half const &>(storage);
+    #else
+    __half_raw raw;
+    std::memcpy(&raw.x, &storage, sizeof(raw.x));
+    return half(raw);
+    #endif
+  }
+
+  /// Accesses raw internal state
+  MUTLASS_HOST_DEVICE
+  uint16_t& raw() {
+    return storage;
+  }
+
+  /// Accesses raw internal state
+  MUTLASS_HOST_DEVICE
+  uint16_t raw() const {
+    return storage;
+  }
+
+  /// Returns the sign bit
+  MUTLASS_HOST_DEVICE
+  bool signbit() const {
+    return ((storage & 0x8000) != 0);
+  }
+
+  /// Returns the biased exponent
+  MUTLASS_HOST_DEVICE
+  int exponent_biased() const {
+    return int((storage >> 10) & 0x1f);
+  }
+
+  /// Returns the unbiased exponent
+  MUTLASS_HOST_DEVICE
+  int exponent() const {
+    return exponent_biased() - 15;
+  }
+
+  /// Returns the mantissa
+  MUTLASS_HOST_DEVICE
+  int mantissa() const {
+    return int(storage & 0x3ff);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+MUTLASS_HOST_DEVICE
+bool signbit(mutlass::half_t const& h) {
+  return ((h.raw() & 0x8000) != 0);
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::half_t abs(mutlass::half_t const& h) {
+  return mutlass::half_t::bitcast(h.raw() & 0x7fff);
+}
+
+MUTLASS_HOST_DEVICE
+bool isnan(mutlass::half_t const& h) {
+  return (h.exponent_biased() == 0x1f) && h.mantissa();
+}
+
+MUTLASS_HOST_DEVICE
+bool isfinite(mutlass::half_t const& h) {
+  return (h.exponent_biased() != 0x1f);
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::half_t nanh(const char*) {
+  return mutlass::half_t::bitcast(0x7fff);
+}
+
+MUTLASS_HOST_DEVICE
+bool isinf(mutlass::half_t const& h) {
+  return (h.exponent_biased() == 0x1f) && !h.mantissa();
+}
+
+MUTLASS_HOST_DEVICE
+bool isnormal(mutlass::half_t const& h) {
+  return h.exponent_biased() && h.exponent_biased() != 0x1f;
+}
+
+MUTLASS_HOST_DEVICE
+int fpclassify(mutlass::half_t const& h) {
+  int exp = h.exponent_biased();
+  int mantissa = h.mantissa();
+  if (exp == 0x1f) {
+    if (mantissa) {
+      return FP_NAN;
+    }
+    else {
+      return FP_INFINITE;
+    }
+  }
+  else if (!exp) {
+    if (mantissa) {
+      return FP_SUBNORMAL;
+    }
+    else {
+      return FP_ZERO;
+    }
+  }
+  return FP_NORMAL;
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::half_t sqrt(mutlass::half_t const& h) {
+#if defined(__MUSACC_RTC__)
+  return mutlass::half_t(sqrtf(float(h)));
+#else
+  return mutlass::half_t(std::sqrt(float(h)));
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t copysign(half_t const& a, half_t const& b) {
+
+  uint16_t a_mag = (a.raw() & 0x7fff);
+  uint16_t b_sign = (b.raw() & 0x8000);
+  uint16_t result = (a_mag | b_sign);
+
+  return half_t::bitcast(result);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Standard Library operations and definitions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if !defined(__MUSACC_RTC__)
+namespace std {
+
+/// Numeric limits
+template <>
+struct numeric_limits<mutlass::half_t> {
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_infinity = true;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+  static std::float_denorm_style const has_denorm = std::denorm_present;
+  static bool const has_denorm_loss = true;
+  static std::float_round_style const round_style = std::round_to_nearest;
+  static bool const is_iec559 = true;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = 10;
+
+  /// Least positive value
+  static mutlass::half_t min() { return mutlass::half_t::bitcast(0x0001); }
+
+  /// Minimum finite value
+  static mutlass::half_t lowest() { return mutlass::half_t::bitcast(0xfbff); }
+
+  /// Maximum finite value
+  static mutlass::half_t max() { return mutlass::half_t::bitcast(0x7bff); }
+
+  /// Returns smallest finite value
+  static mutlass::half_t epsilon() { return mutlass::half_t::bitcast(0x1800); }
+
+  /// Returns maximum rounding error
+  static mutlass::half_t round_error() { return mutlass::half_t(0.5f); }
+
+  /// Returns positive infinity value
+  static mutlass::half_t infinity() { return mutlass::half_t::bitcast(0x7c00); }
+
+  /// Returns quiet NaN value
+  static mutlass::half_t quiet_NaN() { return mutlass::half_t::bitcast(0x7fff); }
+
+  /// Returns signaling NaN value
+  static mutlass::half_t signaling_NaN() { return mutlass::half_t::bitcast(0x7fff); }
+
+  /// Returns smallest positive subnormal value
+  static mutlass::half_t denorm_min() { return mutlass::half_t::bitcast(0x0001); }
+};
+}  // namespace std
+#endif
+
+namespace platform {
+
+/// std::numeric_limits
+template <class T>
+struct numeric_limits;
+
+/// Numeric limits
+template <>
+struct numeric_limits<mutlass::half_t> {
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_infinity = true;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+#if !defined(__MUSACC_RTC__)
+  static std::float_denorm_style const has_denorm = std::denorm_present;
+#endif
+  static bool const has_denorm_loss = true;
+#if !defined(__MUSACC_RTC__)
+  static std::float_round_style const round_style = std::round_to_nearest;
+#endif
+  static bool const is_iec559 = true;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = 10;
+
+  /// Least positive value
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t min() { return mutlass::half_t::bitcast(0x0001); }
+
+  /// Minimum finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t lowest() { return mutlass::half_t::bitcast(0xfbff); }
+
+  /// Maximum finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t max() { return mutlass::half_t::bitcast(0x7bff); }
+
+  /// Returns smallest finite value
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t epsilon() { return mutlass::half_t::bitcast(0x1800); }
+
+  /// Returns maximum rounding error
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t round_error() { return mutlass::half_t(0.5f); }
+
+  /// Returns positive infinity value
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t infinity() { return mutlass::half_t::bitcast(0x7c00); }
+
+  /// Returns quiet NaN value
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t quiet_NaN() { return mutlass::half_t::bitcast(0x7fff); }
+
+  /// Returns signaling NaN value
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t signaling_NaN() { return mutlass::half_t::bitcast(0x7fff); }
+
+  /// Returns smallest positive subnormal value
+  MUTLASS_HOST_DEVICE
+  static mutlass::half_t denorm_min() { return mutlass::half_t::bitcast(0x0001); }
+};
+}  // namespace platform
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Arithmetic operators
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+MUTLASS_HOST_DEVICE
+bool operator==(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return __heq(lhs.to_half(), rhs.to_half());
+#else
+  return float(lhs) == float(rhs);
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+bool operator!=(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return __hne(lhs.to_half(), rhs.to_half());
+#else
+  return float(lhs) != float(rhs);
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return __hlt(lhs.to_half(), rhs.to_half());
+#else
+  return float(lhs) < float(rhs);
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<=(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return __hle(lhs.to_half(), rhs.to_half());
+#else
+  return float(lhs) <= float(rhs);
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return __hgt(lhs.to_half(), rhs.to_half());
+#else
+  return float(lhs) > float(rhs);
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>=(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return __hge(lhs.to_half(), rhs.to_half());
+#else
+  return float(lhs) >= float(rhs);
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t operator+(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return half_t(__hadd(lhs.to_half(), rhs.to_half()));
+#else
+  return half_t(float(lhs) + float(rhs));
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t operator-(half_t const& lhs) {
+#if defined(__MUSA_ARCH__)
+  return half_t(__hneg(lhs.to_half()));
+#else
+  return half_t(-float(lhs));
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t operator-(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return half_t(__hsub(lhs.to_half(), rhs.to_half()));
+#else
+  return half_t(float(lhs) - float(rhs));
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t operator*(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return half_t(__hmul(lhs.to_half(), rhs.to_half()));
+#else
+  return half_t(float(lhs) * float(rhs));
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t operator/(half_t const& lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  return half_t(__hdiv(lhs.to_half(), rhs.to_half()));
+#else
+  return half_t(float(lhs) / float(rhs));
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+half_t& operator+=(half_t & lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  lhs = half_t(__hadd(lhs.to_half(), rhs.to_half()));
+#else
+  lhs = half_t(float(lhs) + float(rhs));
+#endif
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+half_t& operator-=(half_t & lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  lhs = half_t(__hsub(lhs.to_half(), rhs.to_half()));
+#else
+  lhs = half_t(float(lhs) - float(rhs));
+#endif
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+half_t& operator*=(half_t & lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  lhs = half_t(__hmul(lhs.to_half(), rhs.to_half()));
+#else
+  lhs = half_t(float(lhs) * float(rhs));
+#endif
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+half_t& operator/=(half_t & lhs, half_t const& rhs) {
+#if defined(__MUSA_ARCH__)
+  lhs = half_t(__hdiv(lhs.to_half(), rhs.to_half()));
+#else
+  lhs = half_t(float(lhs) / float(rhs));
+#endif
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+half_t& operator++(half_t & lhs) {
+#if defined(__MUSA_ARCH__)
+  lhs = half_t(__hadd(lhs.to_half(), half_t(1.0f).to_half()));
+#else
+  float tmp(lhs);
+  ++tmp;
+  lhs = half_t(tmp);
+#endif
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+half_t& operator--(half_t & lhs) {
+#if defined(__MUSA_ARCH__)
+  lhs = half_t(__hsub(lhs.to_half(), half_t(1.0f).to_half()));
+#else
+  float tmp(lhs);
+  --tmp;
+  lhs = half_t(tmp);
+#endif
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+half_t operator++(half_t & lhs, int) {
+  half_t ret(lhs);
+#if defined(__MUSA_ARCH__)
+  lhs = half_t(__hadd(lhs.to_half(), half_t(1.0f).to_half()));
+#else
+  float tmp(lhs);
+  tmp++;
+  lhs = half_t(tmp);
+#endif
+  return ret;
+}
+
+MUTLASS_HOST_DEVICE
+half_t operator--(half_t & lhs, int) {
+  half_t ret(lhs);
+#if defined(__MUSA_ARCH__)
+  lhs = half_t(__hsub(lhs.to_half(), half_t(1.0f).to_half()));
+#else
+  float tmp(lhs);
+  tmp--;
+  lhs = half_t(tmp);
+#endif
+  return ret;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// User-defined literals
+//
+
+MUTLASS_HOST_DEVICE
+mutlass::half_t operator "" _hf(long double x) {
+  return mutlass::half_t(float(x));
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::half_t operator "" _hf(unsigned long long int x) {
+  return mutlass::half_t(int(x));
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/integer_subbyte.h b/vllm_musa/csrc/flash_attn/include/mutlass/integer_subbyte.h
new file mode 100644
index 000000000..f1d32b689
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/integer_subbyte.h
@@ -0,0 +1,268 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a class for using integer types smaller than one byte in host or
+      device code.
+*/
+
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <cstdint>
+#endif
+
+#include "mutlass/mutlass.h"
+#include "mutlass/numeric_size.h"
+#include "mutlass/platform/platform.h"
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int Bits, bool Signed = true>
+struct integer_subbyte {
+  /// Storage type
+  using Storage = uint8_t;
+
+  /// Number of bits
+  static_assert(Bits <= 8*sizeof(Storage), "Require a subbyte of bits in integer_subbyte");
+
+  /// External type
+  using xint_t = typename platform::conditional<Signed, int, unsigned>::type;
+
+  /// Bitmask for truncation from larger integers
+  static constexpr Storage bits_mask_ = Storage(Storage(-1) >> (8 - Bits));
+  /// Bitmask for the sign bit
+  static constexpr Storage sign_mask_ = Storage((Signed ? 1 : 0) << (Bits - 1));
+
+  //
+  // Data members
+  //
+
+  Storage storage;
+
+  //
+  // Methods
+  //
+
+  /// No operation
+  integer_subbyte() = default;
+
+  /// Conversion from integer type
+  MUTLASS_HOST_DEVICE explicit
+  integer_subbyte(int value)
+      : storage(reinterpret_cast<Storage const&>(value) & bits_mask_) {}
+
+  MUTLASS_HOST_DEVICE explicit
+  integer_subbyte(unsigned value)
+      : storage(reinterpret_cast<Storage const&>(value) & bits_mask_) {}
+
+  MUTLASS_HOST_DEVICE explicit
+  integer_subbyte(double value) {
+    xint_t tmp = static_cast<xint_t>(value);
+    storage = reinterpret_cast<Storage const &>(tmp) & bits_mask_;
+  }
+
+  /// Convert to int or unsigned
+  MUTLASS_HOST_DEVICE
+  operator xint_t() const {
+    if (sign_mask_ & storage) {  // Sign extend
+      return xint_t(storage) | ~xint_t(bits_mask_);
+    } else {
+      return xint_t(storage);
+    }
+  }
+
+  /// Equality
+  MUTLASS_HOST_DEVICE
+  bool operator==(integer_subbyte const& rhs) const {
+    return storage == rhs.storage;
+  }
+
+  /// Inequality
+  MUTLASS_HOST_DEVICE
+  bool operator!=(integer_subbyte const& rhs) const {
+    return storage != rhs.storage;
+  }
+
+  /// Less than or equal
+  MUTLASS_HOST_DEVICE
+  bool operator<=(integer_subbyte const& rhs) const {
+    if (sign_mask_ & storage) {
+      return !(rhs.storage < storage);
+    } else {
+      return storage <= rhs.storage;
+    }
+  }
+
+  /// Less than
+  MUTLASS_HOST_DEVICE
+  bool operator<(integer_subbyte const& rhs) const {
+    if (sign_mask_ & storage) {
+      return !(rhs.storage <= storage);
+    } else {
+      return storage < rhs.storage;
+    }
+  }
+
+  /// Greater than or equal
+  MUTLASS_HOST_DEVICE
+  bool operator>=(integer_subbyte const& rhs) const {
+    return !(*this < rhs);
+  }
+
+  /// Greater than
+  MUTLASS_HOST_DEVICE
+  bool operator>(integer_subbyte const& rhs) const {
+    return !(*this <= rhs);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 1-bit Unsigned integer type
+using uint1b_t = integer_subbyte<1, false>;
+
+/// 2-bit Integer type
+using int2b_t = integer_subbyte<2, true>;
+
+/// 2-bit Unsigned integer type
+using uint2b_t = integer_subbyte<2, false>;
+
+/// 4-bit Integer type
+using int4b_t = integer_subbyte<4, true>;
+
+/// 4-bit Unsigned integer type
+using uint4b_t = integer_subbyte<4, false>;
+
+/// 1-bit binary type
+using bin1_t = bool;
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int Bits, bool Signed>
+struct sizeof_bits<integer_subbyte<Bits,Signed>> {
+  static constexpr int value = Bits;
+};
+
+/// Defines the size of an element in bits - specialized for bin1_t
+template <>
+struct sizeof_bits<bin1_t> {
+  static constexpr int value = 1;
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace platform {
+
+template <>
+struct numeric_limits<mutlass::int4b_t> {
+  MUTLASS_HOST_DEVICE static
+  mutlass::int4b_t const lowest() noexcept { return int4b_t{-8};}
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::int4b_t const max() noexcept { return int4b_t{7};}
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::int4b_t const min() noexcept { return lowest();}
+
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = true;
+};
+
+template <>
+struct numeric_limits<mutlass::uint4b_t> {
+  MUTLASS_HOST_DEVICE static
+ mutlass::uint4b_t const lowest() noexcept { return uint4b_t{0};}
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::uint4b_t const max() noexcept { return uint4b_t{15};}
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::uint4b_t const min() noexcept { return lowest();}
+
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+template <>
+struct numeric_limits<mutlass::uint1b_t> {
+  MUTLASS_HOST_DEVICE static
+  mutlass::uint1b_t const lowest() noexcept { return uint1b_t{0};}
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::uint1b_t const max() noexcept { return uint1b_t{1};}
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::uint1b_t const min() noexcept { return lowest();}
+
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+template <>
+struct numeric_limits<mutlass::int2b_t> {
+  MUTLASS_HOST_DEVICE static
+  mutlass::int2b_t lowest() noexcept { return int2b_t{-2}; }
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::int2b_t min() noexcept { return lowest(); }
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::int2b_t max() noexcept { return int2b_t{1}; }
+
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = true;
+};
+
+template <>
+struct numeric_limits<mutlass::uint2b_t> {
+  MUTLASS_HOST_DEVICE static
+  mutlass::uint2b_t const lowest() noexcept { return uint2b_t{0}; }
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::uint2b_t const min() noexcept { return lowest(); }
+
+  MUTLASS_HOST_DEVICE static
+  mutlass::uint2b_t const max() noexcept { return uint2b_t{3}; }
+
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace platform
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/kernel_hardware_info.h b/vllm_musa/csrc/flash_attn/include/mutlass/kernel_hardware_info.h
new file mode 100644
index 000000000..f5714d095
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/kernel_hardware_info.h
@@ -0,0 +1,77 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#if !defined(__MUSACC_RTC__)
+#include "musa_runtime.h"
+
+#include "mutlass/trace.h"
+#endif
+
+namespace mutlass {
+
+struct KernelHardwareInfo {
+  //
+  // Data members
+  //
+  int device_id = 0;
+  int sm_count  = 0;
+
+  //
+  // Methods
+  //
+
+#if !defined(__MUSACC_RTC__)
+  static inline int
+  query_device_multiprocessor_count(int device_id = 0) {
+    musaError_t result = musaGetDevice(&device_id);
+    if (result != musaSuccess) {
+      MUTLASS_TRACE_HOST(
+        "  musaGetDevice() returned error "
+        << musaGetErrorString(result));
+      return 0;
+    }
+    int multiprocessor_count;
+    result = musaDeviceGetAttribute(&multiprocessor_count,
+      musaDevAttrMultiProcessorCount, device_id);
+    if (result != musaSuccess) {
+      MUTLASS_TRACE_HOST(
+        "  musaDeviceGetAttribute() returned error "
+        << musaGetErrorString(result));
+      return 0;
+    }
+    return multiprocessor_count;
+  }
+#endif
+};
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/kernel_hardware_info.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/kernel_hardware_info.hpp
new file mode 100644
index 000000000..50f384266
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/kernel_hardware_info.hpp
@@ -0,0 +1,36 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+// Simply import .h version of header so as to avoid breaking any existing MUTLASS builds
+// after .hpp was changed to .h
+#include "mutlass/kernel_hardware_info.h"
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/kernel_launch.h b/vllm_musa/csrc/flash_attn/include/mutlass/kernel_launch.h
new file mode 100644
index 000000000..749ae7f4c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/kernel_launch.h
@@ -0,0 +1,74 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 structures and helpers to launch MUSA kernels within MUTLASS.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Structure containing the basic launch configuration of a MUSA kernel.
+struct KernelLaunchConfiguration {
+
+  /// MUSA grid dimensions
+  dim3 grid;
+
+  /// MUSA threablock dimensions
+  dim3 block;
+
+  /// Bytes of dynamically allocated SMEM in addition to static SMEM
+  size_t dynamic_smem;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a KernellaunchConfiguration object
+  MUTLASS_HOST_DEVICE
+  KernelLaunchConfiguration(
+    dim3 _grid = dim3(1,1,1),
+    dim3 _block = dim3(1,1,1),
+    size_t _dynamic_smem = 0
+  ):
+    grid(_grid),
+    block(_block),
+    dynamic_smem(_dynamic_smem) { }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/layout/layout.h b/vllm_musa/csrc/flash_attn/include/mutlass/layout/layout.h
new file mode 100644
index 000000000..31ef18ccf
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/layout/layout.h
@@ -0,0 +1,63 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 layout functions used by TensorRef and derived classes.
+
+    Layout functions map logical coordinates to linear memory. They often require additional
+    data to describe strides between elements.
+
+    Layout functions must implement all members in the public interface of IdentityTensorLayout<>
+    defined in mutlass/tensor_ref.h.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/matrix_coord.h"
+#include "mutlass/layout/matrix.h"
+#include "mutlass/layout/pitch_linear.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/layout/vector.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace layout {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace layout
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/layout/matrix.h b/vllm_musa/csrc/flash_attn/include/mutlass/layout/matrix.h
new file mode 100644
index 000000000..4bea87c11
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/layout/matrix.h
@@ -0,0 +1,1350 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 layout functions used by TensorRef and derived classes.
+
+    Layout functions map logical coordinates to linear memory. They often require additional
+    data to describe strides between elements.
+
+    Layout functions must implement all members in the public interface of IdentityTensorLayout<>
+    defined in mutlass/tensor_ref.h.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/fast_math.h"
+#include "mutlass/matrix_coord.h"
+#include "mutlass/pitch_linear_coord.h"
+
+namespace mutlass {
+namespace layout {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Defines data layouts of various matrix formats usable by TensorRef and other classes.
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Mapping function for row-major matrices.
+class RowMajor {
+public:
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  RowMajor(LongIndex ldm = 0): stride_(ldm) { }
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  RowMajor(Stride stride): stride_(stride) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static RowMajor packed(MatrixCoord const &extent) {
+    return RowMajor(extent.column());
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    return LongIndex(coord.row()) * LongIndex(stride_[0]) + coord.column();
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+    return MatrixCoord(Index(offset / stride_[0]), Index(offset % stride_[0]));
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    return LongIndex(extent.row()) * LongIndex(stride_[0]);
+  }
+};
+
+/// Mapping function for column-major matrices.
+class ColumnMajor {
+public:
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  ColumnMajor(LongIndex ldm = 0): stride_(ldm) { }
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  ColumnMajor(Stride stride): stride_(stride) { }
+
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static ColumnMajor packed(MatrixCoord const &extent) {
+    return ColumnMajor(extent.row());
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    return LongIndex(coord.column()) * LongIndex(stride_[0]) + coord.row();
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+    return MatrixCoord(Index(offset % stride_[0]), Index(offset / stride_[0]));
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    return LongIndex(extent.column()) * LongIndex(stride_[0]);
+  }
+};
+
+/// Mapping function for interleaved matrices. Matrix is structured
+/// as row-major arrangement of fixed-size columns.
+template <int Interleave>
+struct RowMajorInterleaved {
+
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+  /// Size of interleaved columns
+  static int const kInterleave = Interleave;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  RowMajorInterleaved(LongIndex ldm = 0): stride_(ldm) { }
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  RowMajorInterleaved(Stride stride): stride_(stride) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static RowMajorInterleaved packed(MatrixCoord const &extent) {
+    return RowMajorInterleaved(extent.column() * kInterleave);
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    Index row_major = coord.row() / kInterleave;
+    Index row_minor = coord.row() % kInterleave;
+    return LongIndex(row_major) * LongIndex(stride_[0]) + LongIndex(coord.column()) * kInterleave + row_minor;
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+
+    Index row_major = Index(offset / stride_[0]);
+    Index residual = Index(offset % stride_[0]);
+
+    Index column = residual / kInterleave;
+    Index row_minor =  residual % kInterleave;
+
+    return MatrixCoord(row_major * kInterleave + row_minor, column);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    return (extent.row() + kInterleave - 1) / kInterleave * stride_[0];
+  }
+};
+
+/// Mapping function for interleaved matrices. Matrix is structured
+/// as column-major arrangement of fixed-size rows.
+template <int Interleave>
+struct ColumnMajorInterleaved {
+
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+  /// Size of interleaved columns
+  static int const kInterleave = Interleave;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  ColumnMajorInterleaved(LongIndex ldm = 0): stride_(ldm) { }
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  ColumnMajorInterleaved(Stride stride): stride_(stride) { }
+
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static ColumnMajorInterleaved packed(MatrixCoord const &extent) {
+    return ColumnMajorInterleaved(extent.row() * kInterleave);
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    Index column_major = coord.column() / kInterleave;
+    Index column_minor = coord.column() % kInterleave;
+    return LongIndex(column_major) * LongIndex(stride_[0]) + LongIndex(coord.row()) * kInterleave + column_minor;
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+
+    Index column_major = Index(offset / stride_[0]);
+    Index residual = Index(offset % stride_[0]);
+
+    Index row = residual / kInterleave;
+    Index column_minor =  residual % kInterleave;
+
+    return MatrixCoord(row, column_major * kInterleave + column_minor);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    return (extent.column() + kInterleave - 1) / kInterleave * stride_[0];
+  }
+};
+
+/// Enumerated type for canonical pitch-linear matrix layouts
+enum class Matrix {
+  kColumnMajor,       ///< leading dimension refers to stride between columns; stride along rows is 1
+  kRowMajor           ///< leading dimension refers to stride between rows; stride along columns is 1
+};
+
+/// Mapping function for scenario in which layout is row-major or column-major but this information
+/// is only available at runtime.
+struct ContiguousMatrix {
+
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+  /// Enumerated type indicating canonical matrix layout
+  Matrix layout_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  ContiguousMatrix(
+    Index ldm = 0,
+    Matrix layout = Matrix::kColumnMajor
+  ):
+    stride_(ldm), layout_(layout) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static ContiguousMatrix packed(
+    MatrixCoord const &extent,
+    Matrix layout = Matrix::kColumnMajor) {
+
+    Index ldm = 0;
+    if (layout == Matrix::kColumnMajor) {
+      ldm = extent.row();
+    }
+    else if (layout == Matrix::kRowMajor) {
+      ldm = extent.column();
+    }
+    return ContiguousMatrix(ldm, layout);
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    if (layout_ == Matrix::kColumnMajor) {
+      return coord.row() + coord.column() * stride_[0];
+    }
+    else if (layout_ == Matrix::kRowMajor) {
+      return coord.row() * stride_[0] + coord.column();
+    }
+    else {
+      // degenerate case
+      return 0;
+    }
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+    MUTLASS_UNUSED(offset);
+    return MatrixCoord(0, 0);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    if (layout_ == Matrix::kColumnMajor) {
+      return stride_[0] * extent.column();
+    }
+    else if (layout_ == Matrix::kRowMajor) {
+      return stride_[0] * extent.row();
+    }
+    else {
+      // degenerate case
+      return 0;
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Mapping function for scenario in which both rows and columns are separated by a stride.
+template <int Rank>
+struct AffineRankN {
+
+  /// Logical rank of tensor
+  static int const kRank = Rank;
+
+  /// Rank of stride vector
+  static int const kStrideRank = kRank;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = Coord<kRank, Index>;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  AffineRankN(
+    Stride const &stride = Stride()
+  ):
+    stride_(stride) { }
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  AffineRankN(
+    Coord<kRank/2, LongIndex> const &stride_m,
+    Coord<kRank/2, LongIndex> const &stride_n
+  ) {
+
+    // Concatenate the strides
+    MUTLASS_PRAGMA_UNROLL
+    for (int m = 0; m < kRank/2; ++m) {
+      stride_[m] = stride_m[m];
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < kRank/2; ++n) {
+      stride_[n + kRank/2] = stride_n[n];
+    }
+  }
+
+  /// Ctor for N = 2
+  MUTLASS_HOST_DEVICE
+  AffineRankN(
+    LongIndex const &stride_m,
+    LongIndex const &stride_n
+  ) {
+      stride_[0] = stride_m;
+      stride_[1] = stride_n;
+  }
+
+  /// Ctor for N = 2
+  MUTLASS_HOST_DEVICE
+  AffineRankN(
+    LongIndex const &stride
+  ) {
+      stride_[0] = stride;
+      stride_[1] = 1;
+  }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static AffineRankN packed(TensorCoord const &extent) {
+
+    AffineRankN layout;
+    layout.stride_[kRank - 1] = 1;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = kRank - 1; i > 0; --i) {
+      layout.stride_[i - 1] = layout.stride_[i] * extent[i];
+    }
+
+    return layout;
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(TensorCoord const &coord) const {
+    return dot(coord, stride_);
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  TensorCoord inverse(LongIndex offset) const {
+    return TensorCoord();
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &extent) const {
+    int idx = stride_.max_dim_index();
+    return extent[idx] * stride_[idx];
+  }
+};
+
+/// Mapping function for scenario in which both rows and columns are separated by a stride.
+/// Row stride is smaller than column stride in AffineRank2ColumnMajor.
+struct AffineRank2ColumnMajor {
+
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 2;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  AffineRank2ColumnMajor(
+    Stride const &stride = Stride()
+  ):
+    stride_(stride) { }
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  AffineRank2ColumnMajor(
+    LongIndex row_stride,           ///< stride between elements in consecutive rows
+    LongIndex column_stride         ///< stride between elements in consecutive columns
+  )
+    { stride_[0] = row_stride; stride_[1] = column_stride;}
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  AffineRank2ColumnMajor(
+    LongIndex stride
+  )
+    { stride_[0] = 1; stride_[1] = stride;}
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static AffineRank2ColumnMajor packed(MatrixCoord const &extent) {
+    return AffineRank2ColumnMajor(1, extent.row());
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    return dot(coord, stride_);
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+    MUTLASS_UNUSED(offset);
+    return MatrixCoord(0, 0);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    return extent.column() * stride_[1];
+  }
+};
+
+/// Mapping function for scenario in which both rows and columns are separated by a stride.
+/// Column stride is smaller than row stride in AffineRank2RowMajor.
+struct AffineRank2RowMajor {
+
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 2;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  AffineRank2RowMajor(
+    Stride const &stride = Stride()
+  ):
+    stride_(stride) { }
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  AffineRank2RowMajor(
+    LongIndex row_stride,           ///< stride between elements in consecutive rows
+    LongIndex column_stride         ///< stride between elements in consecutive columns
+  ) { stride_[0] = row_stride; stride_[1] = column_stride;}
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  AffineRank2RowMajor(
+    LongIndex stride
+  ) { stride_[0] = stride; stride_[1] = 1;}
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static AffineRank2RowMajor packed(MatrixCoord const &extent) {
+    return AffineRank2RowMajor(1, extent.row());
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    return dot(coord, stride_);
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+    MUTLASS_UNUSED(offset);
+    return MatrixCoord(0, 0);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    return extent.row() * stride_[0];
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Utility functions to convert stride_factor to the strides used by the Affine2 layout.
+//
+// stride_factor is the logical distance between two coorinates.
+//
+// All Coodinates used here are matrix coordinates.  stride[0] and extent[0] are for the
+// rows.  stride[1] and extent[1] are for the columns.
+template <typename Affine2Layout>
+  struct Affine2Layout_Factory {
+  MUTLASS_HOST_DEVICE
+  static Affine2Layout layout_factory(mutlass::Coord<2> const &extent, typename Affine2Layout::Stride stride_factor) {
+    return Affine2Layout::packed(extent);
+  }
+};
+
+template <>
+struct Affine2Layout_Factory<mutlass::layout::AffineRank2ColumnMajor> {
+MUTLASS_HOST_DEVICE
+static mutlass::layout::AffineRank2ColumnMajor layout_factory(
+  mutlass::Coord<2> const &extent,
+  typename mutlass::layout::AffineRank2ColumnMajor::Stride stride_factor) {
+    return mutlass::layout::AffineRank2ColumnMajor({ stride_factor[0], stride_factor[0] * stride_factor[1] * extent[0] });
+  }
+};
+
+template <>
+struct Affine2Layout_Factory<mutlass::layout::AffineRank2RowMajor> {
+MUTLASS_HOST_DEVICE
+static mutlass::layout::AffineRank2RowMajor layout_factory(
+  mutlass::Coord<2> const &extent,
+  typename mutlass::layout::AffineRank2RowMajor::Stride stride_factor) {
+    return mutlass::layout::AffineRank2RowMajor({ stride_factor[0] * stride_factor[1] * extent[1], stride_factor[1] });
+  }
+};
+
+// The base layout mutlass::layout::AffineRankN<2> is similar to AffineRank2ColumnMajor
+template <>
+struct Affine2Layout_Factory<mutlass::layout::AffineRankN<2>> {
+MUTLASS_HOST_DEVICE
+static mutlass::layout::AffineRankN<2> layout_factory(
+  mutlass::Coord<2> const &extent,
+  typename mutlass::layout::AffineRankN<2>::Stride stride_factor) {
+    return mutlass::layout::AffineRankN<2>({ stride_factor[0], stride_factor[0] * stride_factor[1] * extent[0] });
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Mapping function for block-linear matrices. Matrix is structured
+/// as column-major arrangement of 2D tiles (that are column-major).
+template <int BlockRows, int BlockColumns>
+struct ColumnMajorBlockLinear {
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+  /// Size of a block in rows
+  static int const kBlockRows = BlockRows;
+
+  /// Size of a block in columns
+  static int const kBlockColumns = BlockColumns;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  ColumnMajorBlockLinear(Index ldm = 0): stride_(ldm) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static ColumnMajorBlockLinear packed(MatrixCoord const &extent) {
+    return ColumnMajorBlockLinear(extent.row() * kBlockRows * kBlockColumns);
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    return
+      (coord.row() % kBlockRows) +
+      (coord.column() % kBlockColumns) * kBlockRows +
+      (coord.row() / kBlockRows) * kBlockRows * kBlockColumns +
+      (coord.column() / kBlockColumns) * stride_[0];
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+
+    return MatrixCoord(0, 0);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    return (extent.column() + kBlockColumns - 1) / kBlockColumns * stride_[0];
+  }
+};
+
+/// Mapping function for block-linear matrices. Matrix is structured
+/// as row-major arrangement of 2D tiles (that are row-major)
+template <int BlockRows, int BlockColumns>
+struct RowMajorBlockLinear {
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+  /// Size of a block in rows
+  static int const kBlockRows = BlockRows;
+
+  /// Size of a block in columns
+  static int const kBlockColumns = BlockColumns;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  RowMajorBlockLinear(Index ldm = 0): stride_(ldm) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static RowMajorBlockLinear packed(MatrixCoord const &extent) {
+    return RowMajorBlockLinear(extent.column() * kBlockRows * kBlockColumns);
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    return
+      (coord.column() % kBlockColumns) +
+      (coord.row() % kBlockRows) * kBlockColumns +
+      (coord.column() / kBlockColumns) * kBlockRows * kBlockColumns +
+      (coord.row() / kBlockRows) * stride_[0];
+  }
+
+  /// Inverse of layout function, mapping linear offset to logical coordinate
+  MUTLASS_HOST_DEVICE
+  MatrixCoord inverse(LongIndex offset) const {
+    return MatrixCoord(0, 0);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    return (extent.row() + kBlockRows - 1) / kBlockRows * stride_[0];
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct GeneralMatrix {
+
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 2;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = MatrixCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, Index>;
+
+private:
+  //
+  // Data members
+  //
+
+  Matrix layout_id_;
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  GeneralMatrix(): layout_id_(Matrix::kColumnMajor), stride_(make_Coord(0, 1)) { }
+
+  /// Ctor
+  MUTLASS_HOST_DEVICE
+  GeneralMatrix(
+    Matrix layout_id,
+    Index ldm,
+    Index interleave): layout_id_(layout_id), stride_(make_Coord(ldm, interleave)) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static GeneralMatrix packed(
+    MatrixCoord const &extent,
+    Matrix layout_id = Matrix::kColumnMajor,
+    Index interleave = 1) {
+
+    Index c;
+    if (layout_id == Matrix::kRowMajor) {
+      c = extent.column();
+    }
+    else {
+      c = extent.row();
+    }
+
+    Index ldm = c * interleave;
+
+    return GeneralMatrix(layout_id, ldm, interleave);
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (row, column)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord const &coord) const {
+    Index c, s;
+    if (layout_id_ == Matrix::kRowMajor) {
+      c = coord.column();
+      s = coord.row();
+    }
+    else {
+      s = coord.column();
+      c = coord.row();
+    }
+
+    Index v = s / stride_[1];
+    Index residual = (s % stride_[1]);
+
+    return LongIndex(c) * LongIndex(stride_[1]) + LongIndex(v) * LongIndex(stride_[0]) + residual;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix layout_id() const {
+    return layout_id_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix & layout_id() {
+    return layout_id_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index stride(int idx) const {
+    return stride_[idx];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  typename Stride::Index & stride(int idx) {
+    return stride_[idx];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(MatrixCoord const &extent) const {
+    Index s;
+    if (layout_id_ == Matrix::kRowMajor) {
+      s = extent.row();
+    }
+    else {
+      s = extent.column();
+    }
+
+    Index v = Index((s + stride_[1] - 1) / stride_[1]);
+    return LongIndex(v) * LongIndex(stride_[0]);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines transposes of matrix layouts
+template <typename Layout>
+struct LayoutTranspose;
+
+/// Transpose of row-major is column-major
+template <>
+struct LayoutTranspose<layout::RowMajor> {
+  using type = layout::ColumnMajor;
+};
+
+/// Transpose of column-major is row-major
+template <>
+struct LayoutTranspose<layout::ColumnMajor> {
+  using type = layout::RowMajor;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace layout
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/layout/permute.h b/vllm_musa/csrc/flash_attn/include/mutlass/layout/permute.h
new file mode 100644
index 000000000..4e1276adc
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/layout/permute.h
@@ -0,0 +1,830 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 layout functions used by GEMM+permute path for common tensor or matrix formats.
+
+    Like Layout functions, permute layout functions map logical coordinates to linear memory. They often require additional
+    data to describe strides between elements.
+
+    Permute layout functions must implement all members in the interface of NoPermute<> defined in this file. Address offset
+    computation lies in operator() with private member variables  {col_permute_, row_permute_ and stride_} as new addresses after permute op.
+*/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cassert>
+#else
+#include "assert.h"
+#endif
+#include "mutlass/mutlass.h"
+#include "mutlass/fast_math.h"
+#include "mutlass/layout/pitch_linear.h"
+#include "mutlass/layout/matrix.h"
+#include "mutlass/coord.h"
+#include "mutlass/tensor_coord.h"
+
+namespace mutlass {
+namespace layout {
+
+// template<PermuteTag, typename Layout, bool Inverse>
+// struct PermuteSelect {
+//   // Try to give a reasonable error message to the user
+//   static_assert(!platform::is_same<Permute, Permute>::value, // aka always_false<T>
+//                 "You've tried to use a layout permutation for which the implementation is not availble. "
+//                 "In order to provide an implementation for a particular combination of matrix layout "
+//                 "and direction (direct/inverse), please specialize PermuteSelect trait.");
+// };
+
+// Base template for defining specializations of permutation inverses
+template<typename Permute>
+struct InversePermute
+{
+  // Try to give a reasonable error message to the user
+  static_assert(!platform::is_same<Permute, Permute>::value, // aka always_false<T>
+                "To apply permutation to a GEMM input operand (A or B), an inverse permutation for the desired "
+                "permute class must be defined and enabled by specializing mutlass::layout::InversePermute trait.");
+};
+
+class PermuteBase {
+public:
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+};
+
+class NoPermute : public PermuteBase {
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor from matrix extent
+  MUTLASS_HOST_DEVICE
+  NoPermute(MatrixCoord extent, Index stride) { };
+
+  /// Constructor from pitch-linear extent
+  MUTLASS_HOST_DEVICE
+  NoPermute(PitchLinearCoord extent, Index stride) { };
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const { return 0; } // not correct but should never be called
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const { return 0; } // not correct but should never be called
+};
+
+template<>
+struct InversePermute<NoPermute> {
+  using type = NoPermute;
+};
+
+/// Helper trait to detect if permute operation is a noop
+template<typename Permute>
+inline bool constexpr is_trivial_permute = platform::is_same<Permute, mutlass::layout::NoPermute>::value;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Defines permute layouts of various tensor formats.
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//  Tensor4DPermute0213
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Permute layout function for 4-D permuted tensors with matrix (dimensions [M, N]) reshaped
+/// as [M/D1, D1, D2, N/D2]. Then perform permute([0, 2, 1, 3]) on the corresponding tensor.
+template <int D1, int D2>
+class Tensor4DPermute0213RowMajor : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index D3_;
+
+  Index stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermute0213RowMajor(MatrixCoord extent, Index stride) {
+
+    assert(extent.row() % D1 == 0);
+    assert(extent.column() % D2 == 0);
+
+    D3_ = extent.column() / D2;
+
+    stride_ = stride * D1 / D2;
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermute0213RowMajor(PitchLinearCoord extent, Index stride)
+  : Tensor4DPermute0213RowMajor(MatrixCoord(extent.strided(), extent.contiguous()), stride) {}
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    // [i,j,k,l] -> [i,k,j,l]
+    Index l = coord.column() % D3_;
+    Index k = coord.column() / D3_;
+    Index j = coord.row() % D1;
+    Index i = coord.row() / D1;
+
+    MatrixCoord permuted{k + i * D2, l + j * D3_};
+
+    return LongIndex(permuted.row()) * LongIndex(stride_) + LongIndex(permuted.column());
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.strided(), coord.contiguous()));
+  }
+};
+
+// Inverse for Tensor4DPermute0213 can be implemented by simply swapping D1 and D2
+template <int D1, int D2>
+class Tensor4DPermute0213RowMajorInverse : public Tensor4DPermute0213RowMajor<D2, D1> {
+public:
+  using Base = Tensor4DPermute0213RowMajor<D2, D1>;
+  using Base::Base;
+};
+
+template<int D1, int D2>
+struct InversePermute<Tensor4DPermute0213RowMajor<D1, D2>> {
+  using type = Tensor4DPermute0213RowMajorInverse<D1, D2>;
+};
+
+template<int D1, int D2>
+struct InversePermute<Tensor4DPermute0213RowMajorInverse<D1, D2>> {
+  using type = Tensor4DPermute0213RowMajor<D1, D2>;
+};
+
+/// Permute layout function for 4-D permuted tensors with matrix (dimensions [M, N]) reshaped
+/// as [M/D1, D1, D2, N/D2]. Then perform permute([0, 2, 1, 3]) on the corresponding tensor.
+template <int D1, int D2>
+class Tensor4DPermute0213ColumnMajor : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index D0_;
+
+  Index stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermute0213ColumnMajor(MatrixCoord extent, Index stride) {
+
+    assert(extent.row() % D1 == 0);
+    assert(extent.column() % D2 == 0);
+
+    D0_ = extent.row() / D1;
+
+    stride_ = stride * D2 / D1;
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermute0213ColumnMajor(PitchLinearCoord extent, Index stride)
+  : Tensor4DPermute0213ColumnMajor(MatrixCoord(extent.contiguous(), extent.strided()), stride) {}
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    // [i,j,k,l] -> [i,k,j,l]
+    Index l = coord.column() / D2;
+    Index k = coord.column() % D2;
+    Index j = coord.row() / D0_;
+    Index i = coord.row() % D0_;
+
+    MatrixCoord permuted{i + k * D0_, j + l * D1};
+
+    return LongIndex(permuted.row()) + LongIndex(permuted.column()) * LongIndex(stride_);
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.contiguous(), coord.strided()));
+  }
+};
+
+// Inverse for Tensor4DPermute0213 can be implemented by simply swapping D1 and D2
+template <int D1, int D2>
+class Tensor4DPermute0213ColumnMajorInverse : public Tensor4DPermute0213ColumnMajor<D2, D1> {
+public:
+  using Base = Tensor4DPermute0213ColumnMajor<D2, D1>;
+  using Base::Base;
+};
+
+template<int D1, int D2>
+struct InversePermute<Tensor4DPermute0213ColumnMajor<D1, D2>> {
+  using type = Tensor4DPermute0213ColumnMajorInverse<D1, D2>;
+};
+
+template<int D1, int D2>
+struct InversePermute<Tensor4DPermute0213ColumnMajorInverse<D1, D2>> {
+  using type = Tensor4DPermute0213ColumnMajor<D1, D2>;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//  Tensor4DPermuteBMM0213
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Permute layout function for 4-D permuted tensors for BMM with BMM tensor (dimensions [B, M, N]) reshaped
+/// as [B/D1, D1, M, N]. Then perform permute([0, 2, 1, 3]) on the corresponding whole BMM tensor.
+template <int D1>
+class Tensor4DPermuteBMM0213RowMajor : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index D3_;
+
+  Index stride_;
+
+  Index batch_stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermuteBMM0213RowMajor(MatrixCoord extent, Index stride) {
+
+    Index D2 = extent.row();
+    D3_ = extent.column();
+
+    stride_ = stride * D1;
+    batch_stride_ = D2 * stride_;
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermuteBMM0213RowMajor(PitchLinearCoord extent, Index stride)
+  : Tensor4DPermuteBMM0213RowMajor(MatrixCoord(extent.strided(), extent.contiguous()), stride) {}
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    // The batch index for BMM
+    Index BMM_batch_idx = blockIdx.z;
+
+    // [i,j,k,l] -> [i,k,j,l]
+    Index l = coord.column();
+    Index k = coord.row();
+    Index j = BMM_batch_idx % D1;
+    Index i = BMM_batch_idx / D1;
+
+    Index pbatch = i;
+    MatrixCoord pcoord{k, l + j * D3_};
+
+    return pbatch * LongIndex(batch_stride_) + pcoord.row() * LongIndex(stride_) + pcoord.column();
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.strided(), coord.contiguous()));
+  }
+};
+
+template <int D1>
+class Tensor4DPermuteBMM0213RowMajorInverse : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index D3_;
+
+  Index stride_;
+
+  Index batch_stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermuteBMM0213RowMajorInverse(MatrixCoord extent, Index stride) {
+
+    assert(extent.column() % D1 == 0);
+
+    Index D2 = extent.row();
+    D3_ = extent.column() / D1;
+
+    stride_ = stride / D1;
+
+    batch_stride_ = D2 * stride_;
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermuteBMM0213RowMajorInverse(PitchLinearCoord extent, Index stride)
+  : Tensor4DPermuteBMM0213RowMajorInverse(MatrixCoord(extent.strided(), extent.contiguous()), stride) {}
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    // The batch index for BMM
+    Index BMM_batch_idx = blockIdx.z;
+
+    // The following assumes grouping [(D0)->batch, (D2)->row, (D1,D3)->col]
+    Index l = coord.column() % D3_;
+    Index j = coord.column() / D3_;
+    Index k = coord.row();
+    Index i = BMM_batch_idx;
+
+    // compute original [batch, row, col] index
+    Index pbatch = j + i * D1;
+    MatrixCoord pcoord{k, l};
+
+    return pbatch * LongIndex(batch_stride_) + pcoord.row() * LongIndex(stride_) + pcoord.column();
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.strided(), coord.contiguous()));
+  }
+};
+
+template<int D1>
+struct InversePermute<Tensor4DPermuteBMM0213RowMajor<D1>> {
+  using type = Tensor4DPermuteBMM0213RowMajorInverse<D1>;
+};
+
+template<int D1>
+struct InversePermute<Tensor4DPermuteBMM0213RowMajorInverse<D1>> {
+  using type = Tensor4DPermuteBMM0213RowMajor<D1>;
+};
+
+/// Permute layout function for 4-D permuted tensors for BMM with BMM tensor (dimensions [B, M, N]) reshaped
+/// as [B/D1, D1, M, N]. Then perform permute([0, 3, 2, 1]) on the corresponding whole BMM tensor.
+template <int D1>
+class Tensor4DPermuteBMM0321ColumnMajor : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index D2_;
+
+  Index stride_;
+
+  Index batch_stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermuteBMM0321ColumnMajor(MatrixCoord extent, Index stride) {
+
+    D2_ = extent.row();
+    Index D3 = extent.column();
+
+    stride_ = stride * D1;
+    batch_stride_ = stride_ * D3;
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermuteBMM0321ColumnMajor(PitchLinearCoord extent, Index stride)
+  : Tensor4DPermuteBMM0321ColumnMajor(MatrixCoord(extent.contiguous(), extent.strided()), stride) {}
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    Index BMM_batch_idx = blockIdx.z;
+
+    // [i,j,k,l] -> [i,k,j,l]
+    Index l = coord.column();
+    Index k = coord.row();
+    Index j = BMM_batch_idx % D1;
+    Index i = BMM_batch_idx / D1;
+
+    Index pbatch = i;
+    MatrixCoord pcoord{k + j * D2_, l};
+
+    return pbatch * LongIndex(batch_stride_) + pcoord.row() + pcoord.column() * LongIndex(stride_);
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.contiguous(), coord.strided()));
+  }
+};
+
+template <int D1>
+class Tensor4DPermuteBMM0321ColumnMajorInverse : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index D2_;
+
+  Index stride_;
+
+  Index batch_stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermuteBMM0321ColumnMajorInverse(MatrixCoord extent, Index stride) {
+
+    assert(extent.row() % D1 == 0);
+
+    D2_ = extent.row() / D1;
+    Index D3 = extent.column();
+
+    stride_ = stride / D1;
+    batch_stride_ = stride_ * D3;
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor4DPermuteBMM0321ColumnMajorInverse(PitchLinearCoord extent, Index stride)
+  : Tensor4DPermuteBMM0321ColumnMajorInverse(MatrixCoord(extent.contiguous(), extent.strided()), stride) {}
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    Index BMM_batch_idx = blockIdx.z;
+
+    // The following assumes grouping [(D0)->batch, (D1,D2)->row, (D3)->col]
+    Index l = coord.column();
+    Index k = coord.row() % D2_;
+    Index j = coord.row() / D2_;
+    Index i = BMM_batch_idx;
+
+    Index pbatch = i * D1 + j;
+    MatrixCoord pcoord{k, l};
+
+    return pbatch * LongIndex(batch_stride_) + pcoord.row() + pcoord.column() * LongIndex(stride_);
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.contiguous(), coord.strided()));
+  }
+};
+
+template<int D1>
+struct InversePermute<Tensor4DPermuteBMM0321ColumnMajor<D1>> {
+  using type = Tensor4DPermuteBMM0321ColumnMajorInverse<D1>;
+};
+
+template<int D1>
+struct InversePermute<Tensor4DPermuteBMM0321ColumnMajorInverse<D1>> {
+  using type = Tensor4DPermuteBMM0321ColumnMajor<D1>;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//  Tensor5DPermute20314
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Permute layout function for 5-D permuted tensors with output matrix (dimension as [M, N]) reshaped
+/// as [M/T1, T1, T2, T3, N/T2/T3]. Then perform permute([2, 0, 3, 1, 4]) on the corresponding output tensor.
+template <int T1, int T2, int T3>
+class Tensor5DPermute20314RowMajor : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index T0_;
+
+  Index T4_;
+
+  Index stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor5DPermute20314RowMajor(MatrixCoord extent, Index stride) {
+
+    assert(extent.row() % T1 == 0);
+    assert(extent.column() % (T2 * T3) == 0);
+
+    T0_ = extent.row() / T1;
+    T4_ = extent.column() / (T2 * T3);
+
+    /// Update stride_permute with stride
+    stride_ = stride / T2 * T1; // stride in Elements
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor5DPermute20314RowMajor(PitchLinearCoord extent, Index stride)
+  : Tensor5DPermute20314RowMajor(MatrixCoord(extent.strided(), extent.contiguous()), stride) {}
+
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    // Permute as torch.permute(X1, [2, 0, 3, 1, 4]) -> 5D Tensor indices as [i,j,k,l,m], the dimension of X
+    // is [T0, T1, T2, T3, T4], after permutation the dim of X1 is [T2, T0, T3, T1, T4].
+
+    Index m = coord.column() % T4_;
+    Index l = (coord.column() / T4_) % T3;
+    Index k = (coord.column() / T4_) / T3;
+    Index j = coord.row() % T1;
+    Index i = coord.row() / T1;
+
+    MatrixCoord permuted{i + k * T0_, m + j * T4_ + l * T1 * T4_};
+
+    return LongIndex(permuted.row()) * LongIndex(stride_) + LongIndex(permuted.column());
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.strided(), coord.contiguous()));
+  }
+};
+
+/// Inverse for Tensor5DPermute20314 (could also be given a proper name, e.g. Tensor5DPermute13024).
+template <int T1, int T2, int T3>
+class Tensor5DPermute20314RowMajorInverse : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index T0_;
+
+  Index T4_;
+
+  // Permuted stride in units of elements
+  Index stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor5DPermute20314RowMajorInverse(MatrixCoord extent, Index stride) {
+
+    assert(extent.row() % T2 == 0);
+    assert(extent.column() % (T1 * T3) == 0);
+
+    T0_ = extent.row() / T2;
+    T4_ = extent.column() / (T1 * T3);
+
+    stride_ = stride / T1 * T2;
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor5DPermute20314RowMajorInverse(PitchLinearCoord extent, Index stride)
+  : Tensor5DPermute20314RowMajorInverse(MatrixCoord(extent.strided(), extent.contiguous()), stride) {}
+
+  /// Computes the offset after the inverse of permute operation in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    Index m = coord.column() % T4_;
+    Index j = (coord.column() / T4_) % T1;
+    Index l = (coord.column() / T4_) / T1;
+    Index i = coord.row() % T0_;
+    Index k = coord.row() / T0_;
+
+    MatrixCoord permuted{j + i * T1, m + l * T4_ + k * T3 * T4_};
+
+    return LongIndex(permuted.row()) * LongIndex(stride_) + LongIndex(permuted.column());
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.strided(), coord.contiguous()));
+  }
+};
+
+template<int T1, int T2, int T3>
+struct InversePermute<Tensor5DPermute20314RowMajor<T1, T2, T3>> {
+  using type = Tensor5DPermute20314RowMajorInverse<T1, T2, T3>;
+};
+
+template<int T1, int T2, int T3>
+struct InversePermute<Tensor5DPermute20314RowMajorInverse<T1, T2, T3>> {
+  using type = Tensor5DPermute20314RowMajor<T1, T2, T3>;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// Tensor5DPermute02413
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Permute layout function for 5-D permuted tensors with matrix (dimensions [M, N]) reshaped
+/// as [M/T1, T1, T2, T3, N/T2/T3]. Then perform permute([0, 2, 4, 1, 3]) on the corresponding tensor.
+template <int T1, int T2, int T3>
+class Tensor5DPermute02413ColumnMajor : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index T0_;
+
+  Index T4_;
+
+  Index stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor5DPermute02413ColumnMajor(MatrixCoord extent, Index stride) {
+
+    assert(extent.row() % T1 == 0);
+    assert(extent.column() % (T2 * T3) == 0);
+
+    T0_ = extent.row() / T1;
+    T4_ = extent.column() / (T2 * T3);
+
+    /// Update stride_permute with stride
+    stride_ = stride / T1 * T2; // stride in Elements
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor5DPermute02413ColumnMajor(PitchLinearCoord extent, Index stride)
+  : Tensor5DPermute02413ColumnMajor(MatrixCoord(extent.contiguous(), extent.strided()), stride) {}
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    // Permute as torch.permute(X1, [2, 0, 3, 1, 4]) -> 5D Tensor indices as [i,j,k,l,m], the dimension of X
+    // is [T0, T1, T2, T3, T4], after permutation the dim of X1 is [T0, T2, T4, T1, T3].
+
+    Index m = (coord.column() / T2) / T3;
+    Index l = (coord.column() / T2) % T3;
+    Index k = coord.column() % T2;
+    Index j = coord.row() / T0_;
+    Index i = coord.row() % T0_;
+
+    MatrixCoord permuted{i + k * T0_, m + j * T4_ + l * T4_ * T1};
+
+    return LongIndex(permuted.row()) + LongIndex(permuted.column()) * LongIndex(stride_);
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.contiguous(), coord.strided()));
+  }
+};
+
+/// Inverse for Tensor5DPermute02413ColumnMajor
+template <int T1, int T2, int T3>
+class Tensor5DPermute02413ColumnMajorInverse : public PermuteBase {
+private:
+  //
+  // Data members
+  //
+
+  Index T0_;
+
+  Index T4_;
+
+  // Permuted stride in units of elements
+  Index stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor5DPermute02413ColumnMajorInverse(MatrixCoord extent, Index stride) {
+
+    assert(extent.row() % T2 == 0);
+    assert(extent.column() % (T1 * T3) == 0);
+
+    T0_ = extent.row() / T2;
+    T4_ = extent.column() / (T1 * T3);
+
+    stride_ = stride / T2 * T1;
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Tensor5DPermute02413ColumnMajorInverse(PitchLinearCoord extent, Index stride)
+  : Tensor5DPermute02413ColumnMajorInverse(MatrixCoord(extent.contiguous(), extent.strided()), stride) {}
+
+  /// Computes the offset after the inverse of permute operation in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(MatrixCoord coord) const {
+
+    Index m = coord.column() % T4_;
+    Index j = (coord.column() / T4_) % T1;
+    Index l = (coord.column() / T4_) / T1;
+    Index i = coord.row() % T0_;
+    Index k = coord.row() / T0_;
+
+    MatrixCoord permuted{i + j * T0_, k + l * T2 + m * T2 * T3};
+
+    return LongIndex(permuted.row()) + LongIndex(permuted.column()) * LongIndex(stride_);
+  }
+
+  /// Computes the offset after Permute Op in logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return operator()(MatrixCoord(coord.contiguous(), coord.strided()));
+  }
+};
+
+template<int T1, int T2, int T3>
+struct InversePermute<Tensor5DPermute02413ColumnMajor<T1, T2, T3>> {
+  using type = Tensor5DPermute02413ColumnMajorInverse<T1, T2, T3>;
+};
+
+template<int T1, int T2, int T3>
+struct InversePermute<Tensor5DPermute02413ColumnMajorInverse<T1, T2, T3>> {
+  using type = Tensor5DPermute02413ColumnMajor<T1, T2, T3>;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace layout
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/layout/pitch_linear.h b/vllm_musa/csrc/flash_attn/include/mutlass/layout/pitch_linear.h
new file mode 100644
index 000000000..fcf92e697
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/layout/pitch_linear.h
@@ -0,0 +1,149 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 layout functions used by TensorRef and derived classes for pitch-linear memory.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+#include "mutlass/pitch_linear_coord.h"
+
+namespace mutlass {
+namespace layout {
+
+template <int Contiguous, int Strided>
+  using PitchLinearShape = mutlass::PitchLinearShape < Contiguous, Strided >;
+  using PitchLinearCoord = PitchLinearCoord;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Mapping function for pitch-linear memory
+class PitchLinear {
+public:
+  /// Logical rank of tensor
+  static int const kRank = 2;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = PitchLinearCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, LongIndex>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  PitchLinear(LongIndex ldm = 0): stride_(ldm) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  PitchLinear(Stride _stride): stride_(_stride) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static PitchLinear packed(TensorCoord const &extent) {
+    return PitchLinear(extent.contiguous());
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  /// Assumes coordinate has convention (contiguous, strided)
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(TensorCoord const &coord) const {
+    return LongIndex(coord.contiguous()) + LongIndex(coord.strided()) * LongIndex(stride_[0]);
+  }
+
+  /// Returns the logical coordinate given an offset.
+  MUTLASS_HOST_DEVICE
+  TensorCoord inverse(LongIndex index) const {
+    return make_Coord(
+      TensorCoord::Index(index % stride_[0]),
+      TensorCoord::Index(index / stride_[0])
+    );
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  LongIndex stride(int rank) const {
+    return stride_[rank];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  LongIndex & stride(int rank) {
+    return stride_[rank];
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &extent) const {
+    return extent.strided() * stride_[0];
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace layout
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/layout/tensor.h b/vllm_musa/csrc/flash_attn/include/mutlass/layout/tensor.h
new file mode 100644
index 000000000..b3b0f7c80
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/layout/tensor.h
@@ -0,0 +1,654 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 layout functions used by TensorRef and derived classes for common 4-D and 5-D
+      tensor formats.
+
+    Layout functions map logical coordinates to linear memory. They often require additional
+    data to describe strides between elements.
+
+    Layout functions must implement all members in the public interface of IdentityTensorLayout<>
+    defined in mutlass/tensor_ref.h.
+*/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cassert>
+#else
+#include "assert.h"
+#endif
+#include "mutlass/mutlass.h"
+#include "mutlass/fast_math.h"
+#include "mutlass/layout/pitch_linear.h"
+#include "mutlass/layout/matrix.h"
+#include "mutlass/coord.h"
+#include "mutlass/tensor_coord.h"
+
+namespace mutlass {
+namespace layout {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Defines data layouts of various tensor formats usable by TensorRef and other classes.
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag used for 3-D NWC tensors for 1D conv, only used in 3.x API
+class TensorNWC {};
+
+/// Tag used for n-D KCSRT tensors for nD conv, only used in 3.x API for wgrad output layouts
+class TensorKCS {};
+class TensorKCSR {};
+class TensorKCSRT {};
+
+/// Mapping function for 4-D NHWC tensors.
+class TensorNHWC {
+public:
+  /// Logical rank of tensor
+  static int const kRank = 4;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 3;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate (n, h, w, c)
+  using TensorCoord = Tensor4DCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member - [stride_w, stride_h, stride_n]
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorNHWC(Stride const &stride = Stride(0)): stride_(stride) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorNHWC(
+    typename Stride::Index stride_w,    ///< number of elements between adjacent W coordinates
+    typename Stride::Index stride_h,    ///< number of elements between adjacent H coordinates
+    typename Stride::Index stride_n     ///< number of elements between adjacent N coordinates
+  ):
+    stride_(make_Coord(stride_w, stride_h, stride_n)) { }
+
+  /// Constructor
+  // Once convolutions implement 64b stride this ctor can be deleted
+  MUTLASS_HOST_DEVICE
+  TensorNHWC(Coord<kStrideRank, LongIndex> const &stride):
+    stride_(make_Coord(
+      static_cast<typename Stride::Index>(stride[0]),
+      static_cast<typename Stride::Index>(stride[1]),
+      static_cast<typename Stride::Index>(stride[2]))
+    ) { }
+
+  /// Helper returns a layout to a tightly packed NHWC tensor.
+  MUTLASS_HOST_DEVICE
+  static TensorNHWC packed(TensorCoord const &extent) {
+    return TensorNHWC(
+      make_Coord(
+        extent.c(),
+        extent.w() * extent.c(),
+        extent.h() * extent.w() * extent.c()
+      )
+    );
+  }
+
+  /// Returns the offset of a coordinate (n, h, w, c) in linear memory.
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(TensorCoord const &coord) const {
+    return coord.c() +
+      LongIndex(stride_[0] * coord.w()) +
+      LongIndex(stride_[1] * coord.h()) +
+      LongIndex(stride_[2] * coord.n());
+  }
+
+  /// Returns the offset of a pitchlinear coordinate in linear memory.
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return coord.contiguous() + LongIndex(coord.strided() * stride_[2]);
+  }
+
+  /// Returns the logical coordinate (n, h, w, c) from a given offset in linear memory.
+  MUTLASS_HOST_DEVICE
+  TensorCoord inverse(LongIndex index) const {
+
+    int n = 0, h = 0, w = 0, c = 0;
+
+    #if defined(__MUSA_ARCH__)
+    int tmp = 0;
+    c = int(index % static_cast<int>(stride_[0]));
+
+    unsigned int hw_mul, hw_shr, w_mul, w_shr, c_mul, c_shr;
+
+    find_divisor(hw_mul, hw_shr, stride_[2]);
+    find_divisor(w_mul, w_shr, stride_[1]);
+    find_divisor(c_mul, c_shr, stride_[0]);
+
+    fast_divmod(n, tmp, index, int(stride_[2]), hw_mul, hw_shr);
+    fast_divmod(h, w, tmp, int(stride_[1]), w_mul, w_shr);
+    fast_divmod(w, tmp, w, int(stride_[0]), c_mul, c_shr);
+    #else
+
+    n = int(index / stride_[2]);
+    LongIndex residual = index % stride_[2];
+
+    h = int(residual / stride_[1]);
+    residual = (residual % stride_[1]);
+
+    w = int(residual / stride_[0]);
+    c = int(residual % stride_[0]);
+
+    #endif
+    return TensorCoord(n, h, w, c);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &extent) const {
+    // it does not make sense if the extent is larger than stride
+    // and we could not rely on the capacity calculation in such cases
+    // we could move this checkers to debug code only
+    if ((extent.c() > stride_[0])
+        || (extent.w() * stride_[0] > stride_[1])
+        || (extent.h() * stride_[1] > stride_[2])) {
+      assert(0);
+    }
+    return extent.n() * stride_[2];
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Mapping function for 4-D NCHW tensors.
+class TensorNCHW {
+public:
+  /// Logical rank of tensor
+  static int const kRank = 4;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 3;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = Tensor4DCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member - [w, hw, chw]
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorNCHW(Stride const &stride = Stride(0)): stride_(stride) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static TensorNCHW packed(TensorCoord const &extent) {
+    return TensorNCHW(
+      make_Coord(
+        extent.w(),
+        extent.w() * extent.h(),
+        extent.h() * extent.w() * extent.c()
+      )
+    );
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(TensorCoord const &coord) const {
+    return coord.w() +
+      LongIndex(stride_[0] * coord.h()) +
+      LongIndex(stride_[1] * coord.c()) +
+      LongIndex(stride_[2] * coord.n());
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &extent) const {
+    return extent.n() * stride_[2];
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Mapping function for 4-D NC/xHWx tensors.
+template <int Interleave>
+class TensorNCxHWx {
+public:
+
+  /// Interleaving quantity
+  static int const kInterleave = Interleave;
+
+  /// Logical rank of tensor
+  static int const kRank = 4;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 3;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = Tensor4DCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member - [Interleave x w, Interleave x wh, hwc]
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorNCxHWx(Stride const &stride = Stride(0)): stride_(stride) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorNCxHWx(
+    typename Stride::Index stride_w,    ///< number of elements between adjacent W coordinates
+    typename Stride::Index stride_h,    ///< number of elements between adjacent H coordinates
+    typename Stride::Index stride_n     ///< number of elements between adjacent N coordinates
+  ):
+    stride_(make_Coord(stride_w, stride_h, stride_n)) { }
+
+  /// Constructor
+  // Once convolutions implement 64b stride this ctor can be deleted
+  MUTLASS_HOST_DEVICE
+  TensorNCxHWx(Coord<kStrideRank, LongIndex> const &stride):
+    stride_(make_Coord(
+      static_cast<typename Stride::Index>(stride[0]),
+      static_cast<typename Stride::Index>(stride[1]),
+      static_cast<typename Stride::Index>(stride[2]))
+    ) { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static TensorNCxHWx packed(TensorCoord const &extent) {
+    return TensorNCxHWx(
+      make_Coord(
+        kInterleave * extent.w(),
+        kInterleave * extent.w() * extent.h(),
+        extent.h() * extent.w() * extent.c()
+      )
+    );
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(TensorCoord const &coord) const {
+
+    Index c_minor = (coord.c() % kInterleave);
+    Index c_major = (coord.c() / kInterleave);
+
+    return c_minor +
+      LongIndex(kInterleave * coord.w()) +
+      LongIndex(stride_[0] * coord.h()) +
+      LongIndex(stride_[1] * c_major) +
+      LongIndex(stride_[2] * coord.n());
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &extent) const {
+    return extent.n() * stride_[2];
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Mapping function for 4-D CxRSKx tensors.
+template <int Interleave>
+class TensorCxRSKx {
+public:
+
+  /// Interleaving quantity
+  static int const kInterleave = Interleave;
+
+  /// Logical rank of tensor
+  static int const kRank = 4;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 3;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = Tensor4DCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member - [Interleave x n, Interleave x nw, Interleave x nwh]
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorCxRSKx(Stride const &stride = Stride(0)): stride_(stride) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorCxRSKx(
+    typename Stride::Index stride_w,    ///< number of elements between adjacent W coordinates
+    typename Stride::Index stride_h,    ///< number of elements between adjacent H coordinates
+    typename Stride::Index stride_n     ///< number of elements between adjacent N coordinates
+  ):
+    stride_(make_Coord(stride_w, stride_h, stride_n)) { }
+
+  /// Constructor
+  // Once convolutions implement 64b stride this ctor can be deleted
+  MUTLASS_HOST_DEVICE
+  TensorCxRSKx(Coord<kStrideRank, LongIndex> const &stride):
+    stride_(make_Coord(
+      static_cast<typename Stride::Index>(stride[0]),
+      static_cast<typename Stride::Index>(stride[1]),
+      static_cast<typename Stride::Index>(stride[2]))
+    ) { }
+
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static TensorCxRSKx packed(TensorCoord const &extent) {
+    return TensorCxRSKx(
+      make_Coord(
+        kInterleave * extent.n(),
+        kInterleave * extent.n() * extent.w(),
+        kInterleave * extent.n() * extent.w() * extent.h()
+      )
+    );
+  }
+
+  /// Returns the offset of a coordinate in linear memory.
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(TensorCoord const &coord) const {
+
+    Index c_minor = (coord.c() % kInterleave);
+    Index c_major = (coord.c() / kInterleave);
+
+    return c_minor +
+      LongIndex(kInterleave * coord.n()) +
+      LongIndex(stride_[0] * coord.w()) +
+      LongIndex(stride_[1] * coord.h()) +
+      LongIndex(stride_[2] * c_major);
+  }
+
+  /// Returns the offset of a pitchlinear coordinate in linear memory.
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord const &coord) const {
+    return (coord.contiguous() % kInterleave) +
+      LongIndex((coord.contiguous() / kInterleave) * stride_[2]) +
+      LongIndex(coord.strided() * kInterleave);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &extent) const {
+    return (extent.c() / kInterleave * stride_[2]);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Mapping function for 5-D NDHWC tensors.
+class TensorNDHWC {
+public:
+  /// Logical rank of tensor
+  static int const kRank = 5;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 4;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate (n, d, h, w, c)
+  using TensorCoord = Tensor5DCoord;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank>;
+
+private:
+  //
+  // Data members
+  //
+
+  /// Stride data member - [c, wc, hwc, dhwc]
+  Stride stride_;
+
+public:
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorNDHWC(Stride const &stride = Stride(0)): stride_(stride) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  TensorNDHWC(
+    typename Stride::Index c,
+    typename Stride::Index wc,
+    typename Stride::Index hwc,
+    typename Stride::Index dhwc):
+  stride_(make_Coord(c, wc, hwc, dhwc)) { }
+
+  /// Constructor
+  // Once convolutions implement 64b stride this ctor can be deleted
+  MUTLASS_HOST_DEVICE
+  TensorNDHWC(Coord<kStrideRank, LongIndex> const &stride):
+    stride_(make_Coord(
+      static_cast<typename Stride::Index>(stride[0]),
+      static_cast<typename Stride::Index>(stride[1]),
+      static_cast<typename Stride::Index>(stride[2]),
+      static_cast<typename Stride::Index>(stride[3]))
+    ) { }
+
+  /// Helper returns a layout to a tightly packed NHWC tensor.
+  MUTLASS_HOST_DEVICE
+  static TensorNDHWC packed(TensorCoord const &extent) {
+    return TensorNDHWC(
+      make_Coord(
+        extent.c(),
+        extent.w() * extent.c(),
+        extent.h() * extent.w() * extent.c(),
+        extent.d() * extent.h() * extent.w() * extent.c()
+      )
+    );
+  }
+
+  /// Returns the offset of a coordinate (n, d, h, w, c) in linear memory.
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(TensorCoord const &coord) const {
+    return coord.c() +
+      LongIndex(stride_[0] * coord.w()) +
+      LongIndex(stride_[1] * coord.h()) +
+      LongIndex(stride_[2] * coord.d()) +
+      LongIndex(stride_[3] * coord.n());
+  }
+
+  /// Returns the offset of a pitchlinear coordinate in linear memory.
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(PitchLinearCoord coord) const {
+    return coord.contiguous() + LongIndex(coord.strided() * stride_[3]);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &extent) const {
+    // it does not make sense if the extent is larger than stride
+    // and we could not rely on the capacity calculation in such cases
+    // we could move this checkers to debug code only
+    if ((extent.c() > stride_[0])
+        || (extent.w() * stride_[0] > stride_[1])
+        || (extent.h() * stride_[1] > stride_[2])
+        || (extent.d() * stride_[2] > stride_[3])) {
+      assert(0);
+    }
+    return extent.n() * stride_[3];
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tag used for linearized tensors with shape (NW, C) for 1D conv, only used in 3.x API
+class TensorLinearizedNWC {};
+/// Tag used for linearized tensors with shape (NHW, C) for 2D conv, only used in 3.x API
+class TensorLinearizedNHWC : public TensorNHWC {};
+/// Tag used for linearized tensors with shape (NDHW, C) for 3D conv, only used in 3.x API
+class TensorLinearizedNDHWC : public TensorNDHWC {};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace layout
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/layout/vector.h b/vllm_musa/csrc/flash_attn/include/mutlass/layout/vector.h
new file mode 100644
index 000000000..98a9ea9a3
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/layout/vector.h
@@ -0,0 +1,106 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 layout functions used for rank=1 vectors.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+
+namespace mutlass {
+namespace layout {
+
+/// Tensor layout for densely packed vectors.
+class PackedVectorLayout {
+public:
+  /// Logical rank of tensor
+  static int const kRank = 1;
+
+  /// Rank of stride vector
+  static int const kStrideRank = 1;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = Coord<kRank, Index>;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, Index>;
+
+private:
+
+  //
+  // No actual stride vector stored
+  //
+
+public:
+
+  //
+  // Methods
+  //
+
+  MUTLASS_HOST_DEVICE
+  PackedVectorLayout() { }
+
+  /// Helper returns a layout to a tightly packed tensor
+  MUTLASS_HOST_DEVICE
+  static PackedVectorLayout packed(TensorCoord const &size) {
+    MUTLASS_UNUSED(size);
+    return PackedVectorLayout();
+  }
+
+  /// Returns the offset of a coordinate in linear memory
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(TensorCoord const &coord) const {
+    return coord[0];
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return make_Coord(1);
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &size) const {
+    return size[0];
+  }
+};
+
+} // namespace layout
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/matrix.h b/vllm_musa/csrc/flash_attn/include/mutlass/matrix.h
new file mode 100644
index 000000000..403088163
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/matrix.h
@@ -0,0 +1,14130 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Matrix classes with value semantics.
+*/
+
+#pragma once
+
+#if !defined(__MUSACC_RTC__)
+#include <iosfwd>
+#include <cmath>
+#endif
+
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+#include "mutlass/coord.h"
+#include "mutlass/fast_math.h"
+#include "mutlass/layout/matrix.h"
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Primary template with partial specializations to follow
+template <typename Element, int Rows, int Columns> struct Matrix;
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 1-by-2 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 1, 2> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 1;
+
+  /// Number of columns in matrix
+  static int const kColumns = 2;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 2;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 1-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 1-by-2 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> transpose() const {
+    Matrix<Element, 2, 1> mt;
+
+    mt.data[0] = data[0];
+    mt.data[1] = data[1];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 1 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 1 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> row(int i) const {
+    return slice_1x2(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 2> const &v, int i = 0) {
+    return set_slice_1x2(v, i, 0);
+  }
+
+  /// Forms a 1-by-2 matrix by horizontally concatenating an Element with an Element
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Element lhs, Element rhs) {
+    return Matrix(
+      lhs, rhs);
+  }
+
+  /// Concatenates this matrix with a an Element to form a 1-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> hcat(Element rhs) const {
+    return Matrix<Element, 1, 3>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 1-by-2 matrix to form a 1-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> hcat(Matrix<Element, 1, 2> const & rhs) const {
+    return Matrix<Element, 1, 4>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 1-by-2 matrix to form a 2-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> vcat(Matrix<Element, 1, 2> const & rhs) const {
+    return Matrix<Element, 2, 2>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-2 matrix to form a 3-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> vcat(Matrix<Element, 2, 2> const & rhs) const {
+    return Matrix<Element, 3, 2>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 3-by-2 matrix to form a 4-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> vcat(Matrix<Element, 3, 2> const & rhs) const {
+    return Matrix<Element, 4, 2>::vcat(*this, rhs);
+  }
+
+  /// Elementwise add operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+
+    return result;
+  }
+
+  /// Elementwise add operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+
+    return result;
+  }
+
+  /// Scalar divide operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (1-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+
+    return m;
+  }
+
+  /// Matrix product of size 1-by-1-by-2
+  MUTLASS_HOST_DEVICE
+  Element product(Matrix<Element, 2, 1> const &rhs, Element accum = Element()) const {
+
+    // k=0
+    accum += data[0] * rhs.data[0];
+
+    // k=1
+    accum += data[1] * rhs.data[1];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-1-by-2
+  MUTLASS_HOST_DEVICE
+  Element operator*(Matrix<Element, 2, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> product(
+    Matrix<Element, 2, 2> const &rhs,
+    Matrix<Element, 1, 2> accum = Matrix<Element, 1, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> operator*(Matrix<Element, 2, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 2, 2> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 1-by-3-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> product(
+    Matrix<Element, 2, 3> const &rhs,
+    Matrix<Element, 1, 3> accum = Matrix<Element, 1, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-3-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> operator*(Matrix<Element, 2, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-4-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> product(
+    Matrix<Element, 2, 4> const &rhs,
+    Matrix<Element, 1, 4> accum = Matrix<Element, 1, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-4-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> operator*(Matrix<Element, 2, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Dot product of vectors with extent 2
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 2, 1> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    return accum;
+  }
+
+  /// Dot product of vectors with extent 2
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 1, 2> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    return accum;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 1-by-2 matrix
+template <typename Element>
+using Matrix1x2 = Matrix<Element, 1, 2>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix1x2<Element> make_Matrix1x2(
+    Element _0_0, Element _0_1
+) {
+  return Matrix1x2<Element>(
+  _0_0, _0_1
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 1-by-3 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 1, 3> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 1;
+
+  /// Number of columns in matrix
+  static int const kColumns = 3;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 3;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 1-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 1-by-3 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1, Element _0_2
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;  data[2] = _0_2;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> transpose() const {
+    Matrix<Element, 3, 1> mt;
+
+    mt.data[0] = data[0];
+    mt.data[1] = data[1];
+    mt.data[2] = data[2];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 1 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 1 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> slice_1x3(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x3(Matrix<Element, 1, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> row(int i) const {
+    return slice_1x3(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 3> const &v, int i = 0) {
+    return set_slice_1x3(v, i, 0);
+  }
+
+  /// Forms a 1-by-3 matrix by horizontally concatenating an Element with a 1-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Element lhs, Matrix<Element, 1, 2> const & rhs) {
+    return Matrix(
+      lhs, rhs.at(0, 0), rhs.at(0, 1));
+  }
+
+  /// Forms a 1-by-3 matrix by horizontally concatenating a 1-by-2 matrix with an Element
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 1, 2> const & lhs, Element rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), rhs);
+  }
+
+  /// Concatenates this matrix with a an Element to form a 1-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> hcat(Element rhs) const {
+    return Matrix<Element, 1, 4>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 1-by-3 matrix to form a 2-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> vcat(Matrix<Element, 1, 3> const & rhs) const {
+    return Matrix<Element, 2, 3>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-3 matrix to form a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> vcat(Matrix<Element, 2, 3> const & rhs) const {
+    return Matrix<Element, 3, 3>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 3-by-3 matrix to form a 4-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> vcat(Matrix<Element, 3, 3> const & rhs) const {
+    return Matrix<Element, 4, 3>::vcat(*this, rhs);
+  }
+
+  /// Elementwise add operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+    result.data[2] = data[2] + rhs.data[2];
+
+    return result;
+  }
+
+  /// Elementwise add operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+    data[2] += rhs.data[2];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+    result.data[2] = data[2] - rhs.data[2];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+    data[2] -= rhs.data[2];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+    result.data[2] = data[2] * rhs.data[2];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+    result.data[2] = data[2] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+    data[2] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+    result.data[2] = data[2] / rhs.data[2];
+
+    return result;
+  }
+
+  /// Scalar divide operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+    result.data[2] = data[2] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+    data[2] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (1-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+    data[2] /= rhs.data[2];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+
+    return m;
+  }
+
+  /// Matrix product of size 1-by-1-by-3
+  MUTLASS_HOST_DEVICE
+  Element product(Matrix<Element, 3, 1> const &rhs, Element accum = Element()) const {
+
+    // k=0
+    accum += data[0] * rhs.data[0];
+
+    // k=1
+    accum += data[1] * rhs.data[1];
+
+    // k=2
+    accum += data[2] * rhs.data[2];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-1-by-3
+  MUTLASS_HOST_DEVICE
+  Element operator*(Matrix<Element, 3, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-2-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> product(
+    Matrix<Element, 3, 2> const &rhs,
+    Matrix<Element, 1, 2> accum = Matrix<Element, 1, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[4];
+    accum.data[1] += data[2] * rhs.data[5];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-2-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> operator*(Matrix<Element, 3, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> product(
+    Matrix<Element, 3, 3> const &rhs,
+    Matrix<Element, 1, 3> accum = Matrix<Element, 1, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[6];
+    accum.data[1] += data[2] * rhs.data[7];
+    accum.data[2] += data[2] * rhs.data[8];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> operator*(Matrix<Element, 3, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 3, 3> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 1-by-4-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> product(
+    Matrix<Element, 3, 4> const &rhs,
+    Matrix<Element, 1, 4> accum = Matrix<Element, 1, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[8];
+    accum.data[1] += data[2] * rhs.data[9];
+    accum.data[2] += data[2] * rhs.data[10];
+    accum.data[3] += data[2] * rhs.data[11];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-4-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> operator*(Matrix<Element, 3, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Dot product of vectors with extent 3
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 3, 1> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    accum += data[2] * rhs.data[2];
+    return accum;
+  }
+
+  /// Dot product of vectors with extent 3
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 1, 3> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    accum += data[2] * rhs.data[2];
+    return accum;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+
+    return accum;
+  }
+
+  /// Cross product
+  MUTLASS_HOST_DEVICE
+  Matrix cross(Matrix const &rhs) const {
+    return Matrix(
+      data[1] * rhs.data[2] - data[2] * rhs.data[1],
+      data[0] * rhs.data[2] - data[2] * rhs.data[1],
+      data[0] * rhs.data[1] - data[1] * rhs.data[0]
+    );
+  }
+
+};
+
+/// Template alias for 1-by-3 matrix
+template <typename Element>
+using Matrix1x3 = Matrix<Element, 1, 3>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix1x3<Element> make_Matrix1x3(
+    Element _0_0, Element _0_1, Element _0_2
+) {
+  return Matrix1x3<Element>(
+  _0_0, _0_1, _0_2
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 1-by-4 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 1, 4> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 1;
+
+  /// Number of columns in matrix
+  static int const kColumns = 4;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 4;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 1-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 1-by-4 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1, Element _0_2, Element _0_3
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;  data[2] = _0_2;  data[3] = _0_3;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> transpose() const {
+    Matrix<Element, 4, 1> mt;
+
+    mt.data[0] = data[0];
+    mt.data[1] = data[1];
+    mt.data[2] = data[2];
+    mt.data[3] = data[3];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 1 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 1 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> slice_1x3(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x3(Matrix<Element, 1, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> slice_1x4(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x4(Matrix<Element, 1, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> row(int i) const {
+    return slice_1x4(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 4> const &v, int i = 0) {
+    return set_slice_1x4(v, i, 0);
+  }
+
+  /// Forms a 1-by-4 matrix by horizontally concatenating an Element with a 1-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Element lhs, Matrix<Element, 1, 3> const & rhs) {
+    return Matrix(
+      lhs, rhs.at(0, 0), rhs.at(0, 1), rhs.at(0, 2));
+  }
+
+  /// Forms a 1-by-4 matrix by horizontally concatenating a 1-by-2 matrix with a 1-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 1, 2> const & lhs, Matrix<Element, 1, 2> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0), rhs.at(0, 1));
+  }
+
+  /// Forms a 1-by-4 matrix by horizontally concatenating a 1-by-3 matrix with an Element
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 1, 3> const & lhs, Element rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), lhs.at(0, 2), rhs);
+  }
+
+  /// Concatenates this matrix with a a 1-by-4 matrix to form a 2-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> vcat(Matrix<Element, 1, 4> const & rhs) const {
+    return Matrix<Element, 2, 4>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-4 matrix to form a 3-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> vcat(Matrix<Element, 2, 4> const & rhs) const {
+    return Matrix<Element, 3, 4>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 3-by-4 matrix to form a 4-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> vcat(Matrix<Element, 3, 4> const & rhs) const {
+    return Matrix<Element, 4, 4>::vcat(*this, rhs);
+  }
+
+  /// Elementwise add operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+    result.data[2] = data[2] + rhs.data[2];
+    result.data[3] = data[3] + rhs.data[3];
+
+    return result;
+  }
+
+  /// Elementwise add operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+    data[2] += rhs.data[2];
+    data[3] += rhs.data[3];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+    result.data[2] = data[2] - rhs.data[2];
+    result.data[3] = data[3] - rhs.data[3];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+    data[2] -= rhs.data[2];
+    data[3] -= rhs.data[3];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+    result.data[2] = data[2] * rhs.data[2];
+    result.data[3] = data[3] * rhs.data[3];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+    result.data[2] = data[2] * s;
+    result.data[3] = data[3] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+    data[2] *= s;
+    data[3] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+    result.data[2] = data[2] / rhs.data[2];
+    result.data[3] = data[3] / rhs.data[3];
+
+    return result;
+  }
+
+  /// Scalar divide operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+    result.data[2] = data[2] / s;
+    result.data[3] = data[3] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+    data[2] /= s;
+    data[3] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (1-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+    data[2] /= rhs.data[2];
+    data[3] /= rhs.data[3];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+
+    return m;
+  }
+
+  /// Matrix product of size 1-by-1-by-4
+  MUTLASS_HOST_DEVICE
+  Element product(Matrix<Element, 4, 1> const &rhs, Element accum = Element()) const {
+
+    // k=0
+    accum += data[0] * rhs.data[0];
+
+    // k=1
+    accum += data[1] * rhs.data[1];
+
+    // k=2
+    accum += data[2] * rhs.data[2];
+
+    // k=3
+    accum += data[3] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-1-by-4
+  MUTLASS_HOST_DEVICE
+  Element operator*(Matrix<Element, 4, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-2-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> product(
+    Matrix<Element, 4, 2> const &rhs,
+    Matrix<Element, 1, 2> accum = Matrix<Element, 1, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[4];
+    accum.data[1] += data[2] * rhs.data[5];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[6];
+    accum.data[1] += data[3] * rhs.data[7];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-2-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> operator*(Matrix<Element, 4, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-3-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> product(
+    Matrix<Element, 4, 3> const &rhs,
+    Matrix<Element, 1, 3> accum = Matrix<Element, 1, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[6];
+    accum.data[1] += data[2] * rhs.data[7];
+    accum.data[2] += data[2] * rhs.data[8];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[9];
+    accum.data[1] += data[3] * rhs.data[10];
+    accum.data[2] += data[3] * rhs.data[11];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-3-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> operator*(Matrix<Element, 4, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> product(
+    Matrix<Element, 4, 4> const &rhs,
+    Matrix<Element, 1, 4> accum = Matrix<Element, 1, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[8];
+    accum.data[1] += data[2] * rhs.data[9];
+    accum.data[2] += data[2] * rhs.data[10];
+    accum.data[3] += data[2] * rhs.data[11];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[12];
+    accum.data[1] += data[3] * rhs.data[13];
+    accum.data[2] += data[3] * rhs.data[14];
+    accum.data[3] += data[3] * rhs.data[15];
+
+    return accum;
+  }
+
+  /// Matrix product of size 1-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> operator*(Matrix<Element, 4, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 1-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 4, 4> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Dot product of vectors with extent 4
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 4, 1> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    accum += data[2] * rhs.data[2];
+    accum += data[3] * rhs.data[3];
+    return accum;
+  }
+
+  /// Dot product of vectors with extent 4
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 1, 4> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    accum += data[2] * rhs.data[2];
+    accum += data[3] * rhs.data[3];
+    return accum;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 1-by-4 matrix
+template <typename Element>
+using Matrix1x4 = Matrix<Element, 1, 4>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix1x4<Element> make_Matrix1x4(
+    Element _0_0, Element _0_1, Element _0_2, Element _0_3
+) {
+  return Matrix1x4<Element>(
+  _0_0, _0_1, _0_2, _0_3
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 2-by-1 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 2, 1> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 2;
+
+  /// Number of columns in matrix
+  static int const kColumns = 1;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 2;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 2-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 2-by-1 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0,
+    Element _1_0
+  ) {
+
+    data[0] = _0_0;
+    data[1] = _1_0;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> transpose() const {
+    Matrix<Element, 1, 2> mt;
+
+    mt.data[0] = data[0];
+    mt.data[1] = data[1];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 2 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 2 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 1 + j + 0];
+    m.data[1] = data[i * 1 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 1 + j + 0] = m.data[0];
+    data[i * 1 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> column(int j) const {
+    return slice_2x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 2, 1> const &v, int j =0) {
+    return set_slice_2x1(v, 0, j);
+  }
+
+  /// Concatenates this matrix with a a 2-by-1 matrix to form a 2-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> hcat(Matrix<Element, 2, 1> const & rhs) const {
+    return Matrix<Element, 2, 2>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-2 matrix to form a 2-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> hcat(Matrix<Element, 2, 2> const & rhs) const {
+    return Matrix<Element, 2, 3>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-3 matrix to form a 2-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> hcat(Matrix<Element, 2, 3> const & rhs) const {
+    return Matrix<Element, 2, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 2-by-1 matrix by vertically concatenating an Element with an Element
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Element upper, Element lower) {
+    return Matrix(
+      upper
+      , lower);
+  }
+
+  /// Concatenates this matrix with a an Element to form a 3-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> vcat(Element rhs) const {
+    return Matrix<Element, 3, 1>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-1 matrix to form a 4-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> vcat(Matrix<Element, 2, 1> const & rhs) const {
+    return Matrix<Element, 4, 1>::vcat(*this, rhs);
+  }
+
+  /// Elementwise add operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+
+    result.data[1] = data[1] + rhs.data[1];
+
+    return result;
+  }
+
+  /// Elementwise add operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+
+    data[1] += rhs.data[1];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+
+    result.data[1] = data[1] - rhs.data[1];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+
+    data[1] -= rhs.data[1];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+
+    result.data[1] = data[1] * rhs.data[1];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+
+    result.data[1] = data[1] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+
+    data[1] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+
+    result.data[1] = data[1] / rhs.data[1];
+
+    return result;
+  }
+
+  /// Scalar divide operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+
+    result.data[1] = data[1] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+
+    data[1] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (2-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+
+    data[1] /= rhs.data[1];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+
+    return m;
+  }
+
+  /// Matrix product of size 2-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> product(
+    Matrix<Element, 1, 1> const &rhs,
+    Matrix<Element, 2, 1> accum = Matrix<Element, 2, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[1] * rhs.data[0];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> operator*(Matrix<Element, 1, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 1, 1> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 2-by-2-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> product(
+    Matrix<Element, 1, 2> const &rhs,
+    Matrix<Element, 2, 2> accum = Matrix<Element, 2, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[1] * rhs.data[0];
+    accum.data[3] += data[1] * rhs.data[1];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-2-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> operator*(Matrix<Element, 1, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-3-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> product(
+    Matrix<Element, 1, 3> const &rhs,
+    Matrix<Element, 2, 3> accum = Matrix<Element, 2, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[1] * rhs.data[0];
+    accum.data[4] += data[1] * rhs.data[1];
+    accum.data[5] += data[1] * rhs.data[2];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-3-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> operator*(Matrix<Element, 1, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-4-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> product(
+    Matrix<Element, 1, 4> const &rhs,
+    Matrix<Element, 2, 4> accum = Matrix<Element, 2, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[1] * rhs.data[0];
+    accum.data[5] += data[1] * rhs.data[1];
+    accum.data[6] += data[1] * rhs.data[2];
+    accum.data[7] += data[1] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-4-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> operator*(Matrix<Element, 1, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Dot product of vectors with extent 2
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 2, 1> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    return accum;
+  }
+
+  /// Dot product of vectors with extent 2
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 1, 2> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    return accum;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 2-by-1 matrix
+template <typename Element>
+using Matrix2x1 = Matrix<Element, 2, 1>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix2x1<Element> make_Matrix2x1(
+    Element _0_0,
+    Element _1_0
+) {
+  return Matrix2x1<Element>(
+  _0_0,
+  _1_0
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 2-by-2 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 2, 2> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 2;
+
+  /// Number of columns in matrix
+  static int const kColumns = 2;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 4;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 2-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 2-by-2 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1,
+    Element _1_0, Element _1_1
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;
+    data[2] = _1_0;  data[3] = _1_1;
+  }
+
+  /// Constucts a 2-by-2 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 2> const &row_0,
+    Matrix<Element, 1, 2> const &row_1
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_1.data[0];
+    data[3] = row_1.data[1];
+  }
+
+  /// Static method to construct a 2-by-2 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 2, 1> const &column_0,
+    Matrix<Element, 2, 1> const &column_1
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_0.data[1];
+    result.data[3] = column_1.data[1];
+    return result;
+  }
+
+  /// Constructs an identity matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix identity() {
+    Matrix m;
+
+    m.data[0] = Element(1);
+    m.data[3] = Element(1);
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 2, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[3] = diag.data[1];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 2> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[3] = diag.data[1];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> diagonal() const {
+    Matrix<Element, 2, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[3];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> transpose() const {
+    Matrix<Element, 2, 2> mt;
+
+    mt.data[0] = data[0];
+    mt.data[2] = data[1];
+    mt.data[1] = data[2];
+    mt.data[3] = data[3];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 2 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 2 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> row(int i) const {
+    return slice_1x2(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 2> const &v, int i = 0) {
+    return set_slice_1x2(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 2] = m.data[1];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> column(int j) const {
+    return slice_2x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 2, 1> const &v, int j =0) {
+    return set_slice_2x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+    m.data[2] = data[i * 2 + j + 2];
+    m.data[3] = data[i * 2 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+    data[i * 2 + j + 2] = m.data[2];
+    data[i * 2 + j + 3] = m.data[3];
+
+    return *this;
+  }
+
+  /// Forms a 2-by-2 matrix by horizontally concatenating a 2-by-1 matrix with a 2-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 2, 1> const & lhs, Matrix<Element, 2, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0)
+      , lhs.at(1, 0), rhs.at(1, 0));
+  }
+
+  /// Concatenates this matrix with a a 2-by-1 matrix to form a 2-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> hcat(Matrix<Element, 2, 1> const & rhs) const {
+    return Matrix<Element, 2, 3>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-2 matrix to form a 2-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> hcat(Matrix<Element, 2, 2> const & rhs) const {
+    return Matrix<Element, 2, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 2-by-2 matrix by vertically concatenating a 1-by-2 matrix with a 1-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 2> const & upper, Matrix<Element, 1, 2> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1)
+      , lower.at(0, 0), lower.at(0, 1));
+  }
+
+  /// Concatenates this matrix with a a 1-by-2 matrix to form a 3-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> vcat(Matrix<Element, 1, 2> const & rhs) const {
+    return Matrix<Element, 3, 2>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-2 matrix to form a 4-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> vcat(Matrix<Element, 2, 2> const & rhs) const {
+    return Matrix<Element, 4, 2>::vcat(*this, rhs);
+  }
+
+  /// Forms a 2-by-2 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Element                         B,
+    Element                         C, Element                         D) {
+    return Matrix(
+      A, B
+      , C, D
+    );
+  }
+
+  /// Elementwise add operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+
+    result.data[2] = data[2] + rhs.data[2];
+    result.data[3] = data[3] + rhs.data[3];
+
+    return result;
+  }
+
+  /// Elementwise add operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+
+    data[2] += rhs.data[2];
+    data[3] += rhs.data[3];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+
+    result.data[2] = data[2] - rhs.data[2];
+    result.data[3] = data[3] - rhs.data[3];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+
+    data[2] -= rhs.data[2];
+    data[3] -= rhs.data[3];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+
+    result.data[2] = data[2] * rhs.data[2];
+    result.data[3] = data[3] * rhs.data[3];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+
+    result.data[2] = data[2] * s;
+    result.data[3] = data[3] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+
+    data[2] *= s;
+    data[3] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+
+    result.data[2] = data[2] / rhs.data[2];
+    result.data[3] = data[3] / rhs.data[3];
+
+    return result;
+  }
+
+  /// Scalar divide operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+
+    result.data[2] = data[2] / s;
+    result.data[3] = data[3] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+
+    data[2] /= s;
+    data[3] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (2-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+
+    data[2] /= rhs.data[2];
+    data[3] /= rhs.data[3];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+
+    return m;
+  }
+
+  /// Matrix product of size 2-by-1-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> product(
+    Matrix<Element, 2, 1> const &rhs,
+    Matrix<Element, 2, 1> accum = Matrix<Element, 2, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[2] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[3] * rhs.data[1];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-1-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> operator*(Matrix<Element, 2, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> product(
+    Matrix<Element, 2, 2> const &rhs,
+    Matrix<Element, 2, 2> accum = Matrix<Element, 2, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[2] * rhs.data[0];
+    accum.data[3] += data[2] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[3] * rhs.data[2];
+    accum.data[3] += data[3] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> operator*(Matrix<Element, 2, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 2, 2> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 2-by-3-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> product(
+    Matrix<Element, 2, 3> const &rhs,
+    Matrix<Element, 2, 3> accum = Matrix<Element, 2, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[2] * rhs.data[0];
+    accum.data[4] += data[2] * rhs.data[1];
+    accum.data[5] += data[2] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[3] * rhs.data[3];
+    accum.data[4] += data[3] * rhs.data[4];
+    accum.data[5] += data[3] * rhs.data[5];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-3-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> operator*(Matrix<Element, 2, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-4-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> product(
+    Matrix<Element, 2, 4> const &rhs,
+    Matrix<Element, 2, 4> accum = Matrix<Element, 2, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[2] * rhs.data[0];
+    accum.data[5] += data[2] * rhs.data[1];
+    accum.data[6] += data[2] * rhs.data[2];
+    accum.data[7] += data[2] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[3] * rhs.data[4];
+    accum.data[5] += data[3] * rhs.data[5];
+    accum.data[6] += data[3] * rhs.data[6];
+    accum.data[7] += data[3] * rhs.data[7];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-4-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> operator*(Matrix<Element, 2, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[3];
+
+    return accum;
+  }
+
+  /// Returns 2-by-2 rotation matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation(Element theta) {
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    return Matrix(
+      c, -s,
+      s,  c
+    );
+  }
+
+  /// Computes the determinant of a 2-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Element determinant(Element accum = Element()) const {
+        accum += data[0] * data[3] - data[1] * data[2];
+
+    return accum;
+  }
+
+  /// Computes the inverse of a 2-by-2 matrix given
+  /// the matrix's determinant
+  MUTLASS_HOST_DEVICE
+  Matrix inverse(Element det) const {
+    return Matrix(
+      data[3], -data[1],
+      -data[2], data[0]
+    ) * (Element(1) / det);
+  }
+
+  /// Computes the inverse of a 2-by-2 matrix.
+  MUTLASS_HOST_DEVICE
+  Matrix inverse() const {
+    return inverse(determinant());
+  }
+
+};
+
+/// Template alias for 2-by-2 matrix
+template <typename Element>
+using Matrix2x2 = Matrix<Element, 2, 2>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix2x2<Element> make_Matrix2x2(
+    Element _0_0, Element _0_1,
+    Element _1_0, Element _1_1
+) {
+  return Matrix2x2<Element>(
+  _0_0, _0_1,
+  _1_0, _1_1
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 2-by-3 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 2, 3> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 2;
+
+  /// Number of columns in matrix
+  static int const kColumns = 3;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 6;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 2-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 2-by-3 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1, Element _0_2,
+    Element _1_0, Element _1_1, Element _1_2
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;  data[2] = _0_2;
+    data[3] = _1_0;  data[4] = _1_1;  data[5] = _1_2;
+  }
+
+  /// Constucts a 2-by-3 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 3> const &row_0,
+    Matrix<Element, 1, 3> const &row_1
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_0.data[2];
+    data[3] = row_1.data[0];
+    data[4] = row_1.data[1];
+    data[5] = row_1.data[2];
+  }
+
+  /// Static method to construct a 2-by-3 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 3, 1> const &column_0,
+    Matrix<Element, 3, 1> const &column_1,
+    Matrix<Element, 3, 1> const &column_2
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_2.data[0];
+    result.data[3] = column_0.data[1];
+    result.data[4] = column_1.data[1];
+    result.data[5] = column_2.data[1];
+    return result;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+    m.data[4] = s;
+    m.data[5] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 2, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[3] = diag.data[1];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 2> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[3] = diag.data[1];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> diagonal() const {
+    Matrix<Element, 2, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[3];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> transpose() const {
+    Matrix<Element, 3, 2> mt;
+
+    mt.data[0] = data[0];
+    mt.data[2] = data[1];
+    mt.data[4] = data[2];
+    mt.data[1] = data[3];
+    mt.data[3] = data[4];
+    mt.data[5] = data[5];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 2 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 2 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> slice_1x3(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x3(Matrix<Element, 1, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> row(int i) const {
+    return slice_1x3(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 3> const &v, int i = 0) {
+    return set_slice_1x3(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 3] = m.data[1];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> column(int j) const {
+    return slice_2x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 2, 1> const &v, int j =0) {
+    return set_slice_2x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 3];
+    m.data[3] = data[i * 3 + j + 4];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 3] = m.data[2];
+    data[i * 3 + j + 4] = m.data[3];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> slice_2x3(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+    m.data[3] = data[i * 3 + j + 3];
+    m.data[4] = data[i * 3 + j + 4];
+    m.data[5] = data[i * 3 + j + 5];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x3(Matrix<Element, 2, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+    data[i * 3 + j + 3] = m.data[3];
+    data[i * 3 + j + 4] = m.data[4];
+    data[i * 3 + j + 5] = m.data[5];
+
+    return *this;
+  }
+
+  /// Forms a 2-by-3 matrix by horizontally concatenating a 2-by-1 matrix with a 2-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 2, 1> const & lhs, Matrix<Element, 2, 2> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1)
+      , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1));
+  }
+
+  /// Forms a 2-by-3 matrix by horizontally concatenating a 2-by-2 matrix with a 2-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 2, 2> const & lhs, Matrix<Element, 2, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0)
+      , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0));
+  }
+
+  /// Concatenates this matrix with a a 2-by-1 matrix to form a 2-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> hcat(Matrix<Element, 2, 1> const & rhs) const {
+    return Matrix<Element, 2, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 2-by-3 matrix by vertically concatenating a 1-by-3 matrix with a 1-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 3> const & upper, Matrix<Element, 1, 3> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2));
+  }
+
+  /// Concatenates this matrix with a a 1-by-3 matrix to form a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> vcat(Matrix<Element, 1, 3> const & rhs) const {
+    return Matrix<Element, 3, 3>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-3 matrix to form a 4-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> vcat(Matrix<Element, 2, 3> const & rhs) const {
+    return Matrix<Element, 4, 3>::vcat(*this, rhs);
+  }
+
+  /// Forms a 2-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Matrix<Element, 1, 2> const & B,
+    Element                         C, Matrix<Element, 1, 2> const & D) {
+    return Matrix(
+      A, B.at(0, 0), B.at(0, 1)
+      , C, D.at(0, 0), D.at(0, 1)
+    );
+  }
+
+  /// Forms a 2-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 2> const & A, Element                         B,
+    Matrix<Element, 1, 2> const & C, Element                         D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B
+      , C.at(0, 0), C.at(0, 1), D
+    );
+  }
+
+  /// Elementwise add operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+    result.data[2] = data[2] + rhs.data[2];
+
+    result.data[3] = data[3] + rhs.data[3];
+    result.data[4] = data[4] + rhs.data[4];
+    result.data[5] = data[5] + rhs.data[5];
+
+    return result;
+  }
+
+  /// Elementwise add operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+    data[2] += rhs.data[2];
+
+    data[3] += rhs.data[3];
+    data[4] += rhs.data[4];
+    data[5] += rhs.data[5];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+    result.data[2] = data[2] - rhs.data[2];
+
+    result.data[3] = data[3] - rhs.data[3];
+    result.data[4] = data[4] - rhs.data[4];
+    result.data[5] = data[5] - rhs.data[5];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+    data[2] -= rhs.data[2];
+
+    data[3] -= rhs.data[3];
+    data[4] -= rhs.data[4];
+    data[5] -= rhs.data[5];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+    result.data[2] = data[2] * rhs.data[2];
+
+    result.data[3] = data[3] * rhs.data[3];
+    result.data[4] = data[4] * rhs.data[4];
+    result.data[5] = data[5] * rhs.data[5];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+    result.data[2] = data[2] * s;
+
+    result.data[3] = data[3] * s;
+    result.data[4] = data[4] * s;
+    result.data[5] = data[5] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+    data[2] *= s;
+
+    data[3] *= s;
+    data[4] *= s;
+    data[5] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+    result.data[2] = data[2] / rhs.data[2];
+
+    result.data[3] = data[3] / rhs.data[3];
+    result.data[4] = data[4] / rhs.data[4];
+    result.data[5] = data[5] / rhs.data[5];
+
+    return result;
+  }
+
+  /// Scalar divide operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+    result.data[2] = data[2] / s;
+
+    result.data[3] = data[3] / s;
+    result.data[4] = data[4] / s;
+    result.data[5] = data[5] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+    data[2] /= s;
+
+    data[3] /= s;
+    data[4] /= s;
+    data[5] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (2-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+    data[2] /= rhs.data[2];
+
+    data[3] /= rhs.data[3];
+    data[4] /= rhs.data[4];
+    data[5] /= rhs.data[5];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+    m.data[4] = -m.data[4];
+    m.data[5] = -m.data[5];
+
+    return m;
+  }
+
+  /// Matrix product of size 2-by-1-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> product(
+    Matrix<Element, 3, 1> const &rhs,
+    Matrix<Element, 2, 1> accum = Matrix<Element, 2, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[3] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[4] * rhs.data[1];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[2];
+    accum.data[1] += data[5] * rhs.data[2];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-1-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> operator*(Matrix<Element, 3, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-2-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> product(
+    Matrix<Element, 3, 2> const &rhs,
+    Matrix<Element, 2, 2> accum = Matrix<Element, 2, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[3] * rhs.data[0];
+    accum.data[3] += data[3] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[4] * rhs.data[2];
+    accum.data[3] += data[4] * rhs.data[3];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[4];
+    accum.data[1] += data[2] * rhs.data[5];
+    accum.data[2] += data[5] * rhs.data[4];
+    accum.data[3] += data[5] * rhs.data[5];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-2-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> operator*(Matrix<Element, 3, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> product(
+    Matrix<Element, 3, 3> const &rhs,
+    Matrix<Element, 2, 3> accum = Matrix<Element, 2, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[3] * rhs.data[0];
+    accum.data[4] += data[3] * rhs.data[1];
+    accum.data[5] += data[3] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[4] * rhs.data[3];
+    accum.data[4] += data[4] * rhs.data[4];
+    accum.data[5] += data[4] * rhs.data[5];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[6];
+    accum.data[1] += data[2] * rhs.data[7];
+    accum.data[2] += data[2] * rhs.data[8];
+    accum.data[3] += data[5] * rhs.data[6];
+    accum.data[4] += data[5] * rhs.data[7];
+    accum.data[5] += data[5] * rhs.data[8];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> operator*(Matrix<Element, 3, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 3, 3> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 2-by-4-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> product(
+    Matrix<Element, 3, 4> const &rhs,
+    Matrix<Element, 2, 4> accum = Matrix<Element, 2, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[3] * rhs.data[0];
+    accum.data[5] += data[3] * rhs.data[1];
+    accum.data[6] += data[3] * rhs.data[2];
+    accum.data[7] += data[3] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[4] * rhs.data[4];
+    accum.data[5] += data[4] * rhs.data[5];
+    accum.data[6] += data[4] * rhs.data[6];
+    accum.data[7] += data[4] * rhs.data[7];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[8];
+    accum.data[1] += data[2] * rhs.data[9];
+    accum.data[2] += data[2] * rhs.data[10];
+    accum.data[3] += data[2] * rhs.data[11];
+    accum.data[4] += data[5] * rhs.data[8];
+    accum.data[5] += data[5] * rhs.data[9];
+    accum.data[6] += data[5] * rhs.data[10];
+    accum.data[7] += data[5] * rhs.data[11];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-4-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> operator*(Matrix<Element, 3, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+    accum += data[4];
+    accum += data[5];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+    accum += data[4] * data[4];
+    accum += data[5] * data[5];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[4];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 2-by-3 matrix
+template <typename Element>
+using Matrix2x3 = Matrix<Element, 2, 3>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix2x3<Element> make_Matrix2x3(
+    Element _0_0, Element _0_1, Element _0_2,
+    Element _1_0, Element _1_1, Element _1_2
+) {
+  return Matrix2x3<Element>(
+  _0_0, _0_1, _0_2,
+  _1_0, _1_1, _1_2
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 2-by-4 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 2, 4> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 2;
+
+  /// Number of columns in matrix
+  static int const kColumns = 4;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 8;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 2-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 2-by-4 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1, Element _0_2, Element _0_3,
+    Element _1_0, Element _1_1, Element _1_2, Element _1_3
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;  data[2] = _0_2;  data[3] = _0_3;
+    data[4] = _1_0;  data[5] = _1_1;  data[6] = _1_2;  data[7] = _1_3;
+  }
+
+  /// Constucts a 2-by-4 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 4> const &row_0,
+    Matrix<Element, 1, 4> const &row_1
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_0.data[2];
+    data[3] = row_0.data[3];
+    data[4] = row_1.data[0];
+    data[5] = row_1.data[1];
+    data[6] = row_1.data[2];
+    data[7] = row_1.data[3];
+  }
+
+  /// Static method to construct a 2-by-4 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 4, 1> const &column_0,
+    Matrix<Element, 4, 1> const &column_1,
+    Matrix<Element, 4, 1> const &column_2,
+    Matrix<Element, 4, 1> const &column_3
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_2.data[0];
+    result.data[3] = column_3.data[0];
+    result.data[4] = column_0.data[1];
+    result.data[5] = column_1.data[1];
+    result.data[6] = column_2.data[1];
+    result.data[7] = column_3.data[1];
+    return result;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+    m.data[4] = s;
+    m.data[5] = s;
+    m.data[6] = s;
+    m.data[7] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 2, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[3] = diag.data[1];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 2> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[3] = diag.data[1];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> diagonal() const {
+    Matrix<Element, 2, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[3];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> transpose() const {
+    Matrix<Element, 4, 2> mt;
+
+    mt.data[0] = data[0];
+    mt.data[2] = data[1];
+    mt.data[4] = data[2];
+    mt.data[6] = data[3];
+    mt.data[1] = data[4];
+    mt.data[3] = data[5];
+    mt.data[5] = data[6];
+    mt.data[7] = data[7];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 2 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 2 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> slice_1x3(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x3(Matrix<Element, 1, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> slice_1x4(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x4(Matrix<Element, 1, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> row(int i) const {
+    return slice_1x4(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 4> const &v, int i = 0) {
+    return set_slice_1x4(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 4];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 4] = m.data[1];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> column(int j) const {
+    return slice_2x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 2, 1> const &v, int j =0) {
+    return set_slice_2x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 4];
+    m.data[3] = data[i * 4 + j + 5];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 4] = m.data[2];
+    data[i * 4 + j + 5] = m.data[3];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> slice_2x3(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 4];
+    m.data[4] = data[i * 4 + j + 5];
+    m.data[5] = data[i * 4 + j + 6];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x3(Matrix<Element, 2, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 4] = m.data[3];
+    data[i * 4 + j + 5] = m.data[4];
+    data[i * 4 + j + 6] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> slice_2x4(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+    m.data[4] = data[i * 4 + j + 4];
+    m.data[5] = data[i * 4 + j + 5];
+    m.data[6] = data[i * 4 + j + 6];
+    m.data[7] = data[i * 4 + j + 7];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x4(Matrix<Element, 2, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+    data[i * 4 + j + 4] = m.data[4];
+    data[i * 4 + j + 5] = m.data[5];
+    data[i * 4 + j + 6] = m.data[6];
+    data[i * 4 + j + 7] = m.data[7];
+
+    return *this;
+  }
+
+  /// Forms a 2-by-4 matrix by horizontally concatenating a 2-by-1 matrix with a 2-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 2, 1> const & lhs, Matrix<Element, 2, 3> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1), rhs.at(0, 2)
+      , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1), rhs.at(1, 2));
+  }
+
+  /// Forms a 2-by-4 matrix by horizontally concatenating a 2-by-2 matrix with a 2-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 2, 2> const & lhs, Matrix<Element, 2, 2> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0), rhs.at(0, 1)
+      , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0), rhs.at(1, 1));
+  }
+
+  /// Forms a 2-by-4 matrix by horizontally concatenating a 2-by-3 matrix with a 2-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 2, 3> const & lhs, Matrix<Element, 2, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), lhs.at(0, 2), rhs.at(0, 0)
+      , lhs.at(1, 0), lhs.at(1, 1), lhs.at(1, 2), rhs.at(1, 0));
+  }
+
+  /// Forms a 2-by-4 matrix by vertically concatenating a 1-by-4 matrix with a 1-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 4> const & upper, Matrix<Element, 1, 4> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3));
+  }
+
+  /// Concatenates this matrix with a a 1-by-4 matrix to form a 3-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> vcat(Matrix<Element, 1, 4> const & rhs) const {
+    return Matrix<Element, 3, 4>::vcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 2-by-4 matrix to form a 4-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> vcat(Matrix<Element, 2, 4> const & rhs) const {
+    return Matrix<Element, 4, 4>::vcat(*this, rhs);
+  }
+
+  /// Forms a 2-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Matrix<Element, 1, 3> const & B,
+    Element                         C, Matrix<Element, 1, 3> const & D) {
+    return Matrix(
+      A, B.at(0, 0), B.at(0, 1), B.at(0, 2)
+      , C, D.at(0, 0), D.at(0, 1), D.at(0, 2)
+    );
+  }
+
+  /// Forms a 2-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 2> const & A, Matrix<Element, 1, 2> const & B,
+    Matrix<Element, 1, 2> const & C, Matrix<Element, 1, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1)
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1)
+    );
+  }
+
+  /// Forms a 2-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 3> const & A, Element                         B,
+    Matrix<Element, 1, 3> const & C, Element                         D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), A.at(0, 2), B
+      , C.at(0, 0), C.at(0, 1), C.at(0, 2), D
+    );
+  }
+
+  /// Elementwise add operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+    result.data[2] = data[2] + rhs.data[2];
+    result.data[3] = data[3] + rhs.data[3];
+
+    result.data[4] = data[4] + rhs.data[4];
+    result.data[5] = data[5] + rhs.data[5];
+    result.data[6] = data[6] + rhs.data[6];
+    result.data[7] = data[7] + rhs.data[7];
+
+    return result;
+  }
+
+  /// Elementwise add operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+    data[2] += rhs.data[2];
+    data[3] += rhs.data[3];
+
+    data[4] += rhs.data[4];
+    data[5] += rhs.data[5];
+    data[6] += rhs.data[6];
+    data[7] += rhs.data[7];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+    result.data[2] = data[2] - rhs.data[2];
+    result.data[3] = data[3] - rhs.data[3];
+
+    result.data[4] = data[4] - rhs.data[4];
+    result.data[5] = data[5] - rhs.data[5];
+    result.data[6] = data[6] - rhs.data[6];
+    result.data[7] = data[7] - rhs.data[7];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+    data[2] -= rhs.data[2];
+    data[3] -= rhs.data[3];
+
+    data[4] -= rhs.data[4];
+    data[5] -= rhs.data[5];
+    data[6] -= rhs.data[6];
+    data[7] -= rhs.data[7];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+    result.data[2] = data[2] * rhs.data[2];
+    result.data[3] = data[3] * rhs.data[3];
+
+    result.data[4] = data[4] * rhs.data[4];
+    result.data[5] = data[5] * rhs.data[5];
+    result.data[6] = data[6] * rhs.data[6];
+    result.data[7] = data[7] * rhs.data[7];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+    result.data[2] = data[2] * s;
+    result.data[3] = data[3] * s;
+
+    result.data[4] = data[4] * s;
+    result.data[5] = data[5] * s;
+    result.data[6] = data[6] * s;
+    result.data[7] = data[7] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+    data[2] *= s;
+    data[3] *= s;
+
+    data[4] *= s;
+    data[5] *= s;
+    data[6] *= s;
+    data[7] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+    result.data[2] = data[2] / rhs.data[2];
+    result.data[3] = data[3] / rhs.data[3];
+
+    result.data[4] = data[4] / rhs.data[4];
+    result.data[5] = data[5] / rhs.data[5];
+    result.data[6] = data[6] / rhs.data[6];
+    result.data[7] = data[7] / rhs.data[7];
+
+    return result;
+  }
+
+  /// Scalar divide operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+    result.data[2] = data[2] / s;
+    result.data[3] = data[3] / s;
+
+    result.data[4] = data[4] / s;
+    result.data[5] = data[5] / s;
+    result.data[6] = data[6] / s;
+    result.data[7] = data[7] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+    data[2] /= s;
+    data[3] /= s;
+
+    data[4] /= s;
+    data[5] /= s;
+    data[6] /= s;
+    data[7] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (2-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+    data[2] /= rhs.data[2];
+    data[3] /= rhs.data[3];
+
+    data[4] /= rhs.data[4];
+    data[5] /= rhs.data[5];
+    data[6] /= rhs.data[6];
+    data[7] /= rhs.data[7];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+    m.data[4] = -m.data[4];
+    m.data[5] = -m.data[5];
+    m.data[6] = -m.data[6];
+    m.data[7] = -m.data[7];
+
+    return m;
+  }
+
+  /// Matrix product of size 2-by-1-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> product(
+    Matrix<Element, 4, 1> const &rhs,
+    Matrix<Element, 2, 1> accum = Matrix<Element, 2, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[4] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[5] * rhs.data[1];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[2];
+    accum.data[1] += data[6] * rhs.data[2];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[3];
+    accum.data[1] += data[7] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-1-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> operator*(Matrix<Element, 4, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-2-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> product(
+    Matrix<Element, 4, 2> const &rhs,
+    Matrix<Element, 2, 2> accum = Matrix<Element, 2, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[4] * rhs.data[0];
+    accum.data[3] += data[4] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[5] * rhs.data[2];
+    accum.data[3] += data[5] * rhs.data[3];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[4];
+    accum.data[1] += data[2] * rhs.data[5];
+    accum.data[2] += data[6] * rhs.data[4];
+    accum.data[3] += data[6] * rhs.data[5];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[6];
+    accum.data[1] += data[3] * rhs.data[7];
+    accum.data[2] += data[7] * rhs.data[6];
+    accum.data[3] += data[7] * rhs.data[7];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-2-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> operator*(Matrix<Element, 4, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-3-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> product(
+    Matrix<Element, 4, 3> const &rhs,
+    Matrix<Element, 2, 3> accum = Matrix<Element, 2, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[4] * rhs.data[0];
+    accum.data[4] += data[4] * rhs.data[1];
+    accum.data[5] += data[4] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[5] * rhs.data[3];
+    accum.data[4] += data[5] * rhs.data[4];
+    accum.data[5] += data[5] * rhs.data[5];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[6];
+    accum.data[1] += data[2] * rhs.data[7];
+    accum.data[2] += data[2] * rhs.data[8];
+    accum.data[3] += data[6] * rhs.data[6];
+    accum.data[4] += data[6] * rhs.data[7];
+    accum.data[5] += data[6] * rhs.data[8];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[9];
+    accum.data[1] += data[3] * rhs.data[10];
+    accum.data[2] += data[3] * rhs.data[11];
+    accum.data[3] += data[7] * rhs.data[9];
+    accum.data[4] += data[7] * rhs.data[10];
+    accum.data[5] += data[7] * rhs.data[11];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-3-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> operator*(Matrix<Element, 4, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> product(
+    Matrix<Element, 4, 4> const &rhs,
+    Matrix<Element, 2, 4> accum = Matrix<Element, 2, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[4] * rhs.data[0];
+    accum.data[5] += data[4] * rhs.data[1];
+    accum.data[6] += data[4] * rhs.data[2];
+    accum.data[7] += data[4] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[5] * rhs.data[4];
+    accum.data[5] += data[5] * rhs.data[5];
+    accum.data[6] += data[5] * rhs.data[6];
+    accum.data[7] += data[5] * rhs.data[7];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[8];
+    accum.data[1] += data[2] * rhs.data[9];
+    accum.data[2] += data[2] * rhs.data[10];
+    accum.data[3] += data[2] * rhs.data[11];
+    accum.data[4] += data[6] * rhs.data[8];
+    accum.data[5] += data[6] * rhs.data[9];
+    accum.data[6] += data[6] * rhs.data[10];
+    accum.data[7] += data[6] * rhs.data[11];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[12];
+    accum.data[1] += data[3] * rhs.data[13];
+    accum.data[2] += data[3] * rhs.data[14];
+    accum.data[3] += data[3] * rhs.data[15];
+    accum.data[4] += data[7] * rhs.data[12];
+    accum.data[5] += data[7] * rhs.data[13];
+    accum.data[6] += data[7] * rhs.data[14];
+    accum.data[7] += data[7] * rhs.data[15];
+
+    return accum;
+  }
+
+  /// Matrix product of size 2-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> operator*(Matrix<Element, 4, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 2-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 4, 4> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+    accum += data[4];
+    accum += data[5];
+    accum += data[6];
+    accum += data[7];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+    accum += data[4] * data[4];
+    accum += data[5] * data[5];
+    accum += data[6] * data[6];
+    accum += data[7] * data[7];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[5];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 2-by-4 matrix
+template <typename Element>
+using Matrix2x4 = Matrix<Element, 2, 4>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix2x4<Element> make_Matrix2x4(
+    Element _0_0, Element _0_1, Element _0_2, Element _0_3,
+    Element _1_0, Element _1_1, Element _1_2, Element _1_3
+) {
+  return Matrix2x4<Element>(
+  _0_0, _0_1, _0_2, _0_3,
+  _1_0, _1_1, _1_2, _1_3
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 3-by-1 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 3, 1> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 3;
+
+  /// Number of columns in matrix
+  static int const kColumns = 1;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 3;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 3-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 3-by-1 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0,
+    Element _1_0,
+    Element _2_0
+  ) {
+
+    data[0] = _0_0;
+    data[1] = _1_0;
+    data[2] = _2_0;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> transpose() const {
+    Matrix<Element, 1, 3> mt;
+
+    mt.data[0] = data[0];
+    mt.data[1] = data[1];
+    mt.data[2] = data[2];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 3 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 3 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 1 + j + 0];
+    m.data[1] = data[i * 1 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 1 + j + 0] = m.data[0];
+    data[i * 1 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> slice_3x1(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 1> m;
+
+    m.data[0] = data[i * 1 + j + 0];
+    m.data[1] = data[i * 1 + j + 1];
+    m.data[2] = data[i * 1 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x1(Matrix<Element, 3, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 1 + j + 0] = m.data[0];
+    data[i * 1 + j + 1] = m.data[1];
+    data[i * 1 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> column(int j) const {
+    return slice_3x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 3, 1> const &v, int j =0) {
+    return set_slice_3x1(v, 0, j);
+  }
+
+  /// Concatenates this matrix with a a 3-by-1 matrix to form a 3-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> hcat(Matrix<Element, 3, 1> const & rhs) const {
+    return Matrix<Element, 3, 2>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 3-by-2 matrix to form a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> hcat(Matrix<Element, 3, 2> const & rhs) const {
+    return Matrix<Element, 3, 3>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 3-by-3 matrix to form a 3-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> hcat(Matrix<Element, 3, 3> const & rhs) const {
+    return Matrix<Element, 3, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 3-by-1 matrix by vertically concatenating an Element with a 2-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Element upper, Matrix<Element, 2, 1> const & lower) {
+    return Matrix(
+      upper
+      , lower.at(0, 0)
+      , lower.at(1, 0));
+  }
+
+  /// Forms a 3-by-1 matrix by vertically concatenating a 2-by-1 matrix with an Element
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 2, 1> const & upper, Element lower) {
+    return Matrix(
+      upper.at(0, 0)
+      , upper.at(1, 0)
+      , lower);
+  }
+
+  /// Concatenates this matrix with a an Element to form a 4-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> vcat(Element rhs) const {
+    return Matrix<Element, 4, 1>::vcat(*this, rhs);
+  }
+
+  /// Elementwise add operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+
+    result.data[1] = data[1] + rhs.data[1];
+
+    result.data[2] = data[2] + rhs.data[2];
+
+    return result;
+  }
+
+  /// Elementwise add operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+
+    data[1] += rhs.data[1];
+
+    data[2] += rhs.data[2];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+
+    result.data[1] = data[1] - rhs.data[1];
+
+    result.data[2] = data[2] - rhs.data[2];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+
+    data[1] -= rhs.data[1];
+
+    data[2] -= rhs.data[2];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+
+    result.data[1] = data[1] * rhs.data[1];
+
+    result.data[2] = data[2] * rhs.data[2];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+
+    result.data[1] = data[1] * s;
+
+    result.data[2] = data[2] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+
+    data[1] *= s;
+
+    data[2] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+
+    result.data[1] = data[1] / rhs.data[1];
+
+    result.data[2] = data[2] / rhs.data[2];
+
+    return result;
+  }
+
+  /// Scalar divide operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+
+    result.data[1] = data[1] / s;
+
+    result.data[2] = data[2] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+
+    data[1] /= s;
+
+    data[2] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (3-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+
+    data[1] /= rhs.data[1];
+
+    data[2] /= rhs.data[2];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+
+    return m;
+  }
+
+  /// Matrix product of size 3-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> product(
+    Matrix<Element, 1, 1> const &rhs,
+    Matrix<Element, 3, 1> accum = Matrix<Element, 3, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[1] * rhs.data[0];
+    accum.data[2] += data[2] * rhs.data[0];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> operator*(Matrix<Element, 1, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 1, 1> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 3-by-2-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> product(
+    Matrix<Element, 1, 2> const &rhs,
+    Matrix<Element, 3, 2> accum = Matrix<Element, 3, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[1] * rhs.data[0];
+    accum.data[3] += data[1] * rhs.data[1];
+    accum.data[4] += data[2] * rhs.data[0];
+    accum.data[5] += data[2] * rhs.data[1];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-2-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> operator*(Matrix<Element, 1, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-3-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> product(
+    Matrix<Element, 1, 3> const &rhs,
+    Matrix<Element, 3, 3> accum = Matrix<Element, 3, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[1] * rhs.data[0];
+    accum.data[4] += data[1] * rhs.data[1];
+    accum.data[5] += data[1] * rhs.data[2];
+    accum.data[6] += data[2] * rhs.data[0];
+    accum.data[7] += data[2] * rhs.data[1];
+    accum.data[8] += data[2] * rhs.data[2];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-3-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> operator*(Matrix<Element, 1, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-4-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> product(
+    Matrix<Element, 1, 4> const &rhs,
+    Matrix<Element, 3, 4> accum = Matrix<Element, 3, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[1] * rhs.data[0];
+    accum.data[5] += data[1] * rhs.data[1];
+    accum.data[6] += data[1] * rhs.data[2];
+    accum.data[7] += data[1] * rhs.data[3];
+    accum.data[8] += data[2] * rhs.data[0];
+    accum.data[9] += data[2] * rhs.data[1];
+    accum.data[10] += data[2] * rhs.data[2];
+    accum.data[11] += data[2] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-4-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> operator*(Matrix<Element, 1, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Dot product of vectors with extent 3
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 3, 1> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    accum += data[2] * rhs.data[2];
+    return accum;
+  }
+
+  /// Dot product of vectors with extent 3
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 1, 3> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    accum += data[2] * rhs.data[2];
+    return accum;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+
+    return accum;
+  }
+
+  /// Cross product
+  MUTLASS_HOST_DEVICE
+  Matrix cross(Matrix const &rhs) const {
+    return Matrix(
+      data[1] * rhs.data[2] - data[2] * rhs.data[1],
+      data[0] * rhs.data[2] - data[2] * rhs.data[1],
+      data[0] * rhs.data[1] - data[1] * rhs.data[0]
+    );
+  }
+
+};
+
+/// Template alias for 3-by-1 matrix
+template <typename Element>
+using Matrix3x1 = Matrix<Element, 3, 1>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix3x1<Element> make_Matrix3x1(
+    Element _0_0,
+    Element _1_0,
+    Element _2_0
+) {
+  return Matrix3x1<Element>(
+  _0_0,
+  _1_0,
+  _2_0
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 3-by-2 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 3, 2> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 3;
+
+  /// Number of columns in matrix
+  static int const kColumns = 2;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 6;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 3-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 3-by-2 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1,
+    Element _1_0, Element _1_1,
+    Element _2_0, Element _2_1
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;
+    data[2] = _1_0;  data[3] = _1_1;
+    data[4] = _2_0;  data[5] = _2_1;
+  }
+
+  /// Constucts a 3-by-2 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 2> const &row_0,
+    Matrix<Element, 1, 2> const &row_1,
+    Matrix<Element, 1, 2> const &row_2
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_1.data[0];
+    data[3] = row_1.data[1];
+    data[4] = row_2.data[0];
+    data[5] = row_2.data[1];
+  }
+
+  /// Static method to construct a 3-by-2 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 2, 1> const &column_0,
+    Matrix<Element, 2, 1> const &column_1
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_0.data[1];
+    result.data[3] = column_1.data[1];
+    result.data[4] = column_0.data[2];
+    result.data[5] = column_1.data[2];
+    return result;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+    m.data[4] = s;
+    m.data[5] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 2, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[4] = diag.data[1];
+    m.data[8] = diag.data[2];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 2> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[4] = diag.data[1];
+    m.data[8] = diag.data[2];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> diagonal() const {
+    Matrix<Element, 2, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[4];
+    diag.data[2] = data[8];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> transpose() const {
+    Matrix<Element, 2, 3> mt;
+
+    mt.data[0] = data[0];
+    mt.data[3] = data[1];
+    mt.data[1] = data[2];
+    mt.data[4] = data[3];
+    mt.data[2] = data[4];
+    mt.data[5] = data[5];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 3 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 3 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> row(int i) const {
+    return slice_1x2(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 2> const &v, int i = 0) {
+    return set_slice_1x2(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 2] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+    m.data[2] = data[i * 2 + j + 2];
+    m.data[3] = data[i * 2 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+    data[i * 2 + j + 2] = m.data[2];
+    data[i * 2 + j + 3] = m.data[3];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> slice_3x1(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 1> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 2];
+    m.data[2] = data[i * 2 + j + 4];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x1(Matrix<Element, 3, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 2] = m.data[1];
+    data[i * 2 + j + 4] = m.data[2];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> column(int j) const {
+    return slice_3x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 3, 1> const &v, int j =0) {
+    return set_slice_3x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> slice_3x2(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+    m.data[2] = data[i * 2 + j + 2];
+    m.data[3] = data[i * 2 + j + 3];
+    m.data[4] = data[i * 2 + j + 4];
+    m.data[5] = data[i * 2 + j + 5];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x2(Matrix<Element, 3, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+    data[i * 2 + j + 2] = m.data[2];
+    data[i * 2 + j + 3] = m.data[3];
+    data[i * 2 + j + 4] = m.data[4];
+    data[i * 2 + j + 5] = m.data[5];
+
+    return *this;
+  }
+
+  /// Forms a 3-by-2 matrix by horizontally concatenating a 3-by-1 matrix with a 3-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 3, 1> const & lhs, Matrix<Element, 3, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0)
+      , lhs.at(1, 0), rhs.at(1, 0)
+      , lhs.at(2, 0), rhs.at(2, 0));
+  }
+
+  /// Concatenates this matrix with a a 3-by-1 matrix to form a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> hcat(Matrix<Element, 3, 1> const & rhs) const {
+    return Matrix<Element, 3, 3>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 3-by-2 matrix to form a 3-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> hcat(Matrix<Element, 3, 2> const & rhs) const {
+    return Matrix<Element, 3, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 3-by-2 matrix by vertically concatenating a 1-by-2 matrix with a 2-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 2> const & upper, Matrix<Element, 2, 2> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1)
+      , lower.at(0, 0), lower.at(0, 1)
+      , lower.at(1, 0), lower.at(1, 1));
+  }
+
+  /// Forms a 3-by-2 matrix by vertically concatenating a 2-by-2 matrix with a 1-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 2, 2> const & upper, Matrix<Element, 1, 2> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1)
+      , upper.at(1, 0), upper.at(1, 1)
+      , lower.at(0, 0), lower.at(0, 1));
+  }
+
+  /// Concatenates this matrix with a a 1-by-2 matrix to form a 4-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> vcat(Matrix<Element, 1, 2> const & rhs) const {
+    return Matrix<Element, 4, 2>::vcat(*this, rhs);
+  }
+
+  /// Forms a 3-by-2 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Element                         B,
+    Matrix<Element, 2, 1> const & C, Matrix<Element, 2, 1> const & D) {
+    return Matrix(
+      A, B
+      , C.at(0, 0), D.at(0, 0)
+      , C.at(1, 0), D.at(1, 0)
+    );
+  }
+
+  /// Forms a 3-by-2 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 1> const & A, Matrix<Element, 2, 1> const & B,
+    Element                         C, Element                         D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0)
+      , A.at(1, 0), B.at(1, 0)
+      , C, D
+    );
+  }
+
+  /// Elementwise add operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+
+    result.data[2] = data[2] + rhs.data[2];
+    result.data[3] = data[3] + rhs.data[3];
+
+    result.data[4] = data[4] + rhs.data[4];
+    result.data[5] = data[5] + rhs.data[5];
+
+    return result;
+  }
+
+  /// Elementwise add operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+
+    data[2] += rhs.data[2];
+    data[3] += rhs.data[3];
+
+    data[4] += rhs.data[4];
+    data[5] += rhs.data[5];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+
+    result.data[2] = data[2] - rhs.data[2];
+    result.data[3] = data[3] - rhs.data[3];
+
+    result.data[4] = data[4] - rhs.data[4];
+    result.data[5] = data[5] - rhs.data[5];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+
+    data[2] -= rhs.data[2];
+    data[3] -= rhs.data[3];
+
+    data[4] -= rhs.data[4];
+    data[5] -= rhs.data[5];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+
+    result.data[2] = data[2] * rhs.data[2];
+    result.data[3] = data[3] * rhs.data[3];
+
+    result.data[4] = data[4] * rhs.data[4];
+    result.data[5] = data[5] * rhs.data[5];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+
+    result.data[2] = data[2] * s;
+    result.data[3] = data[3] * s;
+
+    result.data[4] = data[4] * s;
+    result.data[5] = data[5] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+
+    data[2] *= s;
+    data[3] *= s;
+
+    data[4] *= s;
+    data[5] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+
+    result.data[2] = data[2] / rhs.data[2];
+    result.data[3] = data[3] / rhs.data[3];
+
+    result.data[4] = data[4] / rhs.data[4];
+    result.data[5] = data[5] / rhs.data[5];
+
+    return result;
+  }
+
+  /// Scalar divide operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+
+    result.data[2] = data[2] / s;
+    result.data[3] = data[3] / s;
+
+    result.data[4] = data[4] / s;
+    result.data[5] = data[5] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+
+    data[2] /= s;
+    data[3] /= s;
+
+    data[4] /= s;
+    data[5] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (3-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+
+    data[2] /= rhs.data[2];
+    data[3] /= rhs.data[3];
+
+    data[4] /= rhs.data[4];
+    data[5] /= rhs.data[5];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+    m.data[4] = -m.data[4];
+    m.data[5] = -m.data[5];
+
+    return m;
+  }
+
+  /// Matrix product of size 3-by-1-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> product(
+    Matrix<Element, 2, 1> const &rhs,
+    Matrix<Element, 3, 1> accum = Matrix<Element, 3, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[2] * rhs.data[0];
+    accum.data[2] += data[4] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[3] * rhs.data[1];
+    accum.data[2] += data[5] * rhs.data[1];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-1-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> operator*(Matrix<Element, 2, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> product(
+    Matrix<Element, 2, 2> const &rhs,
+    Matrix<Element, 3, 2> accum = Matrix<Element, 3, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[2] * rhs.data[0];
+    accum.data[3] += data[2] * rhs.data[1];
+    accum.data[4] += data[4] * rhs.data[0];
+    accum.data[5] += data[4] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[3] * rhs.data[2];
+    accum.data[3] += data[3] * rhs.data[3];
+    accum.data[4] += data[5] * rhs.data[2];
+    accum.data[5] += data[5] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> operator*(Matrix<Element, 2, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 2, 2> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 3-by-3-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> product(
+    Matrix<Element, 2, 3> const &rhs,
+    Matrix<Element, 3, 3> accum = Matrix<Element, 3, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[2] * rhs.data[0];
+    accum.data[4] += data[2] * rhs.data[1];
+    accum.data[5] += data[2] * rhs.data[2];
+    accum.data[6] += data[4] * rhs.data[0];
+    accum.data[7] += data[4] * rhs.data[1];
+    accum.data[8] += data[4] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[3] * rhs.data[3];
+    accum.data[4] += data[3] * rhs.data[4];
+    accum.data[5] += data[3] * rhs.data[5];
+    accum.data[6] += data[5] * rhs.data[3];
+    accum.data[7] += data[5] * rhs.data[4];
+    accum.data[8] += data[5] * rhs.data[5];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-3-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> operator*(Matrix<Element, 2, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-4-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> product(
+    Matrix<Element, 2, 4> const &rhs,
+    Matrix<Element, 3, 4> accum = Matrix<Element, 3, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[2] * rhs.data[0];
+    accum.data[5] += data[2] * rhs.data[1];
+    accum.data[6] += data[2] * rhs.data[2];
+    accum.data[7] += data[2] * rhs.data[3];
+    accum.data[8] += data[4] * rhs.data[0];
+    accum.data[9] += data[4] * rhs.data[1];
+    accum.data[10] += data[4] * rhs.data[2];
+    accum.data[11] += data[4] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[3] * rhs.data[4];
+    accum.data[5] += data[3] * rhs.data[5];
+    accum.data[6] += data[3] * rhs.data[6];
+    accum.data[7] += data[3] * rhs.data[7];
+    accum.data[8] += data[5] * rhs.data[4];
+    accum.data[9] += data[5] * rhs.data[5];
+    accum.data[10] += data[5] * rhs.data[6];
+    accum.data[11] += data[5] * rhs.data[7];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-4-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> operator*(Matrix<Element, 2, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+    accum += data[4];
+    accum += data[5];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+    accum += data[4] * data[4];
+    accum += data[5] * data[5];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[3];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 3-by-2 matrix
+template <typename Element>
+using Matrix3x2 = Matrix<Element, 3, 2>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix3x2<Element> make_Matrix3x2(
+    Element _0_0, Element _0_1,
+    Element _1_0, Element _1_1,
+    Element _2_0, Element _2_1
+) {
+  return Matrix3x2<Element>(
+  _0_0, _0_1,
+  _1_0, _1_1,
+  _2_0, _2_1
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 3-by-3 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 3, 3> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 3;
+
+  /// Number of columns in matrix
+  static int const kColumns = 3;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 9;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 3-by-3 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1, Element _0_2,
+    Element _1_0, Element _1_1, Element _1_2,
+    Element _2_0, Element _2_1, Element _2_2
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;  data[2] = _0_2;
+    data[3] = _1_0;  data[4] = _1_1;  data[5] = _1_2;
+    data[6] = _2_0;  data[7] = _2_1;  data[8] = _2_2;
+  }
+
+  /// Constucts a 3-by-3 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 3> const &row_0,
+    Matrix<Element, 1, 3> const &row_1,
+    Matrix<Element, 1, 3> const &row_2
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_0.data[2];
+    data[3] = row_1.data[0];
+    data[4] = row_1.data[1];
+    data[5] = row_1.data[2];
+    data[6] = row_2.data[0];
+    data[7] = row_2.data[1];
+    data[8] = row_2.data[2];
+  }
+
+  /// Static method to construct a 3-by-3 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 3, 1> const &column_0,
+    Matrix<Element, 3, 1> const &column_1,
+    Matrix<Element, 3, 1> const &column_2
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_2.data[0];
+    result.data[3] = column_0.data[1];
+    result.data[4] = column_1.data[1];
+    result.data[5] = column_2.data[1];
+    result.data[6] = column_0.data[2];
+    result.data[7] = column_1.data[2];
+    result.data[8] = column_2.data[2];
+    return result;
+  }
+
+  /// Constructs an identity matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix identity() {
+    Matrix m;
+
+    m.data[0] = Element(1);
+    m.data[4] = Element(1);
+    m.data[8] = Element(1);
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+    m.data[4] = s;
+    m.data[5] = s;
+    m.data[6] = s;
+    m.data[7] = s;
+    m.data[8] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 3, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[4] = diag.data[1];
+    m.data[8] = diag.data[2];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 3> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[4] = diag.data[1];
+    m.data[8] = diag.data[2];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> diagonal() const {
+    Matrix<Element, 3, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[4];
+    diag.data[2] = data[8];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> transpose() const {
+    Matrix<Element, 3, 3> mt;
+
+    mt.data[0] = data[0];
+    mt.data[3] = data[1];
+    mt.data[6] = data[2];
+    mt.data[1] = data[3];
+    mt.data[4] = data[4];
+    mt.data[7] = data[5];
+    mt.data[2] = data[6];
+    mt.data[5] = data[7];
+    mt.data[8] = data[8];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 3 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 3 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> slice_1x3(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x3(Matrix<Element, 1, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> row(int i) const {
+    return slice_1x3(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 3> const &v, int i = 0) {
+    return set_slice_1x3(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 3] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 3];
+    m.data[3] = data[i * 3 + j + 4];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 3] = m.data[2];
+    data[i * 3 + j + 4] = m.data[3];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> slice_2x3(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+    m.data[3] = data[i * 3 + j + 3];
+    m.data[4] = data[i * 3 + j + 4];
+    m.data[5] = data[i * 3 + j + 5];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x3(Matrix<Element, 2, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+    data[i * 3 + j + 3] = m.data[3];
+    data[i * 3 + j + 4] = m.data[4];
+    data[i * 3 + j + 5] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> slice_3x1(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 1> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 3];
+    m.data[2] = data[i * 3 + j + 6];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x1(Matrix<Element, 3, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 3] = m.data[1];
+    data[i * 3 + j + 6] = m.data[2];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> column(int j) const {
+    return slice_3x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 3, 1> const &v, int j =0) {
+    return set_slice_3x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> slice_3x2(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 3];
+    m.data[3] = data[i * 3 + j + 4];
+    m.data[4] = data[i * 3 + j + 6];
+    m.data[5] = data[i * 3 + j + 7];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x2(Matrix<Element, 3, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 3] = m.data[2];
+    data[i * 3 + j + 4] = m.data[3];
+    data[i * 3 + j + 6] = m.data[4];
+    data[i * 3 + j + 7] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> slice_3x3(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+    m.data[3] = data[i * 3 + j + 3];
+    m.data[4] = data[i * 3 + j + 4];
+    m.data[5] = data[i * 3 + j + 5];
+    m.data[6] = data[i * 3 + j + 6];
+    m.data[7] = data[i * 3 + j + 7];
+    m.data[8] = data[i * 3 + j + 8];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x3(Matrix<Element, 3, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+    data[i * 3 + j + 3] = m.data[3];
+    data[i * 3 + j + 4] = m.data[4];
+    data[i * 3 + j + 5] = m.data[5];
+    data[i * 3 + j + 6] = m.data[6];
+    data[i * 3 + j + 7] = m.data[7];
+    data[i * 3 + j + 8] = m.data[8];
+
+    return *this;
+  }
+
+  /// Forms a 3-by-3 matrix by horizontally concatenating a 3-by-1 matrix with a 3-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 3, 1> const & lhs, Matrix<Element, 3, 2> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1)
+      , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1)
+      , lhs.at(2, 0), rhs.at(2, 0), rhs.at(2, 1));
+  }
+
+  /// Forms a 3-by-3 matrix by horizontally concatenating a 3-by-2 matrix with a 3-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 3, 2> const & lhs, Matrix<Element, 3, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0)
+      , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0)
+      , lhs.at(2, 0), lhs.at(2, 1), rhs.at(2, 0));
+  }
+
+  /// Concatenates this matrix with a a 3-by-1 matrix to form a 3-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> hcat(Matrix<Element, 3, 1> const & rhs) const {
+    return Matrix<Element, 3, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 3-by-3 matrix by vertically concatenating a 1-by-3 matrix with a 2-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 3> const & upper, Matrix<Element, 2, 3> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2)
+      , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2));
+  }
+
+  /// Forms a 3-by-3 matrix by vertically concatenating a 2-by-3 matrix with a 1-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 2, 3> const & upper, Matrix<Element, 1, 3> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2)
+      , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2));
+  }
+
+  /// Concatenates this matrix with a a 1-by-3 matrix to form a 4-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> vcat(Matrix<Element, 1, 3> const & rhs) const {
+    return Matrix<Element, 4, 3>::vcat(*this, rhs);
+  }
+
+  /// Forms a 3-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Matrix<Element, 1, 2> const & B,
+    Matrix<Element, 2, 1> const & C, Matrix<Element, 2, 2> const & D) {
+    return Matrix(
+      A, B.at(0, 0), B.at(0, 1)
+      , C.at(0, 0), D.at(0, 0), D.at(0, 1)
+      , C.at(1, 0), D.at(1, 0), D.at(1, 1)
+    );
+  }
+
+  /// Forms a 3-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 2> const & A, Element                         B,
+    Matrix<Element, 2, 2> const & C, Matrix<Element, 2, 1> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0)
+      , C.at(1, 0), C.at(1, 1), D.at(1, 0)
+    );
+  }
+
+  /// Forms a 3-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 1> const & A, Matrix<Element, 2, 2> const & B,
+    Element                         C, Matrix<Element, 1, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0), B.at(0, 1)
+      , A.at(1, 0), B.at(1, 0), B.at(1, 1)
+      , C, D.at(0, 0), D.at(0, 1)
+    );
+  }
+
+  /// Forms a 3-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 2> const & A, Matrix<Element, 2, 1> const & B,
+    Matrix<Element, 1, 2> const & C, Element                         D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0)
+      , A.at(1, 0), A.at(1, 1), B.at(1, 0)
+      , C.at(0, 0), C.at(0, 1), D
+    );
+  }
+
+  /// Elementwise add operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+    result.data[2] = data[2] + rhs.data[2];
+
+    result.data[3] = data[3] + rhs.data[3];
+    result.data[4] = data[4] + rhs.data[4];
+    result.data[5] = data[5] + rhs.data[5];
+
+    result.data[6] = data[6] + rhs.data[6];
+    result.data[7] = data[7] + rhs.data[7];
+    result.data[8] = data[8] + rhs.data[8];
+
+    return result;
+  }
+
+  /// Elementwise add operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+    data[2] += rhs.data[2];
+
+    data[3] += rhs.data[3];
+    data[4] += rhs.data[4];
+    data[5] += rhs.data[5];
+
+    data[6] += rhs.data[6];
+    data[7] += rhs.data[7];
+    data[8] += rhs.data[8];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+    result.data[2] = data[2] - rhs.data[2];
+
+    result.data[3] = data[3] - rhs.data[3];
+    result.data[4] = data[4] - rhs.data[4];
+    result.data[5] = data[5] - rhs.data[5];
+
+    result.data[6] = data[6] - rhs.data[6];
+    result.data[7] = data[7] - rhs.data[7];
+    result.data[8] = data[8] - rhs.data[8];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+    data[2] -= rhs.data[2];
+
+    data[3] -= rhs.data[3];
+    data[4] -= rhs.data[4];
+    data[5] -= rhs.data[5];
+
+    data[6] -= rhs.data[6];
+    data[7] -= rhs.data[7];
+    data[8] -= rhs.data[8];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+    result.data[2] = data[2] * rhs.data[2];
+
+    result.data[3] = data[3] * rhs.data[3];
+    result.data[4] = data[4] * rhs.data[4];
+    result.data[5] = data[5] * rhs.data[5];
+
+    result.data[6] = data[6] * rhs.data[6];
+    result.data[7] = data[7] * rhs.data[7];
+    result.data[8] = data[8] * rhs.data[8];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+    result.data[2] = data[2] * s;
+
+    result.data[3] = data[3] * s;
+    result.data[4] = data[4] * s;
+    result.data[5] = data[5] * s;
+
+    result.data[6] = data[6] * s;
+    result.data[7] = data[7] * s;
+    result.data[8] = data[8] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+    data[2] *= s;
+
+    data[3] *= s;
+    data[4] *= s;
+    data[5] *= s;
+
+    data[6] *= s;
+    data[7] *= s;
+    data[8] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+    result.data[2] = data[2] / rhs.data[2];
+
+    result.data[3] = data[3] / rhs.data[3];
+    result.data[4] = data[4] / rhs.data[4];
+    result.data[5] = data[5] / rhs.data[5];
+
+    result.data[6] = data[6] / rhs.data[6];
+    result.data[7] = data[7] / rhs.data[7];
+    result.data[8] = data[8] / rhs.data[8];
+
+    return result;
+  }
+
+  /// Scalar divide operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+    result.data[2] = data[2] / s;
+
+    result.data[3] = data[3] / s;
+    result.data[4] = data[4] / s;
+    result.data[5] = data[5] / s;
+
+    result.data[6] = data[6] / s;
+    result.data[7] = data[7] / s;
+    result.data[8] = data[8] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+    data[2] /= s;
+
+    data[3] /= s;
+    data[4] /= s;
+    data[5] /= s;
+
+    data[6] /= s;
+    data[7] /= s;
+    data[8] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (3-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+    data[2] /= rhs.data[2];
+
+    data[3] /= rhs.data[3];
+    data[4] /= rhs.data[4];
+    data[5] /= rhs.data[5];
+
+    data[6] /= rhs.data[6];
+    data[7] /= rhs.data[7];
+    data[8] /= rhs.data[8];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+    m.data[4] = -m.data[4];
+    m.data[5] = -m.data[5];
+    m.data[6] = -m.data[6];
+    m.data[7] = -m.data[7];
+    m.data[8] = -m.data[8];
+
+    return m;
+  }
+
+  /// Matrix product of size 3-by-1-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> product(
+    Matrix<Element, 3, 1> const &rhs,
+    Matrix<Element, 3, 1> accum = Matrix<Element, 3, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[3] * rhs.data[0];
+    accum.data[2] += data[6] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[4] * rhs.data[1];
+    accum.data[2] += data[7] * rhs.data[1];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[2];
+    accum.data[1] += data[5] * rhs.data[2];
+    accum.data[2] += data[8] * rhs.data[2];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-1-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> operator*(Matrix<Element, 3, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-2-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> product(
+    Matrix<Element, 3, 2> const &rhs,
+    Matrix<Element, 3, 2> accum = Matrix<Element, 3, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[3] * rhs.data[0];
+    accum.data[3] += data[3] * rhs.data[1];
+    accum.data[4] += data[6] * rhs.data[0];
+    accum.data[5] += data[6] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[4] * rhs.data[2];
+    accum.data[3] += data[4] * rhs.data[3];
+    accum.data[4] += data[7] * rhs.data[2];
+    accum.data[5] += data[7] * rhs.data[3];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[4];
+    accum.data[1] += data[2] * rhs.data[5];
+    accum.data[2] += data[5] * rhs.data[4];
+    accum.data[3] += data[5] * rhs.data[5];
+    accum.data[4] += data[8] * rhs.data[4];
+    accum.data[5] += data[8] * rhs.data[5];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-2-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> operator*(Matrix<Element, 3, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> product(
+    Matrix<Element, 3, 3> const &rhs,
+    Matrix<Element, 3, 3> accum = Matrix<Element, 3, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[3] * rhs.data[0];
+    accum.data[4] += data[3] * rhs.data[1];
+    accum.data[5] += data[3] * rhs.data[2];
+    accum.data[6] += data[6] * rhs.data[0];
+    accum.data[7] += data[6] * rhs.data[1];
+    accum.data[8] += data[6] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[4] * rhs.data[3];
+    accum.data[4] += data[4] * rhs.data[4];
+    accum.data[5] += data[4] * rhs.data[5];
+    accum.data[6] += data[7] * rhs.data[3];
+    accum.data[7] += data[7] * rhs.data[4];
+    accum.data[8] += data[7] * rhs.data[5];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[6];
+    accum.data[1] += data[2] * rhs.data[7];
+    accum.data[2] += data[2] * rhs.data[8];
+    accum.data[3] += data[5] * rhs.data[6];
+    accum.data[4] += data[5] * rhs.data[7];
+    accum.data[5] += data[5] * rhs.data[8];
+    accum.data[6] += data[8] * rhs.data[6];
+    accum.data[7] += data[8] * rhs.data[7];
+    accum.data[8] += data[8] * rhs.data[8];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> operator*(Matrix<Element, 3, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 3, 3> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 3-by-4-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> product(
+    Matrix<Element, 3, 4> const &rhs,
+    Matrix<Element, 3, 4> accum = Matrix<Element, 3, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[3] * rhs.data[0];
+    accum.data[5] += data[3] * rhs.data[1];
+    accum.data[6] += data[3] * rhs.data[2];
+    accum.data[7] += data[3] * rhs.data[3];
+    accum.data[8] += data[6] * rhs.data[0];
+    accum.data[9] += data[6] * rhs.data[1];
+    accum.data[10] += data[6] * rhs.data[2];
+    accum.data[11] += data[6] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[4] * rhs.data[4];
+    accum.data[5] += data[4] * rhs.data[5];
+    accum.data[6] += data[4] * rhs.data[6];
+    accum.data[7] += data[4] * rhs.data[7];
+    accum.data[8] += data[7] * rhs.data[4];
+    accum.data[9] += data[7] * rhs.data[5];
+    accum.data[10] += data[7] * rhs.data[6];
+    accum.data[11] += data[7] * rhs.data[7];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[8];
+    accum.data[1] += data[2] * rhs.data[9];
+    accum.data[2] += data[2] * rhs.data[10];
+    accum.data[3] += data[2] * rhs.data[11];
+    accum.data[4] += data[5] * rhs.data[8];
+    accum.data[5] += data[5] * rhs.data[9];
+    accum.data[6] += data[5] * rhs.data[10];
+    accum.data[7] += data[5] * rhs.data[11];
+    accum.data[8] += data[8] * rhs.data[8];
+    accum.data[9] += data[8] * rhs.data[9];
+    accum.data[10] += data[8] * rhs.data[10];
+    accum.data[11] += data[8] * rhs.data[11];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-4-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> operator*(Matrix<Element, 3, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+    accum += data[4];
+    accum += data[5];
+    accum += data[6];
+    accum += data[7];
+    accum += data[8];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+    accum += data[4] * data[4];
+    accum += data[5] * data[5];
+    accum += data[6] * data[6];
+    accum += data[7] * data[7];
+    accum += data[8] * data[8];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[4];
+    accum += data[8];
+
+    return accum;
+  }
+
+  /// Returns 3-by-3 rotation matrix around the X axis
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation_X(Element theta) {
+    Matrix m = identity();
+
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    m.at(1, 1) = c;
+    m.at(1, 2) = -s;
+    m.at(2, 1) = s;
+    m.at(2, 2) = c;
+
+    return m;
+  }
+
+  /// Returns 3-by-3 rotation matrix around the Y axis
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation_Y(Element theta) {
+    Matrix m = identity();
+
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    m.at(0, 0) = c;
+    m.at(2, 0) = -s;
+    m.at(0, 2) = s;
+    m.at(2, 2) = c;
+
+    return m;
+  }
+
+  /// Returns 3-by-3 rotation matrix around the Z axis
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation_Z(Element theta) {
+    Matrix m = Matrix::identity();
+
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    m.at(0, 0) = c;
+    m.at(0, 1) = -s;
+    m.at(1, 0) = s;
+    m.at(1, 1) = c;
+
+    return m;
+  }
+
+  /// Returns a 3-by-3 rotation matrix around a unit-length axis
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation(Element theta, Matrix<Element, 3, 1> const &u) {
+    Element x = u.data[0];
+    Element y = u.data[1];
+    Element z = u.data[2];
+
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    Element one_minus_cos = Element(1) - fast_cos(theta);
+
+    Matrix m;
+
+    m.set_slice3x3({
+      c + x * x * one_minus_cos, x * y * one_minus_cos - z * s, x * z * one_minus_cos + y * s,
+      y * x * one_minus_cos * z * s, c + y * y * one_minus_cos, y * z * one_minus_cos - x * s,
+      z * x * one_minus_cos - y * s, z * y * one_minus_cos + x * s, c + z * z * one_minus_cos
+    });
+
+    return m;
+  }
+
+  /// Returns a 3-by-3 reflection about the plane specified by the
+  /// unit-length normal vector n_unit
+  MUTLASS_HOST_DEVICE
+  static Matrix reflection(Matrix<Element, 3, 1> const &n_unit) {
+
+    Element a = n_unit.data[0];
+    Element b = n_unit.data[1];
+    Element c = n_unit.data[2];
+
+    Matrix m = Matrix::identity();
+
+    m.set_slice3x3({
+      Element(1) - Element(2) * a * a, Element(-2) * a * b, Element(-2) * a * c,
+      Element(-2) * a * b, Element(1) - Element(2) * b * b, Element(-2) * b * c,
+      Element(-2) * a * c, Element(-2) * b * c, Element(1) - Element(2) * c * c
+    });
+
+    return m;
+  }
+
+  /// Computes the determinant of a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Element determinant(Element accum = Element()) const {
+
+    accum += at(0, 0) * Matrix<Element, 2, 2>({ at(1, 1), at(1, 2), at(2, 1), at(2, 2) }).determinant();
+    accum -= at(0, 1) * Matrix<Element, 2, 2>({ at(1, 0), at(1, 2), at(2, 0), at(2, 2) }).determinant();
+    accum += at(0, 2) * Matrix<Element, 2, 2>({ at(1, 0), at(1, 1), at(2, 0), at(2, 1) }).determinant();
+
+    return accum;
+  }
+
+  /// Computes the inverse of a 3-by-3 matrix given
+  /// the matrix's determinant
+  MUTLASS_HOST_DEVICE
+  Matrix inverse(Element det) const {
+    return Matrix(
+      at(1, 1) * at(2, 2) - at(1, 2) * at(2, 1),
+      at(0, 2) * at(2, 1) - at(0, 1) * at(2, 2),
+      at(0, 1) * at(1, 2) - at(0, 2) * at(1, 1),
+
+      at(1, 2) * at(2, 0) - at(1, 0) * at(2, 2),
+      at(0, 0) * at(2, 2) - at(0, 2) * at(2, 0),
+      at(0, 2) * at(1, 0) - at(0, 0) * at(1, 2),
+
+      at(1, 0) * at(2, 1) - at(1, 1) * at(2, 0),
+      at(0, 1) * at(2, 0) - at(0, 0) * at(2, 1),
+      at(0, 0) * at(1, 1) - at(0, 1) * at(1, 0)
+    ) * (Element(1) / det);
+  }
+  /// Computes the inverse of a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix inverse() const {
+    return inverse(determinant());
+  }
+
+};
+
+/// Template alias for 3-by-3 matrix
+template <typename Element>
+using Matrix3x3 = Matrix<Element, 3, 3>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix3x3<Element> make_Matrix3x3(
+    Element _0_0, Element _0_1, Element _0_2,
+    Element _1_0, Element _1_1, Element _1_2,
+    Element _2_0, Element _2_1, Element _2_2
+) {
+  return Matrix3x3<Element>(
+  _0_0, _0_1, _0_2,
+  _1_0, _1_1, _1_2,
+  _2_0, _2_1, _2_2
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 3-by-4 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 3, 4> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 3;
+
+  /// Number of columns in matrix
+  static int const kColumns = 4;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 12;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 3-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 3-by-4 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1, Element _0_2, Element _0_3,
+    Element _1_0, Element _1_1, Element _1_2, Element _1_3,
+    Element _2_0, Element _2_1, Element _2_2, Element _2_3
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;  data[2] = _0_2;  data[3] = _0_3;
+    data[4] = _1_0;  data[5] = _1_1;  data[6] = _1_2;  data[7] = _1_3;
+    data[8] = _2_0;  data[9] = _2_1;  data[10] = _2_2;  data[11] = _2_3;
+  }
+
+  /// Constucts a 3-by-4 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 4> const &row_0,
+    Matrix<Element, 1, 4> const &row_1,
+    Matrix<Element, 1, 4> const &row_2
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_0.data[2];
+    data[3] = row_0.data[3];
+    data[4] = row_1.data[0];
+    data[5] = row_1.data[1];
+    data[6] = row_1.data[2];
+    data[7] = row_1.data[3];
+    data[8] = row_2.data[0];
+    data[9] = row_2.data[1];
+    data[10] = row_2.data[2];
+    data[11] = row_2.data[3];
+  }
+
+  /// Static method to construct a 3-by-4 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 4, 1> const &column_0,
+    Matrix<Element, 4, 1> const &column_1,
+    Matrix<Element, 4, 1> const &column_2,
+    Matrix<Element, 4, 1> const &column_3
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_2.data[0];
+    result.data[3] = column_3.data[0];
+    result.data[4] = column_0.data[1];
+    result.data[5] = column_1.data[1];
+    result.data[6] = column_2.data[1];
+    result.data[7] = column_3.data[1];
+    result.data[8] = column_0.data[2];
+    result.data[9] = column_1.data[2];
+    result.data[10] = column_2.data[2];
+    result.data[11] = column_3.data[2];
+    return result;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+    m.data[4] = s;
+    m.data[5] = s;
+    m.data[6] = s;
+    m.data[7] = s;
+    m.data[8] = s;
+    m.data[9] = s;
+    m.data[10] = s;
+    m.data[11] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 3, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[4] = diag.data[1];
+    m.data[8] = diag.data[2];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 3> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[4] = diag.data[1];
+    m.data[8] = diag.data[2];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> diagonal() const {
+    Matrix<Element, 3, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[4];
+    diag.data[2] = data[8];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> transpose() const {
+    Matrix<Element, 4, 3> mt;
+
+    mt.data[0] = data[0];
+    mt.data[3] = data[1];
+    mt.data[6] = data[2];
+    mt.data[9] = data[3];
+    mt.data[1] = data[4];
+    mt.data[4] = data[5];
+    mt.data[7] = data[6];
+    mt.data[10] = data[7];
+    mt.data[2] = data[8];
+    mt.data[5] = data[9];
+    mt.data[8] = data[10];
+    mt.data[11] = data[11];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 3 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 3 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> slice_1x3(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x3(Matrix<Element, 1, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> slice_1x4(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x4(Matrix<Element, 1, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> row(int i) const {
+    return slice_1x4(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 4> const &v, int i = 0) {
+    return set_slice_1x4(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 4];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 4] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 4];
+    m.data[3] = data[i * 4 + j + 5];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 4] = m.data[2];
+    data[i * 4 + j + 5] = m.data[3];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> slice_2x3(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 4];
+    m.data[4] = data[i * 4 + j + 5];
+    m.data[5] = data[i * 4 + j + 6];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x3(Matrix<Element, 2, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 4] = m.data[3];
+    data[i * 4 + j + 5] = m.data[4];
+    data[i * 4 + j + 6] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> slice_2x4(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+    m.data[4] = data[i * 4 + j + 4];
+    m.data[5] = data[i * 4 + j + 5];
+    m.data[6] = data[i * 4 + j + 6];
+    m.data[7] = data[i * 4 + j + 7];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x4(Matrix<Element, 2, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+    data[i * 4 + j + 4] = m.data[4];
+    data[i * 4 + j + 5] = m.data[5];
+    data[i * 4 + j + 6] = m.data[6];
+    data[i * 4 + j + 7] = m.data[7];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> slice_3x1(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 1> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 4];
+    m.data[2] = data[i * 4 + j + 8];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x1(Matrix<Element, 3, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 4] = m.data[1];
+    data[i * 4 + j + 8] = m.data[2];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> column(int j) const {
+    return slice_3x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 3, 1> const &v, int j =0) {
+    return set_slice_3x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> slice_3x2(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 4];
+    m.data[3] = data[i * 4 + j + 5];
+    m.data[4] = data[i * 4 + j + 8];
+    m.data[5] = data[i * 4 + j + 9];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x2(Matrix<Element, 3, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 4] = m.data[2];
+    data[i * 4 + j + 5] = m.data[3];
+    data[i * 4 + j + 8] = m.data[4];
+    data[i * 4 + j + 9] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> slice_3x3(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 4];
+    m.data[4] = data[i * 4 + j + 5];
+    m.data[5] = data[i * 4 + j + 6];
+    m.data[6] = data[i * 4 + j + 8];
+    m.data[7] = data[i * 4 + j + 9];
+    m.data[8] = data[i * 4 + j + 10];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x3(Matrix<Element, 3, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 4] = m.data[3];
+    data[i * 4 + j + 5] = m.data[4];
+    data[i * 4 + j + 6] = m.data[5];
+    data[i * 4 + j + 8] = m.data[6];
+    data[i * 4 + j + 9] = m.data[7];
+    data[i * 4 + j + 10] = m.data[8];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> slice_3x4(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+    m.data[4] = data[i * 4 + j + 4];
+    m.data[5] = data[i * 4 + j + 5];
+    m.data[6] = data[i * 4 + j + 6];
+    m.data[7] = data[i * 4 + j + 7];
+    m.data[8] = data[i * 4 + j + 8];
+    m.data[9] = data[i * 4 + j + 9];
+    m.data[10] = data[i * 4 + j + 10];
+    m.data[11] = data[i * 4 + j + 11];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x4(Matrix<Element, 3, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+    data[i * 4 + j + 4] = m.data[4];
+    data[i * 4 + j + 5] = m.data[5];
+    data[i * 4 + j + 6] = m.data[6];
+    data[i * 4 + j + 7] = m.data[7];
+    data[i * 4 + j + 8] = m.data[8];
+    data[i * 4 + j + 9] = m.data[9];
+    data[i * 4 + j + 10] = m.data[10];
+    data[i * 4 + j + 11] = m.data[11];
+
+    return *this;
+  }
+
+  /// Forms a 3-by-4 matrix by horizontally concatenating a 3-by-1 matrix with a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 3, 1> const & lhs, Matrix<Element, 3, 3> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1), rhs.at(0, 2)
+      , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1), rhs.at(1, 2)
+      , lhs.at(2, 0), rhs.at(2, 0), rhs.at(2, 1), rhs.at(2, 2));
+  }
+
+  /// Forms a 3-by-4 matrix by horizontally concatenating a 3-by-2 matrix with a 3-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 3, 2> const & lhs, Matrix<Element, 3, 2> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0), rhs.at(0, 1)
+      , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0), rhs.at(1, 1)
+      , lhs.at(2, 0), lhs.at(2, 1), rhs.at(2, 0), rhs.at(2, 1));
+  }
+
+  /// Forms a 3-by-4 matrix by horizontally concatenating a 3-by-3 matrix with a 3-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 3, 3> const & lhs, Matrix<Element, 3, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), lhs.at(0, 2), rhs.at(0, 0)
+      , lhs.at(1, 0), lhs.at(1, 1), lhs.at(1, 2), rhs.at(1, 0)
+      , lhs.at(2, 0), lhs.at(2, 1), lhs.at(2, 2), rhs.at(2, 0));
+  }
+
+  /// Forms a 3-by-4 matrix by vertically concatenating a 1-by-4 matrix with a 2-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 4> const & upper, Matrix<Element, 2, 4> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3)
+      , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2), lower.at(1, 3));
+  }
+
+  /// Forms a 3-by-4 matrix by vertically concatenating a 2-by-4 matrix with a 1-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 2, 4> const & upper, Matrix<Element, 1, 4> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3)
+      , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2), upper.at(1, 3)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3));
+  }
+
+  /// Concatenates this matrix with a a 1-by-4 matrix to form a 4-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> vcat(Matrix<Element, 1, 4> const & rhs) const {
+    return Matrix<Element, 4, 4>::vcat(*this, rhs);
+  }
+
+  /// Forms a 3-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Matrix<Element, 1, 3> const & B,
+    Matrix<Element, 2, 1> const & C, Matrix<Element, 2, 3> const & D) {
+    return Matrix(
+      A, B.at(0, 0), B.at(0, 1), B.at(0, 2)
+      , C.at(0, 0), D.at(0, 0), D.at(0, 1), D.at(0, 2)
+      , C.at(1, 0), D.at(1, 0), D.at(1, 1), D.at(1, 2)
+    );
+  }
+
+  /// Forms a 3-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 2> const & A, Matrix<Element, 1, 2> const & B,
+    Matrix<Element, 2, 2> const & C, Matrix<Element, 2, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1)
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1)
+      , C.at(1, 0), C.at(1, 1), D.at(1, 0), D.at(1, 1)
+    );
+  }
+
+  /// Forms a 3-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 3> const & A, Element                         B,
+    Matrix<Element, 2, 3> const & C, Matrix<Element, 2, 1> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), A.at(0, 2), B
+      , C.at(0, 0), C.at(0, 1), C.at(0, 2), D.at(0, 0)
+      , C.at(1, 0), C.at(1, 1), C.at(1, 2), D.at(1, 0)
+    );
+  }
+
+  /// Forms a 3-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 1> const & A, Matrix<Element, 2, 3> const & B,
+    Element                         C, Matrix<Element, 1, 3> const & D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0), B.at(0, 1), B.at(0, 2)
+      , A.at(1, 0), B.at(1, 0), B.at(1, 1), B.at(1, 2)
+      , C, D.at(0, 0), D.at(0, 1), D.at(0, 2)
+    );
+  }
+
+  /// Forms a 3-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 2> const & A, Matrix<Element, 2, 2> const & B,
+    Matrix<Element, 1, 2> const & C, Matrix<Element, 1, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1)
+      , A.at(1, 0), A.at(1, 1), B.at(1, 0), B.at(1, 1)
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1)
+    );
+  }
+
+  /// Forms a 3-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 3> const & A, Matrix<Element, 2, 1> const & B,
+    Matrix<Element, 1, 3> const & C, Element                         D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), A.at(0, 2), B.at(0, 0)
+      , A.at(1, 0), A.at(1, 1), A.at(1, 2), B.at(1, 0)
+      , C.at(0, 0), C.at(0, 1), C.at(0, 2), D
+    );
+  }
+
+  /// Elementwise add operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+    result.data[2] = data[2] + rhs.data[2];
+    result.data[3] = data[3] + rhs.data[3];
+
+    result.data[4] = data[4] + rhs.data[4];
+    result.data[5] = data[5] + rhs.data[5];
+    result.data[6] = data[6] + rhs.data[6];
+    result.data[7] = data[7] + rhs.data[7];
+
+    result.data[8] = data[8] + rhs.data[8];
+    result.data[9] = data[9] + rhs.data[9];
+    result.data[10] = data[10] + rhs.data[10];
+    result.data[11] = data[11] + rhs.data[11];
+
+    return result;
+  }
+
+  /// Elementwise add operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+    data[2] += rhs.data[2];
+    data[3] += rhs.data[3];
+
+    data[4] += rhs.data[4];
+    data[5] += rhs.data[5];
+    data[6] += rhs.data[6];
+    data[7] += rhs.data[7];
+
+    data[8] += rhs.data[8];
+    data[9] += rhs.data[9];
+    data[10] += rhs.data[10];
+    data[11] += rhs.data[11];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+    result.data[2] = data[2] - rhs.data[2];
+    result.data[3] = data[3] - rhs.data[3];
+
+    result.data[4] = data[4] - rhs.data[4];
+    result.data[5] = data[5] - rhs.data[5];
+    result.data[6] = data[6] - rhs.data[6];
+    result.data[7] = data[7] - rhs.data[7];
+
+    result.data[8] = data[8] - rhs.data[8];
+    result.data[9] = data[9] - rhs.data[9];
+    result.data[10] = data[10] - rhs.data[10];
+    result.data[11] = data[11] - rhs.data[11];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+    data[2] -= rhs.data[2];
+    data[3] -= rhs.data[3];
+
+    data[4] -= rhs.data[4];
+    data[5] -= rhs.data[5];
+    data[6] -= rhs.data[6];
+    data[7] -= rhs.data[7];
+
+    data[8] -= rhs.data[8];
+    data[9] -= rhs.data[9];
+    data[10] -= rhs.data[10];
+    data[11] -= rhs.data[11];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+    result.data[2] = data[2] * rhs.data[2];
+    result.data[3] = data[3] * rhs.data[3];
+
+    result.data[4] = data[4] * rhs.data[4];
+    result.data[5] = data[5] * rhs.data[5];
+    result.data[6] = data[6] * rhs.data[6];
+    result.data[7] = data[7] * rhs.data[7];
+
+    result.data[8] = data[8] * rhs.data[8];
+    result.data[9] = data[9] * rhs.data[9];
+    result.data[10] = data[10] * rhs.data[10];
+    result.data[11] = data[11] * rhs.data[11];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+    result.data[2] = data[2] * s;
+    result.data[3] = data[3] * s;
+
+    result.data[4] = data[4] * s;
+    result.data[5] = data[5] * s;
+    result.data[6] = data[6] * s;
+    result.data[7] = data[7] * s;
+
+    result.data[8] = data[8] * s;
+    result.data[9] = data[9] * s;
+    result.data[10] = data[10] * s;
+    result.data[11] = data[11] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+    data[2] *= s;
+    data[3] *= s;
+
+    data[4] *= s;
+    data[5] *= s;
+    data[6] *= s;
+    data[7] *= s;
+
+    data[8] *= s;
+    data[9] *= s;
+    data[10] *= s;
+    data[11] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+    result.data[2] = data[2] / rhs.data[2];
+    result.data[3] = data[3] / rhs.data[3];
+
+    result.data[4] = data[4] / rhs.data[4];
+    result.data[5] = data[5] / rhs.data[5];
+    result.data[6] = data[6] / rhs.data[6];
+    result.data[7] = data[7] / rhs.data[7];
+
+    result.data[8] = data[8] / rhs.data[8];
+    result.data[9] = data[9] / rhs.data[9];
+    result.data[10] = data[10] / rhs.data[10];
+    result.data[11] = data[11] / rhs.data[11];
+
+    return result;
+  }
+
+  /// Scalar divide operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+    result.data[2] = data[2] / s;
+    result.data[3] = data[3] / s;
+
+    result.data[4] = data[4] / s;
+    result.data[5] = data[5] / s;
+    result.data[6] = data[6] / s;
+    result.data[7] = data[7] / s;
+
+    result.data[8] = data[8] / s;
+    result.data[9] = data[9] / s;
+    result.data[10] = data[10] / s;
+    result.data[11] = data[11] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+    data[2] /= s;
+    data[3] /= s;
+
+    data[4] /= s;
+    data[5] /= s;
+    data[6] /= s;
+    data[7] /= s;
+
+    data[8] /= s;
+    data[9] /= s;
+    data[10] /= s;
+    data[11] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (3-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+    data[2] /= rhs.data[2];
+    data[3] /= rhs.data[3];
+
+    data[4] /= rhs.data[4];
+    data[5] /= rhs.data[5];
+    data[6] /= rhs.data[6];
+    data[7] /= rhs.data[7];
+
+    data[8] /= rhs.data[8];
+    data[9] /= rhs.data[9];
+    data[10] /= rhs.data[10];
+    data[11] /= rhs.data[11];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+    m.data[4] = -m.data[4];
+    m.data[5] = -m.data[5];
+    m.data[6] = -m.data[6];
+    m.data[7] = -m.data[7];
+    m.data[8] = -m.data[8];
+    m.data[9] = -m.data[9];
+    m.data[10] = -m.data[10];
+    m.data[11] = -m.data[11];
+
+    return m;
+  }
+
+  /// Matrix product of size 3-by-1-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> product(
+    Matrix<Element, 4, 1> const &rhs,
+    Matrix<Element, 3, 1> accum = Matrix<Element, 3, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[4] * rhs.data[0];
+    accum.data[2] += data[8] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[5] * rhs.data[1];
+    accum.data[2] += data[9] * rhs.data[1];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[2];
+    accum.data[1] += data[6] * rhs.data[2];
+    accum.data[2] += data[10] * rhs.data[2];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[3];
+    accum.data[1] += data[7] * rhs.data[3];
+    accum.data[2] += data[11] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-1-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> operator*(Matrix<Element, 4, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-2-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> product(
+    Matrix<Element, 4, 2> const &rhs,
+    Matrix<Element, 3, 2> accum = Matrix<Element, 3, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[4] * rhs.data[0];
+    accum.data[3] += data[4] * rhs.data[1];
+    accum.data[4] += data[8] * rhs.data[0];
+    accum.data[5] += data[8] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[5] * rhs.data[2];
+    accum.data[3] += data[5] * rhs.data[3];
+    accum.data[4] += data[9] * rhs.data[2];
+    accum.data[5] += data[9] * rhs.data[3];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[4];
+    accum.data[1] += data[2] * rhs.data[5];
+    accum.data[2] += data[6] * rhs.data[4];
+    accum.data[3] += data[6] * rhs.data[5];
+    accum.data[4] += data[10] * rhs.data[4];
+    accum.data[5] += data[10] * rhs.data[5];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[6];
+    accum.data[1] += data[3] * rhs.data[7];
+    accum.data[2] += data[7] * rhs.data[6];
+    accum.data[3] += data[7] * rhs.data[7];
+    accum.data[4] += data[11] * rhs.data[6];
+    accum.data[5] += data[11] * rhs.data[7];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-2-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> operator*(Matrix<Element, 4, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-3-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> product(
+    Matrix<Element, 4, 3> const &rhs,
+    Matrix<Element, 3, 3> accum = Matrix<Element, 3, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[4] * rhs.data[0];
+    accum.data[4] += data[4] * rhs.data[1];
+    accum.data[5] += data[4] * rhs.data[2];
+    accum.data[6] += data[8] * rhs.data[0];
+    accum.data[7] += data[8] * rhs.data[1];
+    accum.data[8] += data[8] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[5] * rhs.data[3];
+    accum.data[4] += data[5] * rhs.data[4];
+    accum.data[5] += data[5] * rhs.data[5];
+    accum.data[6] += data[9] * rhs.data[3];
+    accum.data[7] += data[9] * rhs.data[4];
+    accum.data[8] += data[9] * rhs.data[5];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[6];
+    accum.data[1] += data[2] * rhs.data[7];
+    accum.data[2] += data[2] * rhs.data[8];
+    accum.data[3] += data[6] * rhs.data[6];
+    accum.data[4] += data[6] * rhs.data[7];
+    accum.data[5] += data[6] * rhs.data[8];
+    accum.data[6] += data[10] * rhs.data[6];
+    accum.data[7] += data[10] * rhs.data[7];
+    accum.data[8] += data[10] * rhs.data[8];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[9];
+    accum.data[1] += data[3] * rhs.data[10];
+    accum.data[2] += data[3] * rhs.data[11];
+    accum.data[3] += data[7] * rhs.data[9];
+    accum.data[4] += data[7] * rhs.data[10];
+    accum.data[5] += data[7] * rhs.data[11];
+    accum.data[6] += data[11] * rhs.data[9];
+    accum.data[7] += data[11] * rhs.data[10];
+    accum.data[8] += data[11] * rhs.data[11];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-3-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> operator*(Matrix<Element, 4, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> product(
+    Matrix<Element, 4, 4> const &rhs,
+    Matrix<Element, 3, 4> accum = Matrix<Element, 3, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[4] * rhs.data[0];
+    accum.data[5] += data[4] * rhs.data[1];
+    accum.data[6] += data[4] * rhs.data[2];
+    accum.data[7] += data[4] * rhs.data[3];
+    accum.data[8] += data[8] * rhs.data[0];
+    accum.data[9] += data[8] * rhs.data[1];
+    accum.data[10] += data[8] * rhs.data[2];
+    accum.data[11] += data[8] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[5] * rhs.data[4];
+    accum.data[5] += data[5] * rhs.data[5];
+    accum.data[6] += data[5] * rhs.data[6];
+    accum.data[7] += data[5] * rhs.data[7];
+    accum.data[8] += data[9] * rhs.data[4];
+    accum.data[9] += data[9] * rhs.data[5];
+    accum.data[10] += data[9] * rhs.data[6];
+    accum.data[11] += data[9] * rhs.data[7];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[8];
+    accum.data[1] += data[2] * rhs.data[9];
+    accum.data[2] += data[2] * rhs.data[10];
+    accum.data[3] += data[2] * rhs.data[11];
+    accum.data[4] += data[6] * rhs.data[8];
+    accum.data[5] += data[6] * rhs.data[9];
+    accum.data[6] += data[6] * rhs.data[10];
+    accum.data[7] += data[6] * rhs.data[11];
+    accum.data[8] += data[10] * rhs.data[8];
+    accum.data[9] += data[10] * rhs.data[9];
+    accum.data[10] += data[10] * rhs.data[10];
+    accum.data[11] += data[10] * rhs.data[11];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[12];
+    accum.data[1] += data[3] * rhs.data[13];
+    accum.data[2] += data[3] * rhs.data[14];
+    accum.data[3] += data[3] * rhs.data[15];
+    accum.data[4] += data[7] * rhs.data[12];
+    accum.data[5] += data[7] * rhs.data[13];
+    accum.data[6] += data[7] * rhs.data[14];
+    accum.data[7] += data[7] * rhs.data[15];
+    accum.data[8] += data[11] * rhs.data[12];
+    accum.data[9] += data[11] * rhs.data[13];
+    accum.data[10] += data[11] * rhs.data[14];
+    accum.data[11] += data[11] * rhs.data[15];
+
+    return accum;
+  }
+
+  /// Matrix product of size 3-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> operator*(Matrix<Element, 4, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 3-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 4, 4> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+    accum += data[4];
+    accum += data[5];
+    accum += data[6];
+    accum += data[7];
+    accum += data[8];
+    accum += data[9];
+    accum += data[10];
+    accum += data[11];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+    accum += data[4] * data[4];
+    accum += data[5] * data[5];
+    accum += data[6] * data[6];
+    accum += data[7] * data[7];
+    accum += data[8] * data[8];
+    accum += data[9] * data[9];
+    accum += data[10] * data[10];
+    accum += data[11] * data[11];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[5];
+    accum += data[10];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 3-by-4 matrix
+template <typename Element>
+using Matrix3x4 = Matrix<Element, 3, 4>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix3x4<Element> make_Matrix3x4(
+    Element _0_0, Element _0_1, Element _0_2, Element _0_3,
+    Element _1_0, Element _1_1, Element _1_2, Element _1_3,
+    Element _2_0, Element _2_1, Element _2_2, Element _2_3
+) {
+  return Matrix3x4<Element>(
+  _0_0, _0_1, _0_2, _0_3,
+  _1_0, _1_1, _1_2, _1_3,
+  _2_0, _2_1, _2_2, _2_3
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 4-by-1 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 4, 1> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 4;
+
+  /// Number of columns in matrix
+  static int const kColumns = 1;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 4;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 4-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 4-by-1 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0,
+    Element _1_0,
+    Element _2_0,
+    Element _3_0
+  ) {
+
+    data[0] = _0_0;
+    data[1] = _1_0;
+    data[2] = _2_0;
+    data[3] = _3_0;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> transpose() const {
+    Matrix<Element, 1, 4> mt;
+
+    mt.data[0] = data[0];
+    mt.data[1] = data[1];
+    mt.data[2] = data[2];
+    mt.data[3] = data[3];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 4 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 4 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 1 + j + 0];
+    m.data[1] = data[i * 1 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 1 + j + 0] = m.data[0];
+    data[i * 1 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> slice_3x1(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 1> m;
+
+    m.data[0] = data[i * 1 + j + 0];
+    m.data[1] = data[i * 1 + j + 1];
+    m.data[2] = data[i * 1 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x1(Matrix<Element, 3, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 1 + j + 0] = m.data[0];
+    data[i * 1 + j + 1] = m.data[1];
+    data[i * 1 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> slice_4x1(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 1> m;
+
+    m.data[0] = data[i * 1 + j + 0];
+    m.data[1] = data[i * 1 + j + 1];
+    m.data[2] = data[i * 1 + j + 2];
+    m.data[3] = data[i * 1 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x1(Matrix<Element, 4, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 1 + j + 0] = m.data[0];
+    data[i * 1 + j + 1] = m.data[1];
+    data[i * 1 + j + 2] = m.data[2];
+    data[i * 1 + j + 3] = m.data[3];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> column(int j) const {
+    return slice_4x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 4, 1> const &v, int j =0) {
+    return set_slice_4x1(v, 0, j);
+  }
+
+  /// Concatenates this matrix with a a 4-by-1 matrix to form a 4-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> hcat(Matrix<Element, 4, 1> const & rhs) const {
+    return Matrix<Element, 4, 2>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 4-by-2 matrix to form a 4-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> hcat(Matrix<Element, 4, 2> const & rhs) const {
+    return Matrix<Element, 4, 3>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 4-by-3 matrix to form a 4-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> hcat(Matrix<Element, 4, 3> const & rhs) const {
+    return Matrix<Element, 4, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 4-by-1 matrix by vertically concatenating an Element with a 3-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Element upper, Matrix<Element, 3, 1> const & lower) {
+    return Matrix(
+      upper
+      , lower.at(0, 0)
+      , lower.at(1, 0)
+      , lower.at(2, 0));
+  }
+
+  /// Forms a 4-by-1 matrix by vertically concatenating a 2-by-1 matrix with a 2-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 2, 1> const & upper, Matrix<Element, 2, 1> const & lower) {
+    return Matrix(
+      upper.at(0, 0)
+      , upper.at(1, 0)
+      , lower.at(0, 0)
+      , lower.at(1, 0));
+  }
+
+  /// Forms a 4-by-1 matrix by vertically concatenating a 3-by-1 matrix with an Element
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 3, 1> const & upper, Element lower) {
+    return Matrix(
+      upper.at(0, 0)
+      , upper.at(1, 0)
+      , upper.at(2, 0)
+      , lower);
+  }
+
+  /// Elementwise add operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+
+    result.data[1] = data[1] + rhs.data[1];
+
+    result.data[2] = data[2] + rhs.data[2];
+
+    result.data[3] = data[3] + rhs.data[3];
+
+    return result;
+  }
+
+  /// Elementwise add operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+
+    data[1] += rhs.data[1];
+
+    data[2] += rhs.data[2];
+
+    data[3] += rhs.data[3];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+
+    result.data[1] = data[1] - rhs.data[1];
+
+    result.data[2] = data[2] - rhs.data[2];
+
+    result.data[3] = data[3] - rhs.data[3];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+
+    data[1] -= rhs.data[1];
+
+    data[2] -= rhs.data[2];
+
+    data[3] -= rhs.data[3];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+
+    result.data[1] = data[1] * rhs.data[1];
+
+    result.data[2] = data[2] * rhs.data[2];
+
+    result.data[3] = data[3] * rhs.data[3];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+
+    result.data[1] = data[1] * s;
+
+    result.data[2] = data[2] * s;
+
+    result.data[3] = data[3] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+
+    data[1] *= s;
+
+    data[2] *= s;
+
+    data[3] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+
+    result.data[1] = data[1] / rhs.data[1];
+
+    result.data[2] = data[2] / rhs.data[2];
+
+    result.data[3] = data[3] / rhs.data[3];
+
+    return result;
+  }
+
+  /// Scalar divide operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+
+    result.data[1] = data[1] / s;
+
+    result.data[2] = data[2] / s;
+
+    result.data[3] = data[3] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+
+    data[1] /= s;
+
+    data[2] /= s;
+
+    data[3] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (4-by-1)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+
+    data[1] /= rhs.data[1];
+
+    data[2] /= rhs.data[2];
+
+    data[3] /= rhs.data[3];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+
+    return m;
+  }
+
+  /// Matrix product of size 4-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> product(
+    Matrix<Element, 1, 1> const &rhs,
+    Matrix<Element, 4, 1> accum = Matrix<Element, 4, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[1] * rhs.data[0];
+    accum.data[2] += data[2] * rhs.data[0];
+    accum.data[3] += data[3] * rhs.data[0];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> operator*(Matrix<Element, 1, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-1-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 1, 1> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 4-by-2-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> product(
+    Matrix<Element, 1, 2> const &rhs,
+    Matrix<Element, 4, 2> accum = Matrix<Element, 4, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[1] * rhs.data[0];
+    accum.data[3] += data[1] * rhs.data[1];
+    accum.data[4] += data[2] * rhs.data[0];
+    accum.data[5] += data[2] * rhs.data[1];
+    accum.data[6] += data[3] * rhs.data[0];
+    accum.data[7] += data[3] * rhs.data[1];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-2-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> operator*(Matrix<Element, 1, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-3-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> product(
+    Matrix<Element, 1, 3> const &rhs,
+    Matrix<Element, 4, 3> accum = Matrix<Element, 4, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[1] * rhs.data[0];
+    accum.data[4] += data[1] * rhs.data[1];
+    accum.data[5] += data[1] * rhs.data[2];
+    accum.data[6] += data[2] * rhs.data[0];
+    accum.data[7] += data[2] * rhs.data[1];
+    accum.data[8] += data[2] * rhs.data[2];
+    accum.data[9] += data[3] * rhs.data[0];
+    accum.data[10] += data[3] * rhs.data[1];
+    accum.data[11] += data[3] * rhs.data[2];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-3-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> operator*(Matrix<Element, 1, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-4-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> product(
+    Matrix<Element, 1, 4> const &rhs,
+    Matrix<Element, 4, 4> accum = Matrix<Element, 4, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[1] * rhs.data[0];
+    accum.data[5] += data[1] * rhs.data[1];
+    accum.data[6] += data[1] * rhs.data[2];
+    accum.data[7] += data[1] * rhs.data[3];
+    accum.data[8] += data[2] * rhs.data[0];
+    accum.data[9] += data[2] * rhs.data[1];
+    accum.data[10] += data[2] * rhs.data[2];
+    accum.data[11] += data[2] * rhs.data[3];
+    accum.data[12] += data[3] * rhs.data[0];
+    accum.data[13] += data[3] * rhs.data[1];
+    accum.data[14] += data[3] * rhs.data[2];
+    accum.data[15] += data[3] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-4-by-1
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> operator*(Matrix<Element, 1, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Dot product of vectors with extent 4
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 4, 1> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    accum += data[2] * rhs.data[2];
+    accum += data[3] * rhs.data[3];
+    return accum;
+  }
+
+  /// Dot product of vectors with extent 4
+  MUTLASS_HOST_DEVICE
+  Element dot(Matrix<Element, 1, 4> const &rhs, Element accum = Element()) const {
+
+    accum += data[0] * rhs.data[0];
+    accum += data[1] * rhs.data[1];
+    accum += data[2] * rhs.data[2];
+    accum += data[3] * rhs.data[3];
+    return accum;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 4-by-1 matrix
+template <typename Element>
+using Matrix4x1 = Matrix<Element, 4, 1>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix4x1<Element> make_Matrix4x1(
+    Element _0_0,
+    Element _1_0,
+    Element _2_0,
+    Element _3_0
+) {
+  return Matrix4x1<Element>(
+  _0_0,
+  _1_0,
+  _2_0,
+  _3_0
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 4-by-2 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 4, 2> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 4;
+
+  /// Number of columns in matrix
+  static int const kColumns = 2;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 8;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 4-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 4-by-2 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1,
+    Element _1_0, Element _1_1,
+    Element _2_0, Element _2_1,
+    Element _3_0, Element _3_1
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;
+    data[2] = _1_0;  data[3] = _1_1;
+    data[4] = _2_0;  data[5] = _2_1;
+    data[6] = _3_0;  data[7] = _3_1;
+  }
+
+  /// Constucts a 4-by-2 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 2> const &row_0,
+    Matrix<Element, 1, 2> const &row_1,
+    Matrix<Element, 1, 2> const &row_2,
+    Matrix<Element, 1, 2> const &row_3
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_1.data[0];
+    data[3] = row_1.data[1];
+    data[4] = row_2.data[0];
+    data[5] = row_2.data[1];
+    data[6] = row_3.data[0];
+    data[7] = row_3.data[1];
+  }
+
+  /// Static method to construct a 4-by-2 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 2, 1> const &column_0,
+    Matrix<Element, 2, 1> const &column_1
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_0.data[1];
+    result.data[3] = column_1.data[1];
+    result.data[4] = column_0.data[2];
+    result.data[5] = column_1.data[2];
+    result.data[6] = column_0.data[3];
+    result.data[7] = column_1.data[3];
+    return result;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+    m.data[4] = s;
+    m.data[5] = s;
+    m.data[6] = s;
+    m.data[7] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 2, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[5] = diag.data[1];
+    m.data[10] = diag.data[2];
+    m.data[15] = diag.data[3];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 2> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[5] = diag.data[1];
+    m.data[10] = diag.data[2];
+    m.data[15] = diag.data[3];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> diagonal() const {
+    Matrix<Element, 2, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[5];
+    diag.data[2] = data[10];
+    diag.data[3] = data[15];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> transpose() const {
+    Matrix<Element, 2, 4> mt;
+
+    mt.data[0] = data[0];
+    mt.data[4] = data[1];
+    mt.data[1] = data[2];
+    mt.data[5] = data[3];
+    mt.data[2] = data[4];
+    mt.data[6] = data[5];
+    mt.data[3] = data[6];
+    mt.data[7] = data[7];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 4 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 4 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> row(int i) const {
+    return slice_1x2(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 2> const &v, int i = 0) {
+    return set_slice_1x2(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 2] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+    m.data[2] = data[i * 2 + j + 2];
+    m.data[3] = data[i * 2 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+    data[i * 2 + j + 2] = m.data[2];
+    data[i * 2 + j + 3] = m.data[3];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> slice_3x1(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 1> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 2];
+    m.data[2] = data[i * 2 + j + 4];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x1(Matrix<Element, 3, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 2] = m.data[1];
+    data[i * 2 + j + 4] = m.data[2];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> slice_3x2(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+    m.data[2] = data[i * 2 + j + 2];
+    m.data[3] = data[i * 2 + j + 3];
+    m.data[4] = data[i * 2 + j + 4];
+    m.data[5] = data[i * 2 + j + 5];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x2(Matrix<Element, 3, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+    data[i * 2 + j + 2] = m.data[2];
+    data[i * 2 + j + 3] = m.data[3];
+    data[i * 2 + j + 4] = m.data[4];
+    data[i * 2 + j + 5] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> slice_4x1(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 1> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 2];
+    m.data[2] = data[i * 2 + j + 4];
+    m.data[3] = data[i * 2 + j + 6];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x1(Matrix<Element, 4, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 2] = m.data[1];
+    data[i * 2 + j + 4] = m.data[2];
+    data[i * 2 + j + 6] = m.data[3];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> column(int j) const {
+    return slice_4x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 4, 1> const &v, int j =0) {
+    return set_slice_4x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> slice_4x2(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 2> m;
+
+    m.data[0] = data[i * 2 + j + 0];
+    m.data[1] = data[i * 2 + j + 1];
+    m.data[2] = data[i * 2 + j + 2];
+    m.data[3] = data[i * 2 + j + 3];
+    m.data[4] = data[i * 2 + j + 4];
+    m.data[5] = data[i * 2 + j + 5];
+    m.data[6] = data[i * 2 + j + 6];
+    m.data[7] = data[i * 2 + j + 7];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x2(Matrix<Element, 4, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 2 + j + 0] = m.data[0];
+    data[i * 2 + j + 1] = m.data[1];
+    data[i * 2 + j + 2] = m.data[2];
+    data[i * 2 + j + 3] = m.data[3];
+    data[i * 2 + j + 4] = m.data[4];
+    data[i * 2 + j + 5] = m.data[5];
+    data[i * 2 + j + 6] = m.data[6];
+    data[i * 2 + j + 7] = m.data[7];
+
+    return *this;
+  }
+
+  /// Forms a 4-by-2 matrix by horizontally concatenating a 4-by-1 matrix with a 4-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 4, 1> const & lhs, Matrix<Element, 4, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0)
+      , lhs.at(1, 0), rhs.at(1, 0)
+      , lhs.at(2, 0), rhs.at(2, 0)
+      , lhs.at(3, 0), rhs.at(3, 0));
+  }
+
+  /// Concatenates this matrix with a a 4-by-1 matrix to form a 4-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> hcat(Matrix<Element, 4, 1> const & rhs) const {
+    return Matrix<Element, 4, 3>::hcat(*this, rhs);
+  }
+
+  /// Concatenates this matrix with a a 4-by-2 matrix to form a 4-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> hcat(Matrix<Element, 4, 2> const & rhs) const {
+    return Matrix<Element, 4, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 4-by-2 matrix by vertically concatenating a 1-by-2 matrix with a 3-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 2> const & upper, Matrix<Element, 3, 2> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1)
+      , lower.at(0, 0), lower.at(0, 1)
+      , lower.at(1, 0), lower.at(1, 1)
+      , lower.at(2, 0), lower.at(2, 1));
+  }
+
+  /// Forms a 4-by-2 matrix by vertically concatenating a 2-by-2 matrix with a 2-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 2, 2> const & upper, Matrix<Element, 2, 2> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1)
+      , upper.at(1, 0), upper.at(1, 1)
+      , lower.at(0, 0), lower.at(0, 1)
+      , lower.at(1, 0), lower.at(1, 1));
+  }
+
+  /// Forms a 4-by-2 matrix by vertically concatenating a 3-by-2 matrix with a 1-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 3, 2> const & upper, Matrix<Element, 1, 2> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1)
+      , upper.at(1, 0), upper.at(1, 1)
+      , upper.at(2, 0), upper.at(2, 1)
+      , lower.at(0, 0), lower.at(0, 1));
+  }
+
+  /// Forms a 4-by-2 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Element                         B,
+    Matrix<Element, 3, 1> const & C, Matrix<Element, 3, 1> const & D) {
+    return Matrix(
+      A, B
+      , C.at(0, 0), D.at(0, 0)
+      , C.at(1, 0), D.at(1, 0)
+      , C.at(2, 0), D.at(2, 0)
+    );
+  }
+
+  /// Forms a 4-by-2 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 1> const & A, Matrix<Element, 2, 1> const & B,
+    Matrix<Element, 2, 1> const & C, Matrix<Element, 2, 1> const & D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0)
+      , A.at(1, 0), B.at(1, 0)
+      , C.at(0, 0), D.at(0, 0)
+      , C.at(1, 0), D.at(1, 0)
+    );
+  }
+
+  /// Forms a 4-by-2 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 3, 1> const & A, Matrix<Element, 3, 1> const & B,
+    Element                         C, Element                         D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0)
+      , A.at(1, 0), B.at(1, 0)
+      , A.at(2, 0), B.at(2, 0)
+      , C, D
+    );
+  }
+
+  /// Elementwise add operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+
+    result.data[2] = data[2] + rhs.data[2];
+    result.data[3] = data[3] + rhs.data[3];
+
+    result.data[4] = data[4] + rhs.data[4];
+    result.data[5] = data[5] + rhs.data[5];
+
+    result.data[6] = data[6] + rhs.data[6];
+    result.data[7] = data[7] + rhs.data[7];
+
+    return result;
+  }
+
+  /// Elementwise add operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+
+    data[2] += rhs.data[2];
+    data[3] += rhs.data[3];
+
+    data[4] += rhs.data[4];
+    data[5] += rhs.data[5];
+
+    data[6] += rhs.data[6];
+    data[7] += rhs.data[7];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+
+    result.data[2] = data[2] - rhs.data[2];
+    result.data[3] = data[3] - rhs.data[3];
+
+    result.data[4] = data[4] - rhs.data[4];
+    result.data[5] = data[5] - rhs.data[5];
+
+    result.data[6] = data[6] - rhs.data[6];
+    result.data[7] = data[7] - rhs.data[7];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+
+    data[2] -= rhs.data[2];
+    data[3] -= rhs.data[3];
+
+    data[4] -= rhs.data[4];
+    data[5] -= rhs.data[5];
+
+    data[6] -= rhs.data[6];
+    data[7] -= rhs.data[7];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+
+    result.data[2] = data[2] * rhs.data[2];
+    result.data[3] = data[3] * rhs.data[3];
+
+    result.data[4] = data[4] * rhs.data[4];
+    result.data[5] = data[5] * rhs.data[5];
+
+    result.data[6] = data[6] * rhs.data[6];
+    result.data[7] = data[7] * rhs.data[7];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+
+    result.data[2] = data[2] * s;
+    result.data[3] = data[3] * s;
+
+    result.data[4] = data[4] * s;
+    result.data[5] = data[5] * s;
+
+    result.data[6] = data[6] * s;
+    result.data[7] = data[7] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+
+    data[2] *= s;
+    data[3] *= s;
+
+    data[4] *= s;
+    data[5] *= s;
+
+    data[6] *= s;
+    data[7] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+
+    result.data[2] = data[2] / rhs.data[2];
+    result.data[3] = data[3] / rhs.data[3];
+
+    result.data[4] = data[4] / rhs.data[4];
+    result.data[5] = data[5] / rhs.data[5];
+
+    result.data[6] = data[6] / rhs.data[6];
+    result.data[7] = data[7] / rhs.data[7];
+
+    return result;
+  }
+
+  /// Scalar divide operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+
+    result.data[2] = data[2] / s;
+    result.data[3] = data[3] / s;
+
+    result.data[4] = data[4] / s;
+    result.data[5] = data[5] / s;
+
+    result.data[6] = data[6] / s;
+    result.data[7] = data[7] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+
+    data[2] /= s;
+    data[3] /= s;
+
+    data[4] /= s;
+    data[5] /= s;
+
+    data[6] /= s;
+    data[7] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (4-by-2)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+
+    data[2] /= rhs.data[2];
+    data[3] /= rhs.data[3];
+
+    data[4] /= rhs.data[4];
+    data[5] /= rhs.data[5];
+
+    data[6] /= rhs.data[6];
+    data[7] /= rhs.data[7];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+    m.data[4] = -m.data[4];
+    m.data[5] = -m.data[5];
+    m.data[6] = -m.data[6];
+    m.data[7] = -m.data[7];
+
+    return m;
+  }
+
+  /// Matrix product of size 4-by-1-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> product(
+    Matrix<Element, 2, 1> const &rhs,
+    Matrix<Element, 4, 1> accum = Matrix<Element, 4, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[2] * rhs.data[0];
+    accum.data[2] += data[4] * rhs.data[0];
+    accum.data[3] += data[6] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[3] * rhs.data[1];
+    accum.data[2] += data[5] * rhs.data[1];
+    accum.data[3] += data[7] * rhs.data[1];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-1-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> operator*(Matrix<Element, 2, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> product(
+    Matrix<Element, 2, 2> const &rhs,
+    Matrix<Element, 4, 2> accum = Matrix<Element, 4, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[2] * rhs.data[0];
+    accum.data[3] += data[2] * rhs.data[1];
+    accum.data[4] += data[4] * rhs.data[0];
+    accum.data[5] += data[4] * rhs.data[1];
+    accum.data[6] += data[6] * rhs.data[0];
+    accum.data[7] += data[6] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[3] * rhs.data[2];
+    accum.data[3] += data[3] * rhs.data[3];
+    accum.data[4] += data[5] * rhs.data[2];
+    accum.data[5] += data[5] * rhs.data[3];
+    accum.data[6] += data[7] * rhs.data[2];
+    accum.data[7] += data[7] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> operator*(Matrix<Element, 2, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-2-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 2, 2> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 4-by-3-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> product(
+    Matrix<Element, 2, 3> const &rhs,
+    Matrix<Element, 4, 3> accum = Matrix<Element, 4, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[2] * rhs.data[0];
+    accum.data[4] += data[2] * rhs.data[1];
+    accum.data[5] += data[2] * rhs.data[2];
+    accum.data[6] += data[4] * rhs.data[0];
+    accum.data[7] += data[4] * rhs.data[1];
+    accum.data[8] += data[4] * rhs.data[2];
+    accum.data[9] += data[6] * rhs.data[0];
+    accum.data[10] += data[6] * rhs.data[1];
+    accum.data[11] += data[6] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[3] * rhs.data[3];
+    accum.data[4] += data[3] * rhs.data[4];
+    accum.data[5] += data[3] * rhs.data[5];
+    accum.data[6] += data[5] * rhs.data[3];
+    accum.data[7] += data[5] * rhs.data[4];
+    accum.data[8] += data[5] * rhs.data[5];
+    accum.data[9] += data[7] * rhs.data[3];
+    accum.data[10] += data[7] * rhs.data[4];
+    accum.data[11] += data[7] * rhs.data[5];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-3-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> operator*(Matrix<Element, 2, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-4-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> product(
+    Matrix<Element, 2, 4> const &rhs,
+    Matrix<Element, 4, 4> accum = Matrix<Element, 4, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[2] * rhs.data[0];
+    accum.data[5] += data[2] * rhs.data[1];
+    accum.data[6] += data[2] * rhs.data[2];
+    accum.data[7] += data[2] * rhs.data[3];
+    accum.data[8] += data[4] * rhs.data[0];
+    accum.data[9] += data[4] * rhs.data[1];
+    accum.data[10] += data[4] * rhs.data[2];
+    accum.data[11] += data[4] * rhs.data[3];
+    accum.data[12] += data[6] * rhs.data[0];
+    accum.data[13] += data[6] * rhs.data[1];
+    accum.data[14] += data[6] * rhs.data[2];
+    accum.data[15] += data[6] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[3] * rhs.data[4];
+    accum.data[5] += data[3] * rhs.data[5];
+    accum.data[6] += data[3] * rhs.data[6];
+    accum.data[7] += data[3] * rhs.data[7];
+    accum.data[8] += data[5] * rhs.data[4];
+    accum.data[9] += data[5] * rhs.data[5];
+    accum.data[10] += data[5] * rhs.data[6];
+    accum.data[11] += data[5] * rhs.data[7];
+    accum.data[12] += data[7] * rhs.data[4];
+    accum.data[13] += data[7] * rhs.data[5];
+    accum.data[14] += data[7] * rhs.data[6];
+    accum.data[15] += data[7] * rhs.data[7];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-4-by-2
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> operator*(Matrix<Element, 2, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+    accum += data[4];
+    accum += data[5];
+    accum += data[6];
+    accum += data[7];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+    accum += data[4] * data[4];
+    accum += data[5] * data[5];
+    accum += data[6] * data[6];
+    accum += data[7] * data[7];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[3];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 4-by-2 matrix
+template <typename Element>
+using Matrix4x2 = Matrix<Element, 4, 2>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix4x2<Element> make_Matrix4x2(
+    Element _0_0, Element _0_1,
+    Element _1_0, Element _1_1,
+    Element _2_0, Element _2_1,
+    Element _3_0, Element _3_1
+) {
+  return Matrix4x2<Element>(
+  _0_0, _0_1,
+  _1_0, _1_1,
+  _2_0, _2_1,
+  _3_0, _3_1
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 4-by-3 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 4, 3> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 4;
+
+  /// Number of columns in matrix
+  static int const kColumns = 3;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 12;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 4-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 4-by-3 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1, Element _0_2,
+    Element _1_0, Element _1_1, Element _1_2,
+    Element _2_0, Element _2_1, Element _2_2,
+    Element _3_0, Element _3_1, Element _3_2
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;  data[2] = _0_2;
+    data[3] = _1_0;  data[4] = _1_1;  data[5] = _1_2;
+    data[6] = _2_0;  data[7] = _2_1;  data[8] = _2_2;
+    data[9] = _3_0;  data[10] = _3_1;  data[11] = _3_2;
+  }
+
+  /// Constucts a 4-by-3 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 3> const &row_0,
+    Matrix<Element, 1, 3> const &row_1,
+    Matrix<Element, 1, 3> const &row_2,
+    Matrix<Element, 1, 3> const &row_3
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_0.data[2];
+    data[3] = row_1.data[0];
+    data[4] = row_1.data[1];
+    data[5] = row_1.data[2];
+    data[6] = row_2.data[0];
+    data[7] = row_2.data[1];
+    data[8] = row_2.data[2];
+    data[9] = row_3.data[0];
+    data[10] = row_3.data[1];
+    data[11] = row_3.data[2];
+  }
+
+  /// Static method to construct a 4-by-3 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 3, 1> const &column_0,
+    Matrix<Element, 3, 1> const &column_1,
+    Matrix<Element, 3, 1> const &column_2
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_2.data[0];
+    result.data[3] = column_0.data[1];
+    result.data[4] = column_1.data[1];
+    result.data[5] = column_2.data[1];
+    result.data[6] = column_0.data[2];
+    result.data[7] = column_1.data[2];
+    result.data[8] = column_2.data[2];
+    result.data[9] = column_0.data[3];
+    result.data[10] = column_1.data[3];
+    result.data[11] = column_2.data[3];
+    return result;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+    m.data[4] = s;
+    m.data[5] = s;
+    m.data[6] = s;
+    m.data[7] = s;
+    m.data[8] = s;
+    m.data[9] = s;
+    m.data[10] = s;
+    m.data[11] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 3, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[5] = diag.data[1];
+    m.data[10] = diag.data[2];
+    m.data[15] = diag.data[3];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 3> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[5] = diag.data[1];
+    m.data[10] = diag.data[2];
+    m.data[15] = diag.data[3];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> diagonal() const {
+    Matrix<Element, 3, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[5];
+    diag.data[2] = data[10];
+    diag.data[3] = data[15];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> transpose() const {
+    Matrix<Element, 3, 4> mt;
+
+    mt.data[0] = data[0];
+    mt.data[4] = data[1];
+    mt.data[8] = data[2];
+    mt.data[1] = data[3];
+    mt.data[5] = data[4];
+    mt.data[9] = data[5];
+    mt.data[2] = data[6];
+    mt.data[6] = data[7];
+    mt.data[10] = data[8];
+    mt.data[3] = data[9];
+    mt.data[7] = data[10];
+    mt.data[11] = data[11];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 4 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 4 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> slice_1x3(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x3(Matrix<Element, 1, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> row(int i) const {
+    return slice_1x3(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 3> const &v, int i = 0) {
+    return set_slice_1x3(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 3] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 3];
+    m.data[3] = data[i * 3 + j + 4];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 3] = m.data[2];
+    data[i * 3 + j + 4] = m.data[3];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> slice_2x3(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+    m.data[3] = data[i * 3 + j + 3];
+    m.data[4] = data[i * 3 + j + 4];
+    m.data[5] = data[i * 3 + j + 5];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x3(Matrix<Element, 2, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+    data[i * 3 + j + 3] = m.data[3];
+    data[i * 3 + j + 4] = m.data[4];
+    data[i * 3 + j + 5] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> slice_3x1(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 1> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 3];
+    m.data[2] = data[i * 3 + j + 6];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x1(Matrix<Element, 3, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 3] = m.data[1];
+    data[i * 3 + j + 6] = m.data[2];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> slice_3x2(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 3];
+    m.data[3] = data[i * 3 + j + 4];
+    m.data[4] = data[i * 3 + j + 6];
+    m.data[5] = data[i * 3 + j + 7];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x2(Matrix<Element, 3, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 3] = m.data[2];
+    data[i * 3 + j + 4] = m.data[3];
+    data[i * 3 + j + 6] = m.data[4];
+    data[i * 3 + j + 7] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> slice_3x3(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+    m.data[3] = data[i * 3 + j + 3];
+    m.data[4] = data[i * 3 + j + 4];
+    m.data[5] = data[i * 3 + j + 5];
+    m.data[6] = data[i * 3 + j + 6];
+    m.data[7] = data[i * 3 + j + 7];
+    m.data[8] = data[i * 3 + j + 8];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x3(Matrix<Element, 3, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+    data[i * 3 + j + 3] = m.data[3];
+    data[i * 3 + j + 4] = m.data[4];
+    data[i * 3 + j + 5] = m.data[5];
+    data[i * 3 + j + 6] = m.data[6];
+    data[i * 3 + j + 7] = m.data[7];
+    data[i * 3 + j + 8] = m.data[8];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> slice_4x1(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 1> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 3];
+    m.data[2] = data[i * 3 + j + 6];
+    m.data[3] = data[i * 3 + j + 9];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x1(Matrix<Element, 4, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 3] = m.data[1];
+    data[i * 3 + j + 6] = m.data[2];
+    data[i * 3 + j + 9] = m.data[3];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> column(int j) const {
+    return slice_4x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 4, 1> const &v, int j =0) {
+    return set_slice_4x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> slice_4x2(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 2> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 3];
+    m.data[3] = data[i * 3 + j + 4];
+    m.data[4] = data[i * 3 + j + 6];
+    m.data[5] = data[i * 3 + j + 7];
+    m.data[6] = data[i * 3 + j + 9];
+    m.data[7] = data[i * 3 + j + 10];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x2(Matrix<Element, 4, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 3] = m.data[2];
+    data[i * 3 + j + 4] = m.data[3];
+    data[i * 3 + j + 6] = m.data[4];
+    data[i * 3 + j + 7] = m.data[5];
+    data[i * 3 + j + 9] = m.data[6];
+    data[i * 3 + j + 10] = m.data[7];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> slice_4x3(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 3> m;
+
+    m.data[0] = data[i * 3 + j + 0];
+    m.data[1] = data[i * 3 + j + 1];
+    m.data[2] = data[i * 3 + j + 2];
+    m.data[3] = data[i * 3 + j + 3];
+    m.data[4] = data[i * 3 + j + 4];
+    m.data[5] = data[i * 3 + j + 5];
+    m.data[6] = data[i * 3 + j + 6];
+    m.data[7] = data[i * 3 + j + 7];
+    m.data[8] = data[i * 3 + j + 8];
+    m.data[9] = data[i * 3 + j + 9];
+    m.data[10] = data[i * 3 + j + 10];
+    m.data[11] = data[i * 3 + j + 11];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x3(Matrix<Element, 4, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 3 + j + 0] = m.data[0];
+    data[i * 3 + j + 1] = m.data[1];
+    data[i * 3 + j + 2] = m.data[2];
+    data[i * 3 + j + 3] = m.data[3];
+    data[i * 3 + j + 4] = m.data[4];
+    data[i * 3 + j + 5] = m.data[5];
+    data[i * 3 + j + 6] = m.data[6];
+    data[i * 3 + j + 7] = m.data[7];
+    data[i * 3 + j + 8] = m.data[8];
+    data[i * 3 + j + 9] = m.data[9];
+    data[i * 3 + j + 10] = m.data[10];
+    data[i * 3 + j + 11] = m.data[11];
+
+    return *this;
+  }
+
+  /// Forms a 4-by-3 matrix by horizontally concatenating a 4-by-1 matrix with a 4-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 4, 1> const & lhs, Matrix<Element, 4, 2> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1)
+      , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1)
+      , lhs.at(2, 0), rhs.at(2, 0), rhs.at(2, 1)
+      , lhs.at(3, 0), rhs.at(3, 0), rhs.at(3, 1));
+  }
+
+  /// Forms a 4-by-3 matrix by horizontally concatenating a 4-by-2 matrix with a 4-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 4, 2> const & lhs, Matrix<Element, 4, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0)
+      , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0)
+      , lhs.at(2, 0), lhs.at(2, 1), rhs.at(2, 0)
+      , lhs.at(3, 0), lhs.at(3, 1), rhs.at(3, 0));
+  }
+
+  /// Concatenates this matrix with a a 4-by-1 matrix to form a 4-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> hcat(Matrix<Element, 4, 1> const & rhs) const {
+    return Matrix<Element, 4, 4>::hcat(*this, rhs);
+  }
+
+  /// Forms a 4-by-3 matrix by vertically concatenating a 1-by-3 matrix with a 3-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 3> const & upper, Matrix<Element, 3, 3> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2)
+      , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2)
+      , lower.at(2, 0), lower.at(2, 1), lower.at(2, 2));
+  }
+
+  /// Forms a 4-by-3 matrix by vertically concatenating a 2-by-3 matrix with a 2-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 2, 3> const & upper, Matrix<Element, 2, 3> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2)
+      , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2)
+      , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2));
+  }
+
+  /// Forms a 4-by-3 matrix by vertically concatenating a 3-by-3 matrix with a 1-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 3, 3> const & upper, Matrix<Element, 1, 3> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2)
+      , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2)
+      , upper.at(2, 0), upper.at(2, 1), upper.at(2, 2)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2));
+  }
+
+  /// Forms a 4-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Matrix<Element, 1, 2> const & B,
+    Matrix<Element, 3, 1> const & C, Matrix<Element, 3, 2> const & D) {
+    return Matrix(
+      A, B.at(0, 0), B.at(0, 1)
+      , C.at(0, 0), D.at(0, 0), D.at(0, 1)
+      , C.at(1, 0), D.at(1, 0), D.at(1, 1)
+      , C.at(2, 0), D.at(2, 0), D.at(2, 1)
+    );
+  }
+
+  /// Forms a 4-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 2> const & A, Element                         B,
+    Matrix<Element, 3, 2> const & C, Matrix<Element, 3, 1> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0)
+      , C.at(1, 0), C.at(1, 1), D.at(1, 0)
+      , C.at(2, 0), C.at(2, 1), D.at(2, 0)
+    );
+  }
+
+  /// Forms a 4-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 1> const & A, Matrix<Element, 2, 2> const & B,
+    Matrix<Element, 2, 1> const & C, Matrix<Element, 2, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0), B.at(0, 1)
+      , A.at(1, 0), B.at(1, 0), B.at(1, 1)
+      , C.at(0, 0), D.at(0, 0), D.at(0, 1)
+      , C.at(1, 0), D.at(1, 0), D.at(1, 1)
+    );
+  }
+
+  /// Forms a 4-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 2> const & A, Matrix<Element, 2, 1> const & B,
+    Matrix<Element, 2, 2> const & C, Matrix<Element, 2, 1> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0)
+      , A.at(1, 0), A.at(1, 1), B.at(1, 0)
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0)
+      , C.at(1, 0), C.at(1, 1), D.at(1, 0)
+    );
+  }
+
+  /// Forms a 4-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 3, 1> const & A, Matrix<Element, 3, 2> const & B,
+    Element                         C, Matrix<Element, 1, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0), B.at(0, 1)
+      , A.at(1, 0), B.at(1, 0), B.at(1, 1)
+      , A.at(2, 0), B.at(2, 0), B.at(2, 1)
+      , C, D.at(0, 0), D.at(0, 1)
+    );
+  }
+
+  /// Forms a 4-by-3 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 3, 2> const & A, Matrix<Element, 3, 1> const & B,
+    Matrix<Element, 1, 2> const & C, Element                         D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0)
+      , A.at(1, 0), A.at(1, 1), B.at(1, 0)
+      , A.at(2, 0), A.at(2, 1), B.at(2, 0)
+      , C.at(0, 0), C.at(0, 1), D
+    );
+  }
+
+  /// Elementwise add operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+    result.data[2] = data[2] + rhs.data[2];
+
+    result.data[3] = data[3] + rhs.data[3];
+    result.data[4] = data[4] + rhs.data[4];
+    result.data[5] = data[5] + rhs.data[5];
+
+    result.data[6] = data[6] + rhs.data[6];
+    result.data[7] = data[7] + rhs.data[7];
+    result.data[8] = data[8] + rhs.data[8];
+
+    result.data[9] = data[9] + rhs.data[9];
+    result.data[10] = data[10] + rhs.data[10];
+    result.data[11] = data[11] + rhs.data[11];
+
+    return result;
+  }
+
+  /// Elementwise add operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+    data[2] += rhs.data[2];
+
+    data[3] += rhs.data[3];
+    data[4] += rhs.data[4];
+    data[5] += rhs.data[5];
+
+    data[6] += rhs.data[6];
+    data[7] += rhs.data[7];
+    data[8] += rhs.data[8];
+
+    data[9] += rhs.data[9];
+    data[10] += rhs.data[10];
+    data[11] += rhs.data[11];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+    result.data[2] = data[2] - rhs.data[2];
+
+    result.data[3] = data[3] - rhs.data[3];
+    result.data[4] = data[4] - rhs.data[4];
+    result.data[5] = data[5] - rhs.data[5];
+
+    result.data[6] = data[6] - rhs.data[6];
+    result.data[7] = data[7] - rhs.data[7];
+    result.data[8] = data[8] - rhs.data[8];
+
+    result.data[9] = data[9] - rhs.data[9];
+    result.data[10] = data[10] - rhs.data[10];
+    result.data[11] = data[11] - rhs.data[11];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+    data[2] -= rhs.data[2];
+
+    data[3] -= rhs.data[3];
+    data[4] -= rhs.data[4];
+    data[5] -= rhs.data[5];
+
+    data[6] -= rhs.data[6];
+    data[7] -= rhs.data[7];
+    data[8] -= rhs.data[8];
+
+    data[9] -= rhs.data[9];
+    data[10] -= rhs.data[10];
+    data[11] -= rhs.data[11];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+    result.data[2] = data[2] * rhs.data[2];
+
+    result.data[3] = data[3] * rhs.data[3];
+    result.data[4] = data[4] * rhs.data[4];
+    result.data[5] = data[5] * rhs.data[5];
+
+    result.data[6] = data[6] * rhs.data[6];
+    result.data[7] = data[7] * rhs.data[7];
+    result.data[8] = data[8] * rhs.data[8];
+
+    result.data[9] = data[9] * rhs.data[9];
+    result.data[10] = data[10] * rhs.data[10];
+    result.data[11] = data[11] * rhs.data[11];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+    result.data[2] = data[2] * s;
+
+    result.data[3] = data[3] * s;
+    result.data[4] = data[4] * s;
+    result.data[5] = data[5] * s;
+
+    result.data[6] = data[6] * s;
+    result.data[7] = data[7] * s;
+    result.data[8] = data[8] * s;
+
+    result.data[9] = data[9] * s;
+    result.data[10] = data[10] * s;
+    result.data[11] = data[11] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+    data[2] *= s;
+
+    data[3] *= s;
+    data[4] *= s;
+    data[5] *= s;
+
+    data[6] *= s;
+    data[7] *= s;
+    data[8] *= s;
+
+    data[9] *= s;
+    data[10] *= s;
+    data[11] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+    result.data[2] = data[2] / rhs.data[2];
+
+    result.data[3] = data[3] / rhs.data[3];
+    result.data[4] = data[4] / rhs.data[4];
+    result.data[5] = data[5] / rhs.data[5];
+
+    result.data[6] = data[6] / rhs.data[6];
+    result.data[7] = data[7] / rhs.data[7];
+    result.data[8] = data[8] / rhs.data[8];
+
+    result.data[9] = data[9] / rhs.data[9];
+    result.data[10] = data[10] / rhs.data[10];
+    result.data[11] = data[11] / rhs.data[11];
+
+    return result;
+  }
+
+  /// Scalar divide operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+    result.data[2] = data[2] / s;
+
+    result.data[3] = data[3] / s;
+    result.data[4] = data[4] / s;
+    result.data[5] = data[5] / s;
+
+    result.data[6] = data[6] / s;
+    result.data[7] = data[7] / s;
+    result.data[8] = data[8] / s;
+
+    result.data[9] = data[9] / s;
+    result.data[10] = data[10] / s;
+    result.data[11] = data[11] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+    data[2] /= s;
+
+    data[3] /= s;
+    data[4] /= s;
+    data[5] /= s;
+
+    data[6] /= s;
+    data[7] /= s;
+    data[8] /= s;
+
+    data[9] /= s;
+    data[10] /= s;
+    data[11] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (4-by-3)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+    data[2] /= rhs.data[2];
+
+    data[3] /= rhs.data[3];
+    data[4] /= rhs.data[4];
+    data[5] /= rhs.data[5];
+
+    data[6] /= rhs.data[6];
+    data[7] /= rhs.data[7];
+    data[8] /= rhs.data[8];
+
+    data[9] /= rhs.data[9];
+    data[10] /= rhs.data[10];
+    data[11] /= rhs.data[11];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+    m.data[4] = -m.data[4];
+    m.data[5] = -m.data[5];
+    m.data[6] = -m.data[6];
+    m.data[7] = -m.data[7];
+    m.data[8] = -m.data[8];
+    m.data[9] = -m.data[9];
+    m.data[10] = -m.data[10];
+    m.data[11] = -m.data[11];
+
+    return m;
+  }
+
+  /// Matrix product of size 4-by-1-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> product(
+    Matrix<Element, 3, 1> const &rhs,
+    Matrix<Element, 4, 1> accum = Matrix<Element, 4, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[3] * rhs.data[0];
+    accum.data[2] += data[6] * rhs.data[0];
+    accum.data[3] += data[9] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[4] * rhs.data[1];
+    accum.data[2] += data[7] * rhs.data[1];
+    accum.data[3] += data[10] * rhs.data[1];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[2];
+    accum.data[1] += data[5] * rhs.data[2];
+    accum.data[2] += data[8] * rhs.data[2];
+    accum.data[3] += data[11] * rhs.data[2];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-1-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> operator*(Matrix<Element, 3, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-2-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> product(
+    Matrix<Element, 3, 2> const &rhs,
+    Matrix<Element, 4, 2> accum = Matrix<Element, 4, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[3] * rhs.data[0];
+    accum.data[3] += data[3] * rhs.data[1];
+    accum.data[4] += data[6] * rhs.data[0];
+    accum.data[5] += data[6] * rhs.data[1];
+    accum.data[6] += data[9] * rhs.data[0];
+    accum.data[7] += data[9] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[4] * rhs.data[2];
+    accum.data[3] += data[4] * rhs.data[3];
+    accum.data[4] += data[7] * rhs.data[2];
+    accum.data[5] += data[7] * rhs.data[3];
+    accum.data[6] += data[10] * rhs.data[2];
+    accum.data[7] += data[10] * rhs.data[3];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[4];
+    accum.data[1] += data[2] * rhs.data[5];
+    accum.data[2] += data[5] * rhs.data[4];
+    accum.data[3] += data[5] * rhs.data[5];
+    accum.data[4] += data[8] * rhs.data[4];
+    accum.data[5] += data[8] * rhs.data[5];
+    accum.data[6] += data[11] * rhs.data[4];
+    accum.data[7] += data[11] * rhs.data[5];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-2-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> operator*(Matrix<Element, 3, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> product(
+    Matrix<Element, 3, 3> const &rhs,
+    Matrix<Element, 4, 3> accum = Matrix<Element, 4, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[3] * rhs.data[0];
+    accum.data[4] += data[3] * rhs.data[1];
+    accum.data[5] += data[3] * rhs.data[2];
+    accum.data[6] += data[6] * rhs.data[0];
+    accum.data[7] += data[6] * rhs.data[1];
+    accum.data[8] += data[6] * rhs.data[2];
+    accum.data[9] += data[9] * rhs.data[0];
+    accum.data[10] += data[9] * rhs.data[1];
+    accum.data[11] += data[9] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[4] * rhs.data[3];
+    accum.data[4] += data[4] * rhs.data[4];
+    accum.data[5] += data[4] * rhs.data[5];
+    accum.data[6] += data[7] * rhs.data[3];
+    accum.data[7] += data[7] * rhs.data[4];
+    accum.data[8] += data[7] * rhs.data[5];
+    accum.data[9] += data[10] * rhs.data[3];
+    accum.data[10] += data[10] * rhs.data[4];
+    accum.data[11] += data[10] * rhs.data[5];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[6];
+    accum.data[1] += data[2] * rhs.data[7];
+    accum.data[2] += data[2] * rhs.data[8];
+    accum.data[3] += data[5] * rhs.data[6];
+    accum.data[4] += data[5] * rhs.data[7];
+    accum.data[5] += data[5] * rhs.data[8];
+    accum.data[6] += data[8] * rhs.data[6];
+    accum.data[7] += data[8] * rhs.data[7];
+    accum.data[8] += data[8] * rhs.data[8];
+    accum.data[9] += data[11] * rhs.data[6];
+    accum.data[10] += data[11] * rhs.data[7];
+    accum.data[11] += data[11] * rhs.data[8];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> operator*(Matrix<Element, 3, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-3-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 3, 3> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Matrix product of size 4-by-4-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> product(
+    Matrix<Element, 3, 4> const &rhs,
+    Matrix<Element, 4, 4> accum = Matrix<Element, 4, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[3] * rhs.data[0];
+    accum.data[5] += data[3] * rhs.data[1];
+    accum.data[6] += data[3] * rhs.data[2];
+    accum.data[7] += data[3] * rhs.data[3];
+    accum.data[8] += data[6] * rhs.data[0];
+    accum.data[9] += data[6] * rhs.data[1];
+    accum.data[10] += data[6] * rhs.data[2];
+    accum.data[11] += data[6] * rhs.data[3];
+    accum.data[12] += data[9] * rhs.data[0];
+    accum.data[13] += data[9] * rhs.data[1];
+    accum.data[14] += data[9] * rhs.data[2];
+    accum.data[15] += data[9] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[4] * rhs.data[4];
+    accum.data[5] += data[4] * rhs.data[5];
+    accum.data[6] += data[4] * rhs.data[6];
+    accum.data[7] += data[4] * rhs.data[7];
+    accum.data[8] += data[7] * rhs.data[4];
+    accum.data[9] += data[7] * rhs.data[5];
+    accum.data[10] += data[7] * rhs.data[6];
+    accum.data[11] += data[7] * rhs.data[7];
+    accum.data[12] += data[10] * rhs.data[4];
+    accum.data[13] += data[10] * rhs.data[5];
+    accum.data[14] += data[10] * rhs.data[6];
+    accum.data[15] += data[10] * rhs.data[7];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[8];
+    accum.data[1] += data[2] * rhs.data[9];
+    accum.data[2] += data[2] * rhs.data[10];
+    accum.data[3] += data[2] * rhs.data[11];
+    accum.data[4] += data[5] * rhs.data[8];
+    accum.data[5] += data[5] * rhs.data[9];
+    accum.data[6] += data[5] * rhs.data[10];
+    accum.data[7] += data[5] * rhs.data[11];
+    accum.data[8] += data[8] * rhs.data[8];
+    accum.data[9] += data[8] * rhs.data[9];
+    accum.data[10] += data[8] * rhs.data[10];
+    accum.data[11] += data[8] * rhs.data[11];
+    accum.data[12] += data[11] * rhs.data[8];
+    accum.data[13] += data[11] * rhs.data[9];
+    accum.data[14] += data[11] * rhs.data[10];
+    accum.data[15] += data[11] * rhs.data[11];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-4-by-3
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> operator*(Matrix<Element, 3, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+    accum += data[4];
+    accum += data[5];
+    accum += data[6];
+    accum += data[7];
+    accum += data[8];
+    accum += data[9];
+    accum += data[10];
+    accum += data[11];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+    accum += data[4] * data[4];
+    accum += data[5] * data[5];
+    accum += data[6] * data[6];
+    accum += data[7] * data[7];
+    accum += data[8] * data[8];
+    accum += data[9] * data[9];
+    accum += data[10] * data[10];
+    accum += data[11] * data[11];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[4];
+    accum += data[8];
+
+    return accum;
+  }
+
+};
+
+/// Template alias for 4-by-3 matrix
+template <typename Element>
+using Matrix4x3 = Matrix<Element, 4, 3>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix4x3<Element> make_Matrix4x3(
+    Element _0_0, Element _0_1, Element _0_2,
+    Element _1_0, Element _1_1, Element _1_2,
+    Element _2_0, Element _2_1, Element _2_2,
+    Element _3_0, Element _3_1, Element _3_2
+) {
+  return Matrix4x3<Element>(
+  _0_0, _0_1, _0_2,
+  _1_0, _1_1, _1_2,
+  _2_0, _2_1, _2_2,
+  _3_0, _3_1, _3_2
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// 4-by-4 matrix template class definition
+template <typename Element_>
+struct Matrix<Element_, 4, 4> {
+
+  //
+  // Type definitions
+  //
+
+  /// Element data type
+  using Element = Element_;
+
+  /// Number of rows in matrix
+  static int const kRows = 4;
+
+  /// Number of columns in matrix
+  static int const kColumns = 4;
+
+  /// Layout of matrix in underlying array
+  using Layout = layout::RowMajor;
+
+  /// Number of elements in matrix
+  static int const kCount = 16;
+
+  //
+  // Data members
+  //
+
+  /// Elements of the matrix in row-major layout
+  Array<Element, kCount> data;
+
+  //
+  // Methods
+  //
+
+  /// Constructs a zero matrix
+  MUTLASS_HOST_DEVICE
+  Matrix() {
+    data.clear();
+  }
+
+  /// Copy constructor for a 4-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Matrix(Matrix const &rhs) {
+    data = rhs.data;
+  }
+
+  /// Constucts a 4-by-4 matrix from scalar elements
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Element _0_0, Element _0_1, Element _0_2, Element _0_3,
+    Element _1_0, Element _1_1, Element _1_2, Element _1_3,
+    Element _2_0, Element _2_1, Element _2_2, Element _2_3,
+    Element _3_0, Element _3_1, Element _3_2, Element _3_3
+  ) {
+
+    data[0] = _0_0;  data[1] = _0_1;  data[2] = _0_2;  data[3] = _0_3;
+    data[4] = _1_0;  data[5] = _1_1;  data[6] = _1_2;  data[7] = _1_3;
+    data[8] = _2_0;  data[9] = _2_1;  data[10] = _2_2;  data[11] = _2_3;
+    data[12] = _3_0;  data[13] = _3_1;  data[14] = _3_2;  data[15] = _3_3;
+  }
+
+  /// Constucts a 4-by-4 matrix from row vectors
+  MUTLASS_HOST_DEVICE
+  Matrix(
+    Matrix<Element, 1, 4> const &row_0,
+    Matrix<Element, 1, 4> const &row_1,
+    Matrix<Element, 1, 4> const &row_2,
+    Matrix<Element, 1, 4> const &row_3
+  ) {
+    data[0] = row_0.data[0];
+    data[1] = row_0.data[1];
+    data[2] = row_0.data[2];
+    data[3] = row_0.data[3];
+    data[4] = row_1.data[0];
+    data[5] = row_1.data[1];
+    data[6] = row_1.data[2];
+    data[7] = row_1.data[3];
+    data[8] = row_2.data[0];
+    data[9] = row_2.data[1];
+    data[10] = row_2.data[2];
+    data[11] = row_2.data[3];
+    data[12] = row_3.data[0];
+    data[13] = row_3.data[1];
+    data[14] = row_3.data[2];
+    data[15] = row_3.data[3];
+  }
+
+  /// Static method to construct a 4-by-4 matrix from column vectors
+  MUTLASS_HOST_DEVICE
+  static Matrix from_columns(
+    Matrix<Element, 4, 1> const &column_0,
+    Matrix<Element, 4, 1> const &column_1,
+    Matrix<Element, 4, 1> const &column_2,
+    Matrix<Element, 4, 1> const &column_3
+  ) {
+    Matrix result;
+
+    result.data[0] = column_0.data[0];
+    result.data[1] = column_1.data[0];
+    result.data[2] = column_2.data[0];
+    result.data[3] = column_3.data[0];
+    result.data[4] = column_0.data[1];
+    result.data[5] = column_1.data[1];
+    result.data[6] = column_2.data[1];
+    result.data[7] = column_3.data[1];
+    result.data[8] = column_0.data[2];
+    result.data[9] = column_1.data[2];
+    result.data[10] = column_2.data[2];
+    result.data[11] = column_3.data[2];
+    result.data[12] = column_0.data[3];
+    result.data[13] = column_1.data[3];
+    result.data[14] = column_2.data[3];
+    result.data[15] = column_3.data[3];
+    return result;
+  }
+
+  /// Constructs an identity matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix identity() {
+    Matrix m;
+
+    m.data[0] = Element(1);
+    m.data[5] = Element(1);
+    m.data[10] = Element(1);
+    m.data[15] = Element(1);
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element
+  MUTLASS_HOST_DEVICE
+  static Matrix uniform(Element s) {
+    Matrix m;
+
+    m.data[0] = s;
+    m.data[1] = s;
+    m.data[2] = s;
+    m.data[3] = s;
+    m.data[4] = s;
+    m.data[5] = s;
+    m.data[6] = s;
+    m.data[7] = s;
+    m.data[8] = s;
+    m.data[9] = s;
+    m.data[10] = s;
+    m.data[11] = s;
+    m.data[12] = s;
+    m.data[13] = s;
+    m.data[14] = s;
+    m.data[15] = s;
+
+    return m;
+  }
+
+  /// Constructs a matrix from a uniform element 1
+  MUTLASS_HOST_DEVICE
+  static Matrix ones() {
+    return uniform(Element(1));
+  }
+
+  /// Constructs a matrix from a uniform element 0
+  MUTLASS_HOST_DEVICE
+  static Matrix zero() {
+    return Matrix();
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 4, 1> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[5] = diag.data[1];
+    m.data[10] = diag.data[2];
+    m.data[15] = diag.data[3];
+
+    return m;
+  }
+
+  /// Constructs a matrix from elements along its diagonal
+  MUTLASS_HOST_DEVICE
+  static Matrix from_diagonal(Matrix<Element, 1, 4> const &diag) {
+    Matrix m;
+
+    m.data[0] = diag.data[0];
+    m.data[5] = diag.data[1];
+    m.data[10] = diag.data[2];
+    m.data[15] = diag.data[3];
+
+    return m;
+  }
+
+  /// Gets an array of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> diagonal() const {
+    Matrix<Element, 4, 1> diag;
+
+    diag.data[0] = data[0];
+    diag.data[1] = data[5];
+    diag.data[2] = data[10];
+    diag.data[3] = data[15];
+
+    return diag;
+  }
+
+  /// Returns a transposed matrix
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> transpose() const {
+    Matrix<Element, 4, 4> mt;
+
+    mt.data[0] = data[0];
+    mt.data[4] = data[1];
+    mt.data[8] = data[2];
+    mt.data[12] = data[3];
+    mt.data[1] = data[4];
+    mt.data[5] = data[5];
+    mt.data[9] = data[6];
+    mt.data[13] = data[7];
+    mt.data[2] = data[8];
+    mt.data[6] = data[9];
+    mt.data[10] = data[10];
+    mt.data[14] = data[11];
+    mt.data[3] = data[12];
+    mt.data[7] = data[13];
+    mt.data[11] = data[14];
+    mt.data[15] = data[15];
+
+    return mt;
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(int i, int j) const {
+    return data[i * 4 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(int i, int j) {
+    return data[i * 4 + j];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element at(Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & at(Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element &at(int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element at(int offset) const {
+    return data[offset];
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element operator[](Coord<2> const &coord) const {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by coordinate
+  MUTLASS_HOST_DEVICE
+  Element & operator[](Coord<2> const &coord) {
+    return at(coord[0], coord[1]);
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element & operator[](int offset) {
+    return data[offset];
+  }
+
+  /// Accesses an element by offset
+  MUTLASS_HOST_DEVICE
+  Element operator[](int offset) const {
+    return data[offset];
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 2> slice_1x2(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x2(Matrix<Element, 1, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 3> slice_1x3(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x3(Matrix<Element, 1, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> slice_1x4(int i = 0, int j = 0) const {
+    Matrix<Element, 1, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_1x4(Matrix<Element, 1, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 1, 4> row(int i) const {
+    return slice_1x4(i, 0);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_row(Matrix<Element, 1, 4> const &v, int i = 0) {
+    return set_slice_1x4(v, i, 0);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 1> slice_2x1(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 1> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 4];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x1(Matrix<Element, 2, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 4] = m.data[1];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 2> slice_2x2(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 4];
+    m.data[3] = data[i * 4 + j + 5];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x2(Matrix<Element, 2, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 4] = m.data[2];
+    data[i * 4 + j + 5] = m.data[3];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 3> slice_2x3(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 4];
+    m.data[4] = data[i * 4 + j + 5];
+    m.data[5] = data[i * 4 + j + 6];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x3(Matrix<Element, 2, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 4] = m.data[3];
+    data[i * 4 + j + 5] = m.data[4];
+    data[i * 4 + j + 6] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 2, 4> slice_2x4(int i = 0, int j = 0) const {
+    Matrix<Element, 2, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+    m.data[4] = data[i * 4 + j + 4];
+    m.data[5] = data[i * 4 + j + 5];
+    m.data[6] = data[i * 4 + j + 6];
+    m.data[7] = data[i * 4 + j + 7];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_2x4(Matrix<Element, 2, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+    data[i * 4 + j + 4] = m.data[4];
+    data[i * 4 + j + 5] = m.data[5];
+    data[i * 4 + j + 6] = m.data[6];
+    data[i * 4 + j + 7] = m.data[7];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 1> slice_3x1(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 1> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 4];
+    m.data[2] = data[i * 4 + j + 8];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x1(Matrix<Element, 3, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 4] = m.data[1];
+    data[i * 4 + j + 8] = m.data[2];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 2> slice_3x2(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 4];
+    m.data[3] = data[i * 4 + j + 5];
+    m.data[4] = data[i * 4 + j + 8];
+    m.data[5] = data[i * 4 + j + 9];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x2(Matrix<Element, 3, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 4] = m.data[2];
+    data[i * 4 + j + 5] = m.data[3];
+    data[i * 4 + j + 8] = m.data[4];
+    data[i * 4 + j + 9] = m.data[5];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 3> slice_3x3(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 4];
+    m.data[4] = data[i * 4 + j + 5];
+    m.data[5] = data[i * 4 + j + 6];
+    m.data[6] = data[i * 4 + j + 8];
+    m.data[7] = data[i * 4 + j + 9];
+    m.data[8] = data[i * 4 + j + 10];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x3(Matrix<Element, 3, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 4] = m.data[3];
+    data[i * 4 + j + 5] = m.data[4];
+    data[i * 4 + j + 6] = m.data[5];
+    data[i * 4 + j + 8] = m.data[6];
+    data[i * 4 + j + 9] = m.data[7];
+    data[i * 4 + j + 10] = m.data[8];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 3, 4> slice_3x4(int i = 0, int j = 0) const {
+    Matrix<Element, 3, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+    m.data[4] = data[i * 4 + j + 4];
+    m.data[5] = data[i * 4 + j + 5];
+    m.data[6] = data[i * 4 + j + 6];
+    m.data[7] = data[i * 4 + j + 7];
+    m.data[8] = data[i * 4 + j + 8];
+    m.data[9] = data[i * 4 + j + 9];
+    m.data[10] = data[i * 4 + j + 10];
+    m.data[11] = data[i * 4 + j + 11];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_3x4(Matrix<Element, 3, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+    data[i * 4 + j + 4] = m.data[4];
+    data[i * 4 + j + 5] = m.data[5];
+    data[i * 4 + j + 6] = m.data[6];
+    data[i * 4 + j + 7] = m.data[7];
+    data[i * 4 + j + 8] = m.data[8];
+    data[i * 4 + j + 9] = m.data[9];
+    data[i * 4 + j + 10] = m.data[10];
+    data[i * 4 + j + 11] = m.data[11];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> slice_4x1(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 1> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 4];
+    m.data[2] = data[i * 4 + j + 8];
+    m.data[3] = data[i * 4 + j + 12];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x1(Matrix<Element, 4, 1> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 4] = m.data[1];
+    data[i * 4 + j + 8] = m.data[2];
+    data[i * 4 + j + 12] = m.data[3];
+
+    return *this;
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> column(int j) const {
+    return slice_4x1(0, j);
+  }
+
+  MUTLASS_HOST_DEVICE
+  Matrix &set_column(Matrix<Element, 4, 1> const &v, int j =0) {
+    return set_slice_4x1(v, 0, j);
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> slice_4x2(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 2> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 4];
+    m.data[3] = data[i * 4 + j + 5];
+    m.data[4] = data[i * 4 + j + 8];
+    m.data[5] = data[i * 4 + j + 9];
+    m.data[6] = data[i * 4 + j + 12];
+    m.data[7] = data[i * 4 + j + 13];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x2(Matrix<Element, 4, 2> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 4] = m.data[2];
+    data[i * 4 + j + 5] = m.data[3];
+    data[i * 4 + j + 8] = m.data[4];
+    data[i * 4 + j + 9] = m.data[5];
+    data[i * 4 + j + 12] = m.data[6];
+    data[i * 4 + j + 13] = m.data[7];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> slice_4x3(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 3> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 4];
+    m.data[4] = data[i * 4 + j + 5];
+    m.data[5] = data[i * 4 + j + 6];
+    m.data[6] = data[i * 4 + j + 8];
+    m.data[7] = data[i * 4 + j + 9];
+    m.data[8] = data[i * 4 + j + 10];
+    m.data[9] = data[i * 4 + j + 12];
+    m.data[10] = data[i * 4 + j + 13];
+    m.data[11] = data[i * 4 + j + 14];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x3(Matrix<Element, 4, 3> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 4] = m.data[3];
+    data[i * 4 + j + 5] = m.data[4];
+    data[i * 4 + j + 6] = m.data[5];
+    data[i * 4 + j + 8] = m.data[6];
+    data[i * 4 + j + 9] = m.data[7];
+    data[i * 4 + j + 10] = m.data[8];
+    data[i * 4 + j + 12] = m.data[9];
+    data[i * 4 + j + 13] = m.data[10];
+    data[i * 4 + j + 14] = m.data[11];
+
+    return *this;
+  }
+
+  /// Gets a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> slice_4x4(int i = 0, int j = 0) const {
+    Matrix<Element, 4, 4> m;
+
+    m.data[0] = data[i * 4 + j + 0];
+    m.data[1] = data[i * 4 + j + 1];
+    m.data[2] = data[i * 4 + j + 2];
+    m.data[3] = data[i * 4 + j + 3];
+    m.data[4] = data[i * 4 + j + 4];
+    m.data[5] = data[i * 4 + j + 5];
+    m.data[6] = data[i * 4 + j + 6];
+    m.data[7] = data[i * 4 + j + 7];
+    m.data[8] = data[i * 4 + j + 8];
+    m.data[9] = data[i * 4 + j + 9];
+    m.data[10] = data[i * 4 + j + 10];
+    m.data[11] = data[i * 4 + j + 11];
+    m.data[12] = data[i * 4 + j + 12];
+    m.data[13] = data[i * 4 + j + 13];
+    m.data[14] = data[i * 4 + j + 14];
+    m.data[15] = data[i * 4 + j + 15];
+
+    return m;
+  }
+
+  /// Overwrites a submatrix with optional offset
+  MUTLASS_HOST_DEVICE
+  Matrix & set_slice_4x4(Matrix<Element, 4, 4> const &m, int i = 0, int j = 0) {
+
+    data[i * 4 + j + 0] = m.data[0];
+    data[i * 4 + j + 1] = m.data[1];
+    data[i * 4 + j + 2] = m.data[2];
+    data[i * 4 + j + 3] = m.data[3];
+    data[i * 4 + j + 4] = m.data[4];
+    data[i * 4 + j + 5] = m.data[5];
+    data[i * 4 + j + 6] = m.data[6];
+    data[i * 4 + j + 7] = m.data[7];
+    data[i * 4 + j + 8] = m.data[8];
+    data[i * 4 + j + 9] = m.data[9];
+    data[i * 4 + j + 10] = m.data[10];
+    data[i * 4 + j + 11] = m.data[11];
+    data[i * 4 + j + 12] = m.data[12];
+    data[i * 4 + j + 13] = m.data[13];
+    data[i * 4 + j + 14] = m.data[14];
+    data[i * 4 + j + 15] = m.data[15];
+
+    return *this;
+  }
+
+  /// Forms a 4-by-4 matrix by horizontally concatenating a 4-by-1 matrix with a 4-by-3 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 4, 1> const & lhs, Matrix<Element, 4, 3> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), rhs.at(0, 0), rhs.at(0, 1), rhs.at(0, 2)
+      , lhs.at(1, 0), rhs.at(1, 0), rhs.at(1, 1), rhs.at(1, 2)
+      , lhs.at(2, 0), rhs.at(2, 0), rhs.at(2, 1), rhs.at(2, 2)
+      , lhs.at(3, 0), rhs.at(3, 0), rhs.at(3, 1), rhs.at(3, 2));
+  }
+
+  /// Forms a 4-by-4 matrix by horizontally concatenating a 4-by-2 matrix with a 4-by-2 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 4, 2> const & lhs, Matrix<Element, 4, 2> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), rhs.at(0, 0), rhs.at(0, 1)
+      , lhs.at(1, 0), lhs.at(1, 1), rhs.at(1, 0), rhs.at(1, 1)
+      , lhs.at(2, 0), lhs.at(2, 1), rhs.at(2, 0), rhs.at(2, 1)
+      , lhs.at(3, 0), lhs.at(3, 1), rhs.at(3, 0), rhs.at(3, 1));
+  }
+
+  /// Forms a 4-by-4 matrix by horizontally concatenating a 4-by-3 matrix with a 4-by-1 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix hcat(Matrix<Element, 4, 3> const & lhs, Matrix<Element, 4, 1> const & rhs) {
+    return Matrix(
+      lhs.at(0, 0), lhs.at(0, 1), lhs.at(0, 2), rhs.at(0, 0)
+      , lhs.at(1, 0), lhs.at(1, 1), lhs.at(1, 2), rhs.at(1, 0)
+      , lhs.at(2, 0), lhs.at(2, 1), lhs.at(2, 2), rhs.at(2, 0)
+      , lhs.at(3, 0), lhs.at(3, 1), lhs.at(3, 2), rhs.at(3, 0));
+  }
+
+  /// Forms a 4-by-4 matrix by vertically concatenating a 1-by-4 matrix with a 3-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 1, 4> const & upper, Matrix<Element, 3, 4> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3)
+      , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2), lower.at(1, 3)
+      , lower.at(2, 0), lower.at(2, 1), lower.at(2, 2), lower.at(2, 3));
+  }
+
+  /// Forms a 4-by-4 matrix by vertically concatenating a 2-by-4 matrix with a 2-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 2, 4> const & upper, Matrix<Element, 2, 4> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3)
+      , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2), upper.at(1, 3)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3)
+      , lower.at(1, 0), lower.at(1, 1), lower.at(1, 2), lower.at(1, 3));
+  }
+
+  /// Forms a 4-by-4 matrix by vertically concatenating a 3-by-4 matrix with a 1-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  static Matrix vcat(Matrix<Element, 3, 4> const & upper, Matrix<Element, 1, 4> const & lower) {
+    return Matrix(
+      upper.at(0, 0), upper.at(0, 1), upper.at(0, 2), upper.at(0, 3)
+      , upper.at(1, 0), upper.at(1, 1), upper.at(1, 2), upper.at(1, 3)
+      , upper.at(2, 0), upper.at(2, 1), upper.at(2, 2), upper.at(2, 3)
+      , lower.at(0, 0), lower.at(0, 1), lower.at(0, 2), lower.at(0, 3));
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Element                         A, Matrix<Element, 1, 3> const & B,
+    Matrix<Element, 3, 1> const & C, Matrix<Element, 3, 3> const & D) {
+    return Matrix(
+      A, B.at(0, 0), B.at(0, 1), B.at(0, 2)
+      , C.at(0, 0), D.at(0, 0), D.at(0, 1), D.at(0, 2)
+      , C.at(1, 0), D.at(1, 0), D.at(1, 1), D.at(1, 2)
+      , C.at(2, 0), D.at(2, 0), D.at(2, 1), D.at(2, 2)
+    );
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 2> const & A, Matrix<Element, 1, 2> const & B,
+    Matrix<Element, 3, 2> const & C, Matrix<Element, 3, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1)
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1)
+      , C.at(1, 0), C.at(1, 1), D.at(1, 0), D.at(1, 1)
+      , C.at(2, 0), C.at(2, 1), D.at(2, 0), D.at(2, 1)
+    );
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 1, 3> const & A, Element                         B,
+    Matrix<Element, 3, 3> const & C, Matrix<Element, 3, 1> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), A.at(0, 2), B
+      , C.at(0, 0), C.at(0, 1), C.at(0, 2), D.at(0, 0)
+      , C.at(1, 0), C.at(1, 1), C.at(1, 2), D.at(1, 0)
+      , C.at(2, 0), C.at(2, 1), C.at(2, 2), D.at(2, 0)
+    );
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 1> const & A, Matrix<Element, 2, 3> const & B,
+    Matrix<Element, 2, 1> const & C, Matrix<Element, 2, 3> const & D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0), B.at(0, 1), B.at(0, 2)
+      , A.at(1, 0), B.at(1, 0), B.at(1, 1), B.at(1, 2)
+      , C.at(0, 0), D.at(0, 0), D.at(0, 1), D.at(0, 2)
+      , C.at(1, 0), D.at(1, 0), D.at(1, 1), D.at(1, 2)
+    );
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 2> const & A, Matrix<Element, 2, 2> const & B,
+    Matrix<Element, 2, 2> const & C, Matrix<Element, 2, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1)
+      , A.at(1, 0), A.at(1, 1), B.at(1, 0), B.at(1, 1)
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1)
+      , C.at(1, 0), C.at(1, 1), D.at(1, 0), D.at(1, 1)
+    );
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 2, 3> const & A, Matrix<Element, 2, 1> const & B,
+    Matrix<Element, 2, 3> const & C, Matrix<Element, 2, 1> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), A.at(0, 2), B.at(0, 0)
+      , A.at(1, 0), A.at(1, 1), A.at(1, 2), B.at(1, 0)
+      , C.at(0, 0), C.at(0, 1), C.at(0, 2), D.at(0, 0)
+      , C.at(1, 0), C.at(1, 1), C.at(1, 2), D.at(1, 0)
+    );
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 3, 1> const & A, Matrix<Element, 3, 3> const & B,
+    Element                         C, Matrix<Element, 1, 3> const & D) {
+    return Matrix(
+      A.at(0, 0), B.at(0, 0), B.at(0, 1), B.at(0, 2)
+      , A.at(1, 0), B.at(1, 0), B.at(1, 1), B.at(1, 2)
+      , A.at(2, 0), B.at(2, 0), B.at(2, 1), B.at(2, 2)
+      , C, D.at(0, 0), D.at(0, 1), D.at(0, 2)
+    );
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 3, 2> const & A, Matrix<Element, 3, 2> const & B,
+    Matrix<Element, 1, 2> const & C, Matrix<Element, 1, 2> const & D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), B.at(0, 0), B.at(0, 1)
+      , A.at(1, 0), A.at(1, 1), B.at(1, 0), B.at(1, 1)
+      , A.at(2, 0), A.at(2, 1), B.at(2, 0), B.at(2, 1)
+      , C.at(0, 0), C.at(0, 1), D.at(0, 0), D.at(0, 1)
+    );
+  }
+
+  /// Forms a 4-by-4 matrix by concatenating four components
+  MUTLASS_HOST_DEVICE
+  static Matrix block(
+    Matrix<Element, 3, 3> const & A, Matrix<Element, 3, 1> const & B,
+    Matrix<Element, 1, 3> const & C, Element                         D) {
+    return Matrix(
+      A.at(0, 0), A.at(0, 1), A.at(0, 2), B.at(0, 0)
+      , A.at(1, 0), A.at(1, 1), A.at(1, 2), B.at(1, 0)
+      , A.at(2, 0), A.at(2, 1), A.at(2, 2), B.at(2, 0)
+      , C.at(0, 0), C.at(0, 1), C.at(0, 2), D
+    );
+  }
+
+  /// Elementwise add operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix add(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] + rhs.data[0];
+    result.data[1] = data[1] + rhs.data[1];
+    result.data[2] = data[2] + rhs.data[2];
+    result.data[3] = data[3] + rhs.data[3];
+
+    result.data[4] = data[4] + rhs.data[4];
+    result.data[5] = data[5] + rhs.data[5];
+    result.data[6] = data[6] + rhs.data[6];
+    result.data[7] = data[7] + rhs.data[7];
+
+    result.data[8] = data[8] + rhs.data[8];
+    result.data[9] = data[9] + rhs.data[9];
+    result.data[10] = data[10] + rhs.data[10];
+    result.data[11] = data[11] + rhs.data[11];
+
+    result.data[12] = data[12] + rhs.data[12];
+    result.data[13] = data[13] + rhs.data[13];
+    result.data[14] = data[14] + rhs.data[14];
+    result.data[15] = data[15] + rhs.data[15];
+
+    return result;
+  }
+
+  /// Elementwise add operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator +(Matrix const &rhs) const {
+    return add(rhs);
+  }
+
+  /// Elementwise add operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator +=(Matrix const &rhs) {
+
+    data[0] += rhs.data[0];
+    data[1] += rhs.data[1];
+    data[2] += rhs.data[2];
+    data[3] += rhs.data[3];
+
+    data[4] += rhs.data[4];
+    data[5] += rhs.data[5];
+    data[6] += rhs.data[6];
+    data[7] += rhs.data[7];
+
+    data[8] += rhs.data[8];
+    data[9] += rhs.data[9];
+    data[10] += rhs.data[10];
+    data[11] += rhs.data[11];
+
+    data[12] += rhs.data[12];
+    data[13] += rhs.data[13];
+    data[14] += rhs.data[14];
+    data[15] += rhs.data[15];
+
+    return *this;
+  }
+
+  /// Elementwise subtract operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix subtract(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] - rhs.data[0];
+    result.data[1] = data[1] - rhs.data[1];
+    result.data[2] = data[2] - rhs.data[2];
+    result.data[3] = data[3] - rhs.data[3];
+
+    result.data[4] = data[4] - rhs.data[4];
+    result.data[5] = data[5] - rhs.data[5];
+    result.data[6] = data[6] - rhs.data[6];
+    result.data[7] = data[7] - rhs.data[7];
+
+    result.data[8] = data[8] - rhs.data[8];
+    result.data[9] = data[9] - rhs.data[9];
+    result.data[10] = data[10] - rhs.data[10];
+    result.data[11] = data[11] - rhs.data[11];
+
+    result.data[12] = data[12] - rhs.data[12];
+    result.data[13] = data[13] - rhs.data[13];
+    result.data[14] = data[14] - rhs.data[14];
+    result.data[15] = data[15] - rhs.data[15];
+
+    return result;
+  }
+
+  /// Elementwise subtract operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator -(Matrix const &rhs) const {
+    return subtract(rhs);
+  }
+
+  /// Elementwise subtract operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator -=(Matrix const &rhs) {
+
+    data[0] -= rhs.data[0];
+    data[1] -= rhs.data[1];
+    data[2] -= rhs.data[2];
+    data[3] -= rhs.data[3];
+
+    data[4] -= rhs.data[4];
+    data[5] -= rhs.data[5];
+    data[6] -= rhs.data[6];
+    data[7] -= rhs.data[7];
+
+    data[8] -= rhs.data[8];
+    data[9] -= rhs.data[9];
+    data[10] -= rhs.data[10];
+    data[11] -= rhs.data[11];
+
+    data[12] -= rhs.data[12];
+    data[13] -= rhs.data[13];
+    data[14] -= rhs.data[14];
+    data[15] -= rhs.data[15];
+
+    return *this;
+  }
+
+  /// Elementwise multiply operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * rhs.data[0];
+    result.data[1] = data[1] * rhs.data[1];
+    result.data[2] = data[2] * rhs.data[2];
+    result.data[3] = data[3] * rhs.data[3];
+
+    result.data[4] = data[4] * rhs.data[4];
+    result.data[5] = data[5] * rhs.data[5];
+    result.data[6] = data[6] * rhs.data[6];
+    result.data[7] = data[7] * rhs.data[7];
+
+    result.data[8] = data[8] * rhs.data[8];
+    result.data[9] = data[9] * rhs.data[9];
+    result.data[10] = data[10] * rhs.data[10];
+    result.data[11] = data[11] * rhs.data[11];
+
+    result.data[12] = data[12] * rhs.data[12];
+    result.data[13] = data[13] * rhs.data[13];
+    result.data[14] = data[14] * rhs.data[14];
+    result.data[15] = data[15] * rhs.data[15];
+
+    return result;
+  }
+
+  /// Scalar multiply operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix multiply(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] * s;
+    result.data[1] = data[1] * s;
+    result.data[2] = data[2] * s;
+    result.data[3] = data[3] * s;
+
+    result.data[4] = data[4] * s;
+    result.data[5] = data[5] * s;
+    result.data[6] = data[6] * s;
+    result.data[7] = data[7] * s;
+
+    result.data[8] = data[8] * s;
+    result.data[9] = data[9] * s;
+    result.data[10] = data[10] * s;
+    result.data[11] = data[11] * s;
+
+    result.data[12] = data[12] * s;
+    result.data[13] = data[13] * s;
+    result.data[14] = data[14] * s;
+    result.data[15] = data[15] * s;
+
+    return result;
+  }
+
+  /// Scalar multiply operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator *(Element const &s) const {
+    return multiply(s);
+  }
+
+  /// Scalar multiply operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator *=(Element const &s) {
+
+    data[0] *= s;
+    data[1] *= s;
+    data[2] *= s;
+    data[3] *= s;
+
+    data[4] *= s;
+    data[5] *= s;
+    data[6] *= s;
+    data[7] *= s;
+
+    data[8] *= s;
+    data[9] *= s;
+    data[10] *= s;
+    data[11] *= s;
+
+    data[12] *= s;
+    data[13] *= s;
+    data[14] *= s;
+    data[15] *= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Matrix const &rhs) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / rhs.data[0];
+    result.data[1] = data[1] / rhs.data[1];
+    result.data[2] = data[2] / rhs.data[2];
+    result.data[3] = data[3] / rhs.data[3];
+
+    result.data[4] = data[4] / rhs.data[4];
+    result.data[5] = data[5] / rhs.data[5];
+    result.data[6] = data[6] / rhs.data[6];
+    result.data[7] = data[7] / rhs.data[7];
+
+    result.data[8] = data[8] / rhs.data[8];
+    result.data[9] = data[9] / rhs.data[9];
+    result.data[10] = data[10] / rhs.data[10];
+    result.data[11] = data[11] / rhs.data[11];
+
+    result.data[12] = data[12] / rhs.data[12];
+    result.data[13] = data[13] / rhs.data[13];
+    result.data[14] = data[14] / rhs.data[14];
+    result.data[15] = data[15] / rhs.data[15];
+
+    return result;
+  }
+
+  /// Scalar divide operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix divide(Element const &s) const {
+
+    Matrix result;
+
+    result.data[0] = data[0] / s;
+    result.data[1] = data[1] / s;
+    result.data[2] = data[2] / s;
+    result.data[3] = data[3] / s;
+
+    result.data[4] = data[4] / s;
+    result.data[5] = data[5] / s;
+    result.data[6] = data[6] / s;
+    result.data[7] = data[7] / s;
+
+    result.data[8] = data[8] / s;
+    result.data[9] = data[9] / s;
+    result.data[10] = data[10] / s;
+    result.data[11] = data[11] / s;
+
+    result.data[12] = data[12] / s;
+    result.data[13] = data[13] / s;
+    result.data[14] = data[14] / s;
+    result.data[15] = data[15] / s;
+
+    return result;
+  }
+
+  /// Scalar divide operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Element const &s) const {
+    return divide(s);
+  }
+
+  /// Scalar divide operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Element const &s) {
+
+    data[0] /= s;
+    data[1] /= s;
+    data[2] /= s;
+    data[3] /= s;
+
+    data[4] /= s;
+    data[5] /= s;
+    data[6] /= s;
+    data[7] /= s;
+
+    data[8] /= s;
+    data[9] /= s;
+    data[10] /= s;
+    data[11] /= s;
+
+    data[12] /= s;
+    data[13] /= s;
+    data[14] /= s;
+    data[15] /= s;
+
+    return *this;
+  }
+
+  /// Elementwise divide operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix operator /(Matrix const &rhs) const {
+    return divide(rhs);
+  }
+
+  /// Elementwise divide operator (4-by-4)
+  MUTLASS_HOST_DEVICE
+  Matrix & operator /=(Matrix const &rhs) {
+
+    data[0] /= rhs.data[0];
+    data[1] /= rhs.data[1];
+    data[2] /= rhs.data[2];
+    data[3] /= rhs.data[3];
+
+    data[4] /= rhs.data[4];
+    data[5] /= rhs.data[5];
+    data[6] /= rhs.data[6];
+    data[7] /= rhs.data[7];
+
+    data[8] /= rhs.data[8];
+    data[9] /= rhs.data[9];
+    data[10] /= rhs.data[10];
+    data[11] /= rhs.data[11];
+
+    data[12] /= rhs.data[12];
+    data[13] /= rhs.data[13];
+    data[14] /= rhs.data[14];
+    data[15] /= rhs.data[15];
+
+    return *this;
+  }
+
+  /// Negates each element of the matrix
+  MUTLASS_HOST_DEVICE
+  Matrix operator-() const {
+    Matrix m;
+
+    m.data[0] = -m.data[0];
+    m.data[1] = -m.data[1];
+    m.data[2] = -m.data[2];
+    m.data[3] = -m.data[3];
+    m.data[4] = -m.data[4];
+    m.data[5] = -m.data[5];
+    m.data[6] = -m.data[6];
+    m.data[7] = -m.data[7];
+    m.data[8] = -m.data[8];
+    m.data[9] = -m.data[9];
+    m.data[10] = -m.data[10];
+    m.data[11] = -m.data[11];
+    m.data[12] = -m.data[12];
+    m.data[13] = -m.data[13];
+    m.data[14] = -m.data[14];
+    m.data[15] = -m.data[15];
+
+    return m;
+  }
+
+  /// Matrix product of size 4-by-1-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> product(
+    Matrix<Element, 4, 1> const &rhs,
+    Matrix<Element, 4, 1> accum = Matrix<Element, 4, 1>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[4] * rhs.data[0];
+    accum.data[2] += data[8] * rhs.data[0];
+    accum.data[3] += data[12] * rhs.data[0];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[1];
+    accum.data[1] += data[5] * rhs.data[1];
+    accum.data[2] += data[9] * rhs.data[1];
+    accum.data[3] += data[13] * rhs.data[1];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[2];
+    accum.data[1] += data[6] * rhs.data[2];
+    accum.data[2] += data[10] * rhs.data[2];
+    accum.data[3] += data[14] * rhs.data[2];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[3];
+    accum.data[1] += data[7] * rhs.data[3];
+    accum.data[2] += data[11] * rhs.data[3];
+    accum.data[3] += data[15] * rhs.data[3];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-1-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 1> operator*(Matrix<Element, 4, 1> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-2-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> product(
+    Matrix<Element, 4, 2> const &rhs,
+    Matrix<Element, 4, 2> accum = Matrix<Element, 4, 2>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[4] * rhs.data[0];
+    accum.data[3] += data[4] * rhs.data[1];
+    accum.data[4] += data[8] * rhs.data[0];
+    accum.data[5] += data[8] * rhs.data[1];
+    accum.data[6] += data[12] * rhs.data[0];
+    accum.data[7] += data[12] * rhs.data[1];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[2];
+    accum.data[1] += data[1] * rhs.data[3];
+    accum.data[2] += data[5] * rhs.data[2];
+    accum.data[3] += data[5] * rhs.data[3];
+    accum.data[4] += data[9] * rhs.data[2];
+    accum.data[5] += data[9] * rhs.data[3];
+    accum.data[6] += data[13] * rhs.data[2];
+    accum.data[7] += data[13] * rhs.data[3];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[4];
+    accum.data[1] += data[2] * rhs.data[5];
+    accum.data[2] += data[6] * rhs.data[4];
+    accum.data[3] += data[6] * rhs.data[5];
+    accum.data[4] += data[10] * rhs.data[4];
+    accum.data[5] += data[10] * rhs.data[5];
+    accum.data[6] += data[14] * rhs.data[4];
+    accum.data[7] += data[14] * rhs.data[5];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[6];
+    accum.data[1] += data[3] * rhs.data[7];
+    accum.data[2] += data[7] * rhs.data[6];
+    accum.data[3] += data[7] * rhs.data[7];
+    accum.data[4] += data[11] * rhs.data[6];
+    accum.data[5] += data[11] * rhs.data[7];
+    accum.data[6] += data[15] * rhs.data[6];
+    accum.data[7] += data[15] * rhs.data[7];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-2-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 2> operator*(Matrix<Element, 4, 2> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-3-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> product(
+    Matrix<Element, 4, 3> const &rhs,
+    Matrix<Element, 4, 3> accum = Matrix<Element, 4, 3>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[4] * rhs.data[0];
+    accum.data[4] += data[4] * rhs.data[1];
+    accum.data[5] += data[4] * rhs.data[2];
+    accum.data[6] += data[8] * rhs.data[0];
+    accum.data[7] += data[8] * rhs.data[1];
+    accum.data[8] += data[8] * rhs.data[2];
+    accum.data[9] += data[12] * rhs.data[0];
+    accum.data[10] += data[12] * rhs.data[1];
+    accum.data[11] += data[12] * rhs.data[2];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[3];
+    accum.data[1] += data[1] * rhs.data[4];
+    accum.data[2] += data[1] * rhs.data[5];
+    accum.data[3] += data[5] * rhs.data[3];
+    accum.data[4] += data[5] * rhs.data[4];
+    accum.data[5] += data[5] * rhs.data[5];
+    accum.data[6] += data[9] * rhs.data[3];
+    accum.data[7] += data[9] * rhs.data[4];
+    accum.data[8] += data[9] * rhs.data[5];
+    accum.data[9] += data[13] * rhs.data[3];
+    accum.data[10] += data[13] * rhs.data[4];
+    accum.data[11] += data[13] * rhs.data[5];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[6];
+    accum.data[1] += data[2] * rhs.data[7];
+    accum.data[2] += data[2] * rhs.data[8];
+    accum.data[3] += data[6] * rhs.data[6];
+    accum.data[4] += data[6] * rhs.data[7];
+    accum.data[5] += data[6] * rhs.data[8];
+    accum.data[6] += data[10] * rhs.data[6];
+    accum.data[7] += data[10] * rhs.data[7];
+    accum.data[8] += data[10] * rhs.data[8];
+    accum.data[9] += data[14] * rhs.data[6];
+    accum.data[10] += data[14] * rhs.data[7];
+    accum.data[11] += data[14] * rhs.data[8];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[9];
+    accum.data[1] += data[3] * rhs.data[10];
+    accum.data[2] += data[3] * rhs.data[11];
+    accum.data[3] += data[7] * rhs.data[9];
+    accum.data[4] += data[7] * rhs.data[10];
+    accum.data[5] += data[7] * rhs.data[11];
+    accum.data[6] += data[11] * rhs.data[9];
+    accum.data[7] += data[11] * rhs.data[10];
+    accum.data[8] += data[11] * rhs.data[11];
+    accum.data[9] += data[15] * rhs.data[9];
+    accum.data[10] += data[15] * rhs.data[10];
+    accum.data[11] += data[15] * rhs.data[11];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-3-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 3> operator*(Matrix<Element, 4, 3> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> product(
+    Matrix<Element, 4, 4> const &rhs,
+    Matrix<Element, 4, 4> accum = Matrix<Element, 4, 4>()
+  ) const {
+
+    // k=0
+    accum.data[0] += data[0] * rhs.data[0];
+    accum.data[1] += data[0] * rhs.data[1];
+    accum.data[2] += data[0] * rhs.data[2];
+    accum.data[3] += data[0] * rhs.data[3];
+    accum.data[4] += data[4] * rhs.data[0];
+    accum.data[5] += data[4] * rhs.data[1];
+    accum.data[6] += data[4] * rhs.data[2];
+    accum.data[7] += data[4] * rhs.data[3];
+    accum.data[8] += data[8] * rhs.data[0];
+    accum.data[9] += data[8] * rhs.data[1];
+    accum.data[10] += data[8] * rhs.data[2];
+    accum.data[11] += data[8] * rhs.data[3];
+    accum.data[12] += data[12] * rhs.data[0];
+    accum.data[13] += data[12] * rhs.data[1];
+    accum.data[14] += data[12] * rhs.data[2];
+    accum.data[15] += data[12] * rhs.data[3];
+
+    // k=1
+    accum.data[0] += data[1] * rhs.data[4];
+    accum.data[1] += data[1] * rhs.data[5];
+    accum.data[2] += data[1] * rhs.data[6];
+    accum.data[3] += data[1] * rhs.data[7];
+    accum.data[4] += data[5] * rhs.data[4];
+    accum.data[5] += data[5] * rhs.data[5];
+    accum.data[6] += data[5] * rhs.data[6];
+    accum.data[7] += data[5] * rhs.data[7];
+    accum.data[8] += data[9] * rhs.data[4];
+    accum.data[9] += data[9] * rhs.data[5];
+    accum.data[10] += data[9] * rhs.data[6];
+    accum.data[11] += data[9] * rhs.data[7];
+    accum.data[12] += data[13] * rhs.data[4];
+    accum.data[13] += data[13] * rhs.data[5];
+    accum.data[14] += data[13] * rhs.data[6];
+    accum.data[15] += data[13] * rhs.data[7];
+
+    // k=2
+    accum.data[0] += data[2] * rhs.data[8];
+    accum.data[1] += data[2] * rhs.data[9];
+    accum.data[2] += data[2] * rhs.data[10];
+    accum.data[3] += data[2] * rhs.data[11];
+    accum.data[4] += data[6] * rhs.data[8];
+    accum.data[5] += data[6] * rhs.data[9];
+    accum.data[6] += data[6] * rhs.data[10];
+    accum.data[7] += data[6] * rhs.data[11];
+    accum.data[8] += data[10] * rhs.data[8];
+    accum.data[9] += data[10] * rhs.data[9];
+    accum.data[10] += data[10] * rhs.data[10];
+    accum.data[11] += data[10] * rhs.data[11];
+    accum.data[12] += data[14] * rhs.data[8];
+    accum.data[13] += data[14] * rhs.data[9];
+    accum.data[14] += data[14] * rhs.data[10];
+    accum.data[15] += data[14] * rhs.data[11];
+
+    // k=3
+    accum.data[0] += data[3] * rhs.data[12];
+    accum.data[1] += data[3] * rhs.data[13];
+    accum.data[2] += data[3] * rhs.data[14];
+    accum.data[3] += data[3] * rhs.data[15];
+    accum.data[4] += data[7] * rhs.data[12];
+    accum.data[5] += data[7] * rhs.data[13];
+    accum.data[6] += data[7] * rhs.data[14];
+    accum.data[7] += data[7] * rhs.data[15];
+    accum.data[8] += data[11] * rhs.data[12];
+    accum.data[9] += data[11] * rhs.data[13];
+    accum.data[10] += data[11] * rhs.data[14];
+    accum.data[11] += data[11] * rhs.data[15];
+    accum.data[12] += data[15] * rhs.data[12];
+    accum.data[13] += data[15] * rhs.data[13];
+    accum.data[14] += data[15] * rhs.data[14];
+    accum.data[15] += data[15] * rhs.data[15];
+
+    return accum;
+  }
+
+  /// Matrix product of size 4-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix<Element, 4, 4> operator*(Matrix<Element, 4, 4> const &rhs) const {
+    return product(rhs);
+  }
+
+  /// Matrix product of size 4-by-4-by-4
+  MUTLASS_HOST_DEVICE
+  Matrix & operator*=(Matrix<Element, 4, 4> const &rhs) {
+    *this = product(rhs);
+    return *this;
+  }
+
+  /// Returns the sum of elements
+  MUTLASS_HOST_DEVICE
+  Element sum(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[1];
+    accum += data[2];
+    accum += data[3];
+    accum += data[4];
+    accum += data[5];
+    accum += data[6];
+    accum += data[7];
+    accum += data[8];
+    accum += data[9];
+    accum += data[10];
+    accum += data[11];
+    accum += data[12];
+    accum += data[13];
+    accum += data[14];
+    accum += data[15];
+
+    return accum;
+  }
+
+  /// Returns the sum of squared elements
+  MUTLASS_HOST_DEVICE
+  Element norm(Element accum = Element()) const {
+
+    accum += data[0] * data[0];
+    accum += data[1] * data[1];
+    accum += data[2] * data[2];
+    accum += data[3] * data[3];
+    accum += data[4] * data[4];
+    accum += data[5] * data[5];
+    accum += data[6] * data[6];
+    accum += data[7] * data[7];
+    accum += data[8] * data[8];
+    accum += data[9] * data[9];
+    accum += data[10] * data[10];
+    accum += data[11] * data[11];
+    accum += data[12] * data[12];
+    accum += data[13] * data[13];
+    accum += data[14] * data[14];
+    accum += data[15] * data[15];
+
+    return accum;
+  }
+
+  /// Returns square root of the norm
+  MUTLASS_HOST_DEVICE
+  Element magnitude() const {
+    return fast_sqrt(norm());
+  }
+
+  /// Returns the sum of diagonal elements
+  MUTLASS_HOST_DEVICE
+  Element trace(Element accum = Element()) const {
+
+    accum += data[0];
+    accum += data[5];
+    accum += data[10];
+    accum += data[15];
+
+    return accum;
+  }
+
+  /// Returns 4-by-4 rotation matrix around the X axis
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation_X(Element theta) {
+    Matrix m = identity();
+
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    m.at(1, 1) = c;
+    m.at(1, 2) = -s;
+    m.at(2, 1) = s;
+    m.at(2, 2) = c;
+
+    return m;
+  }
+
+  /// Returns 4-by-4 rotation matrix around the Y axis
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation_Y(Element theta) {
+    Matrix m = identity();
+
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    m.at(0, 0) = c;
+    m.at(2, 0) = -s;
+    m.at(0, 2) = s;
+    m.at(2, 2) = c;
+
+    return m;
+  }
+
+  /// Returns 4-by-4 rotation matrix around the Z axis
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation_Z(Element theta) {
+    Matrix m = Matrix::identity();
+
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    m.at(0, 0) = c;
+    m.at(0, 1) = -s;
+    m.at(1, 0) = s;
+    m.at(1, 1) = c;
+
+    return m;
+  }
+
+  /// Returns a 4-by-4 rotation matrix around a unit-length axis
+  MUTLASS_HOST_DEVICE
+  static Matrix rotation(Element theta, Matrix<Element, 3, 1> const &u) {
+    Element x = u.data[0];
+    Element y = u.data[1];
+    Element z = u.data[2];
+
+    Element c = fast_cos(theta);
+    Element s = fast_sin(theta);
+
+    Element one_minus_cos = Element(1) - fast_cos(theta);
+
+    Matrix m;
+
+    m.set_slice3x3({
+      c + x * x * one_minus_cos, x * y * one_minus_cos - z * s, x * z * one_minus_cos + y * s,
+      y * x * one_minus_cos * z * s, c + y * y * one_minus_cos, y * z * one_minus_cos - x * s,
+      z * x * one_minus_cos - y * s, z * y * one_minus_cos + x * s, c + z * z * one_minus_cos
+    });
+
+    return m;
+  }
+
+  /// Returns a 4-by-4 reflection about the plane specified by the
+  /// unit-length normal vector n_unit
+  MUTLASS_HOST_DEVICE
+  static Matrix reflection(Matrix<Element, 3, 1> const &n_unit) {
+
+    Element a = n_unit.data[0];
+    Element b = n_unit.data[1];
+    Element c = n_unit.data[2];
+
+    Matrix m = Matrix::identity();
+
+    m.set_slice3x3({
+      Element(1) - Element(2) * a * a, Element(-2) * a * b, Element(-2) * a * c,
+      Element(-2) * a * b, Element(1) - Element(2) * b * b, Element(-2) * b * c,
+      Element(-2) * a * c, Element(-2) * b * c, Element(1) - Element(2) * c * c
+    });
+
+    return m;
+  }
+
+  /// Returns a perspective projection matrix typical of OpenGL applications
+  MUTLASS_HOST_DEVICE
+  static Matrix perspective(Element near_plane, Element far_plane, Element fovH, Element fovV) {
+    Element aspect = fovH / fovV;
+    Element f = Element(cos(fovV)) / Element(fovH);
+    Element Q = near_plane - far_plane;
+
+    return Matrix(
+      f / aspect, 0,                0,                           0,
+      0,          f,                0,                           0,
+      0,          0, (near_plane + far_plane) / Q, Element(2) * far_plane * near_plane / Q,
+      0,          0,                -1,                          0
+    );
+  }
+
+  MUTLASS_HOST_DEVICE
+  static Matrix translation(Matrix<Element, 3, 1> const &v) {
+    return Matrix(
+      1, 0, 0, v.data[0],
+      0, 1, 0, v.data[1],
+      0, 0, 1, v.data[2],
+      0, 0, 0, 1
+    );
+  }
+
+  /// Computes the determinant of a 4-by-4 matrix
+  MUTLASS_HOST_DEVICE
+  Element determinant(Element accum = Element()) const {
+
+    accum += at(0, 0) * Matrix<Element, 3, 3>({ at(1, 1), at(1, 2), at(1, 3), at(2, 1), at(2, 2), at(2, 3), at(3, 1), at(3, 2), at(3, 3) }).determinant();
+    accum -= at(0, 1) * Matrix<Element, 3, 3>({ at(1, 0), at(1, 2), at(1, 3), at(2, 0), at(2, 2), at(2, 3), at(3, 0), at(3, 2), at(3, 3) }).determinant();
+    accum += at(0, 2) * Matrix<Element, 3, 3>({ at(1, 0), at(1, 1), at(1, 3), at(2, 0), at(2, 1), at(2, 3), at(3, 0), at(3, 1), at(3, 3) }).determinant();
+    accum -= at(0, 3) * Matrix<Element, 3, 3>({ at(1, 0), at(1, 1), at(1, 2), at(2, 0), at(2, 1), at(2, 2), at(3, 0), at(3, 1), at(3, 2) }).determinant();
+
+    return accum;
+  }
+
+  /// Computes the inverse of a 4-by-4 matrix (ignores the optional argument)
+  MUTLASS_HOST_DEVICE
+  Matrix inverse(Element ignore = 1) const {
+    Matrix<Element, 2, 2> B = slice_2x2(0, 2);
+    Matrix<Element, 2, 2> A = slice_2x2(0, 0);
+    Matrix<Element, 2, 2> C = slice_2x2(2, 0);
+    Matrix<Element, 2, 2> D = slice_2x2(2, 2);
+
+    Matrix<Element, 2, 2> D_inv = D.inverse();
+
+    Matrix<Element, 2, 2> E = (A - B * D_inv * C).inverse();
+
+    return Matrix::block(
+      E,              -E * B * D_inv,
+      -D_inv * C * E, D_inv + D_inv * C * E * B * D_inv
+    );
+  }
+
+};
+
+/// Template alias for 4-by-4 matrix
+template <typename Element>
+using Matrix4x4 = Matrix<Element, 4, 4>;
+
+
+/// Free funciton to infer element type from template arguments
+template <typename Element>
+MUTLASS_HOST_DEVICE Matrix4x4<Element> make_Matrix4x4(
+    Element _0_0, Element _0_1, Element _0_2, Element _0_3,
+    Element _1_0, Element _1_1, Element _1_2, Element _1_3,
+    Element _2_0, Element _2_1, Element _2_2, Element _2_3,
+    Element _3_0, Element _3_1, Element _3_2, Element _3_3
+) {
+  return Matrix4x4<Element>(
+  _0_0, _0_1, _0_2, _0_3,
+  _1_0, _1_1, _1_2, _1_3,
+  _2_0, _2_1, _2_2, _2_3,
+  _3_0, _3_1, _3_2, _3_3
+  );
+}
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Elementwise scalar multiplication
+template <typename Element, int Rows, int Columns>
+MUTLASS_HOST_DEVICE
+Matrix<Element, Rows, Columns> operator*(Element s, Matrix<Element, Rows, Columns> const &rhs) {
+  return rhs.multiply(s);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/matrix_coord.h b/vllm_musa/csrc/flash_attn/include/mutlass/matrix_coord.h
new file mode 100644
index 000000000..018d06814
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/matrix_coord.h
@@ -0,0 +1,165 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a canonical coordinate for rank=2 matrices offering named indices.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// MatrixCoord wraps Coord<2, int> to provide a helper for accessing named dimensions. Classes
+/// expecting a coordinate in the rank=2 index space of a matrix should use MatrixCoord.
+struct MatrixCoord : public Coord<2, int> {
+
+public:
+
+  /// Integer-valued index
+  using Index = int;
+
+  /// Base type is a Coord of rank=2
+  using Base = Coord<2, Index>;
+
+  /// LongIndex type
+  using LongIndex = typename Base::LongIndex;
+
+private:
+
+  /// Rows dimension
+  static int const kRow = 0;
+
+  /// Columns dimension
+  static int const kColumn = 1;
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  MatrixCoord() { }
+
+  /// Constructs from Coord<2>
+  MUTLASS_HOST_DEVICE
+  MatrixCoord(Coord<2, Index> const &coord): Base(coord) { }
+
+  /// Helper to construct from a row and column
+  MUTLASS_HOST_DEVICE
+  MatrixCoord(Index row, Index column): Base(make_Coord(row, column)) { }
+
+  /// Helper to construct from a row and column, which are LongIndex based
+  MUTLASS_HOST_DEVICE
+  MatrixCoord(LongIndex row, LongIndex column): Base(make_Coord(Index(row), Index(column))) { }
+
+  /// Returns the row of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & row() const { return this->at(kRow); }
+
+  /// Returns the row of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & row() { return this->at(kRow); }
+
+  /// Returns the column of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & column() const { return this->at(kColumn); }
+
+  /// Returns the column of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & column() { return this->at(kColumn); }
+
+  //
+  // Coord operators
+  //
+
+  /// Element-wise addition
+  MUTLASS_HOST_DEVICE
+  MatrixCoord operator+(Base const& b) const {
+    return MatrixCoord(Base::operator+(b));
+  }
+
+  /// Element-wise subtraction
+  MUTLASS_HOST_DEVICE
+  MatrixCoord operator-(Base const& b) const {
+    return MatrixCoord(Base::operator-(b));
+  }
+
+  /// Element-wise multiplication
+  MUTLASS_HOST_DEVICE
+  MatrixCoord operator*(Base const& b) const {
+    return MatrixCoord(Base::operator*(b));
+  }
+
+  /// Element-wise division
+  MUTLASS_HOST_DEVICE
+  MatrixCoord operator/(Base const& b) const {
+    return MatrixCoord(Base::operator/(b));
+  }
+
+  /// In-place addition
+  MUTLASS_HOST_DEVICE
+  MatrixCoord& operator+=(Base const& b) {
+    Base::operator+=(b);
+    return *this;
+  }
+
+  /// In-place subtraction
+  MUTLASS_HOST_DEVICE
+  MatrixCoord& operator-=(Base const& b) {
+    Base::operator-=(b);
+    return *this;
+  }
+
+  /// In-place multiplication
+  MUTLASS_HOST_DEVICE
+  MatrixCoord& operator*=(Base const& b) {
+    Base::operator*=(b);
+    return *this;
+  }
+
+  /// In-place division
+  MUTLASS_HOST_DEVICE
+  MatrixCoord& operator/=(Base const& b) {
+    Base::operator/=(b);
+    return *this;
+  }
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/matrix_shape.h b/vllm_musa/csrc/flash_attn/include/mutlass/matrix_shape.h
new file mode 100644
index 000000000..b7675c2d5
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/matrix_shape.h
@@ -0,0 +1,66 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a Shape template for matrix tiles
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Describes the size of a matrix tile
+template <
+  int Row_,     ///< rows of a matrix
+  int Column_      ///< columns of a matrix
+>
+struct MatrixShape {
+  static int const kRow = Row_;           ///< rows of a matrix
+  static int const kColumn = Column_;           ///< columns of a matrix
+  static int const kCount = Row_ * Column_;  ///< total number of elements in a matrix
+
+  //
+  // Static member functions
+  //
+
+  MUTLASS_HOST_DEVICE
+  static Coord<2> toCoord() {
+    return make_Coord(kRow, kColumn);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/musa_host_adapter.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/musa_host_adapter.hpp
new file mode 100644
index 000000000..71fb533aa
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/musa_host_adapter.hpp
@@ -0,0 +1,218 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Interface betweeen a MUTLASS device-wide operator and MUSA.
+*/
+
+#pragma once
+
+#include <musa_runtime_api.h>
+#include "mutlass/mutlass.h"
+#include "mutlass/trace.h"
+
+#include "mutlass/platform/platform.h"
+#if ! defined(__MUSACC_RTC__)
+#include <cstdio>
+#endif
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// Macro-level guard for MUSA Host Adapter
+//
+#if !defined(MUTLASS_ENABLE_MUSA_HOST_ADAPTER)
+#define MUTLASS_ENABLE_MUSA_HOST_ADAPTER false
+#endif
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// This class defines an object which abstracts interactions between the MUTLASS device-wide GEMM and
+/// MUSA. The intention is to enable MUTLASS to be used with both the MUSA Runtime API and MUSA Driver API.
+struct MusaHostAdapter {
+
+  /// Limit the number of kernels
+  static constexpr int32_t kMaximumKernelCount = 4;
+
+  /// Maximum cluster size
+  static constexpr int MaxClusterSize = 32;
+
+  //
+  // Data members
+  //
+
+  /// Handles
+  void        *kernel_handles[kMaximumKernelCount];
+  int32_t      kernel_count = 0;
+
+  //
+  // Methods
+  //
+
+  /// Ctor
+  MusaHostAdapter() = default;
+
+  /// Dtor
+  virtual ~MusaHostAdapter() {}
+
+  /// Copy Ctor
+  inline MusaHostAdapter(const MusaHostAdapter & rhs):
+    kernel_count(rhs.kernel_count)
+  {
+    MUTLASS_ASSERT(rhs.kernel_count >= 0 && rhs.kernel_count < kMaximumKernelCount);
+    for (int32_t i = 0; i < rhs.kernel_count && i < kMaximumKernelCount; ++i) {
+      kernel_handles[i] = rhs.kernel_handles[i];
+    }
+  }
+
+  /// Copy Assignment
+  inline MusaHostAdapter& operator=(const MusaHostAdapter & rhs) {
+
+    MUTLASS_ASSERT(rhs.kernel_count >= 0 && rhs.kernel_count < kMaximumKernelCount);
+    for (int32_t i = 0; i < rhs.kernel_count && i < kMaximumKernelCount; ++i) {
+      kernel_handles[i] = rhs.kernel_handles[i];
+    }
+    kernel_count = rhs.kernel_count;
+    return *this;
+  }
+
+  /// Move ctor
+  inline MusaHostAdapter(MusaHostAdapter && rhs):
+    kernel_count(rhs.kernel_count)
+  {
+    MUTLASS_ASSERT(rhs.kernel_count >= 0 && rhs.kernel_count < kMaximumKernelCount);
+    for (int32_t i = 0; i < rhs.kernel_count && i < kMaximumKernelCount; ++i) {
+      kernel_handles[i] = rhs.kernel_handles[i];
+    }
+  }
+
+  /// Move assignment
+  inline MusaHostAdapter& operator=(MusaHostAdapter && rhs) {
+
+    MUTLASS_ASSERT(rhs.kernel_count >= 0 && rhs.kernel_count < kMaximumKernelCount);
+    for (int32_t i = 0; i < rhs.kernel_count && i < kMaximumKernelCount; ++i) {
+      kernel_handles[i] = rhs.kernel_handles[i];
+    }
+
+    kernel_count = rhs.kernel_count;
+
+    return *this;
+  }
+
+  /// Ctor
+  inline MusaHostAdapter(
+    void **kernel_handles_,
+    int32_t kernel_count_
+  ):
+    kernel_count(kernel_count_)
+  {
+    MUTLASS_ASSERT(kernel_count >= 0);
+    for (int32_t i = 0; i < kernel_count && i < kMaximumKernelCount; ++i) {
+      kernel_handles[i] = kernel_handles_[i];
+    }
+  }
+
+  /// Returns true if the MusaHostAdapter is empty (kernel_count == 0)
+  inline bool empty() const { return !kernel_count; }
+
+  /// Returns kernel_count
+  inline size_t size() const { return static_cast<size_t>(kernel_count); }
+
+  /// Queries the occupancy of a kernel
+  virtual Status query_occupancy(
+    int32_t *device_sms,
+    int32_t *sm_occupancy,
+    int32_t kernel_index,
+    int32_t thread_count,
+    int32_t smem_size) const = 0;
+
+  /// Launches a kernel without using Threadblock Clusters.
+  virtual Status launch(
+    dim3 const grid_dims,
+    dim3 const block_dims,
+    size_t const smem_size,
+    musaStream_t musa_stream,
+    void** kernel_params,
+    int32_t kernel_index) const = 0;
+
+  /// Launches a kernel using the MUSA Extensible Launch API and Threadblock Clusters.
+  virtual Status launch(
+    dim3 const grid_dims,
+    dim3 const cluster_dims,
+    dim3 const block_dims,
+    size_t const smem_size,
+    musaStream_t musa_stream,
+    void** kernel_params,
+    int32_t kernel_index) const = 0;
+
+protected:
+
+  /**
+   * Fills a buffer in Global Memory with a byte sequence copied from host memory.
+   * This function can be overriden to dispatch to the appropriate muMemsetD*Async API
+  */
+  virtual Status memsetDeviceImpl(
+    void* destination, ///< Device memory pointer to be filled
+    void const* fill_value, ///< Value to be filled in the buffer
+    size_t fill_size, ///< Size of the data type to be used for filling the buffer
+    size_t count, ///< Number of elements of size fill_size
+    musaStream_t stream) const = 0;
+
+public:
+
+  /// Fills a buffer in Global Memory with a byte sequence copied from host memory
+  template<class FillValueType>
+  Status memsetDevice(
+    void* destination,
+    FillValueType fill_value,
+    size_t count,
+    musaStream_t stream) const
+  {
+    return this->memsetDeviceImpl(
+      destination,
+      &fill_value,
+      sizeof(FillValueType),
+      count,
+      stream);
+  }
+
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/mutlass.h b/vllm_musa/csrc/flash_attn/include/mutlass/mutlass.h
new file mode 100644
index 000000000..812bdfe6f
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/mutlass.h
@@ -0,0 +1,165 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Basic include for MUTLASS.
+*/
+
+/*
+  Note:  MUTLASS 3x increases the host compiler requirements to C++17. However, certain
+         existing integrations of MUTLASS require C++11 host compilers.
+
+         Until this requirement can be lifted, certain headers with this annotation are required
+         to be remain consistent with C++11 syntax.
+
+         C++11 compatibility is enforced by `mutlass_test_unit_core_cpp11`.
+*/
+
+#pragma once
+
+#include "mutlass/detail/helper_macros.hpp"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+/// Status code returned by MUTLASS 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 MUTLASS occurred.
+  kErrorArchMismatch,          ///< MUTLASS runs on a device that it was not compiled for.
+  kErrorInsufficientDriver,    ///< MUTLASS runs with a driver that is too old.
+  kErrorMemoryAllocation,      ///< Kernel launch failed due to insufficient device memory.
+  kInvalid                     ///< Status is unspecified.
+};
+
+/// Convert mutlass status to status strings
+MUTLASS_HOST_DEVICE
+static char const* mutlassGetStatusString(mutlass::Status status) {
+  switch (status) {
+    case mutlass::Status::kSuccess:
+      return "Success";
+    case mutlass::Status::kErrorMisalignedOperand:
+      return "Error Misaligned Operand";
+    case mutlass::Status::kErrorInvalidDataType:
+      return "Error Invalid Data Type";
+    case mutlass::Status::kErrorInvalidLayout:
+      return "Error Invalid Layout";
+    case mutlass::Status::kErrorInvalidProblem:
+      return "Error Invalid Problem";
+    case mutlass::Status::kErrorNotSupported:
+      return "Error Not Supported";
+    case mutlass::Status::kErrorWorkspaceNull:
+      return "Error Workspace Null";
+    case mutlass::Status::kErrorInternal:
+      return "Error Internal";
+    case mutlass::Status::kErrorInsufficientDriver:
+      return "Error Insufficient Driver";
+    case mutlass::Status::kErrorArchMismatch:
+      return "Error Architecture Mismatch";
+    case mutlass::Status::kErrorMemoryAllocation:
+      return "Error Memory Allocation failed";
+    case mutlass::Status::kInvalid: break;
+  }
+
+  return "Invalid status";
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static const int NumThreadsPerWarpBeforeMP31 = 128;
+static const int NumThreadsPerWarp = 32;
+static const int NumThreadsPerWarpSquad = 128;
+static const int NumWarpsPerWarpSquad = NumThreadsPerWarpSquad / NumThreadsPerWarp;
+static const int NumThreadsPerHalfWarp = NumThreadsPerWarp / 2;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper function to return true when called by thread 0 of threadblock 0.
+MUTLASS_HOST_DEVICE bool thread0() {
+  #if defined(__MUSA_ARCH__)
+    return (!threadIdx.x && !threadIdx.y && !threadIdx.z) && (!blockIdx.x && !blockIdx.y && !blockIdx.z);
+  #else
+    return false;
+  #endif
+}
+
+/// Returns a lane index in the warp. The threads in warp may not be convergent
+MUTLASS_DEVICE
+int canonical_lane_idx() {
+  #if defined(__MUSA_ARCH__)
+  #if (__MUSA_ARCH__ >= 310)
+    return threadIdx.x % NumThreadsPerWarp;
+  #else
+    return threadIdx.x % NumThreadsPerWarpBeforeMP31;
+  #endif // #if (__MUSA_ARCH__ >= 310)
+  #else
+    return 0;
+  #endif // #if defined(__MUSA_ARCH__)
+}
+
+/// Returns a warp index in the CTA. The threads in warp may not be convergent
+/// As it doesn't sync the warp, it faster and allows forward progress
+MUTLASS_DEVICE
+int canonical_warp_idx() {
+  #if defined(__MUSA_ARCH__)
+  #if (__MUSA_ARCH__ >= 310)
+    return threadIdx.x / NumThreadsPerWarp;
+  #else
+    return threadIdx.x / NumThreadsPerWarpBeforeMP31;
+  #endif // #if (__MUSA_ARCH__ >= 310)
+  #else
+    return 0;
+  #endif // #if defined(__MUSA_ARCH__)
+}
+
+MUTLASS_DEVICE
+int canonical_warp_squad_idx() {
+  #if defined(__MUSA_ARCH__)
+    return threadIdx.x / NumThreadsPerWarpSquad;
+  #else
+    return 0;
+  #endif // #if defined(__MUSA_ARCH__)
+}
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/numeric_conversion.h b/vllm_musa/csrc/flash_attn/include/mutlass/numeric_conversion.h
new file mode 100644
index 000000000..407fa57cd
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/numeric_conversion.h
@@ -0,0 +1,1986 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Boost-like numeric conversion operator for MUTLASS numeric types
+*/
+
+#pragma once
+
+#if !defined(__MUSACC_RTC__)
+#include <cfenv>
+#endif
+
+#include "mutlass/mutlass.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/transform/thread/unary_op.h"
+
+#include "mutlass/array.h"
+#include "mutlass/half.h"
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Floating-point rounding style similare to Standard Library's formats but supporting
+/// additional rounding options.
+enum class FloatRoundStyle {
+  round_indeterminate,          ///< rounding mode unknown
+  round_toward_zero,            ///< round toward zero
+  round_to_nearest,             ///< round to nearest even
+  round_to_nearest_satfinite,   ///< round to nearest even, capping value to min and max of destination type
+  round_toward_infinity,        ///< round toward infinity
+  round_toward_neg_infinity,    ///< round toward negative infinity
+  round_half_ulp_truncate,      ///< add 0.5ulp to integer representation then round toward zero
+  round_half_ulp_trunc_dntz     ///< like round_half_ulp_truncate, except denorms are rounded *toward* zero
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename T,
+  typename S,
+  FloatRoundStyle Round = FloatRoundStyle::round_to_nearest
+>
+struct NumericConverter {
+
+  using result_type = T;
+  using source_type = S;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return static_cast<result_type>(s);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for float <=> int32_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(__MUSA_ARCH__)
+template <>
+struct NumericConverter<int32_t, float, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = int32_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return __float2int_rn(s);
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<int32_t, float, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = int32_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return __float2int_rz(s);
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<float, int32_t, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = float;
+  using source_type = int32_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return __int2float_rn(s);
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<float, int32_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = float;
+  using source_type = int32_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return __int2float_rz(s);
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+
+
+#elif !defined(__MUSACC_RTC__)
+
+template <>
+struct NumericConverter<int32_t, float, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = int32_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  static result_type convert(source_type const & s) {
+    std::fesetround(FE_TONEAREST);
+    return (result_type)std::nearbyint(s);
+  }
+
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<int32_t, float, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = int32_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  static result_type convert(source_type const & s) {
+    std::fesetround(FE_TOWARDZERO);
+    return (result_type)std::nearbyint(s);
+  }
+
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+#endif
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for float <=> int8_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(__MUSA_ARCH__)
+template <>
+struct NumericConverter<int8_t, float, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = int8_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return __float2char_rn_sat(s);
+
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<int8_t, float, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = int8_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return __float2char_rz_sat(s);
+
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<float, int8_t, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = float;
+  using source_type = int8_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    int32_t result;
+    result = static_cast<int32_t>(s);
+    return __int2float_rn(result);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<float, int8_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = float;
+  using source_type = int8_t;
+  static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    int32_t result;
+    result = static_cast<int32_t>(s);
+    return __int2float_rz(result);
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for float <=> uint8_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <>
+struct NumericConverter<uint8_t, float, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = uint8_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return __musa_f2uc_rn_sat(s);
+
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<uint8_t, float, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = uint8_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return __musa_f2uc_rz_sat(s);
+
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<float, uint8_t, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = float;
+  using source_type = uint8_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    int32_t result;
+    result = static_cast<uint32_t>(s);
+    return __uint2float_rn(result);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<float, uint8_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = float;
+  using source_type = uint8_t;
+  static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    int32_t result;
+    result = static_cast<uint32_t>(s);
+    return __uint2float_rz(result);
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+#elif !defined(__MUSACC_RTC__)
+
+template <>
+struct NumericConverter<int8_t, float, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = int8_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  static result_type convert(source_type const & s) {
+    std::fesetround(FE_TONEAREST);
+    int32_t intermediate = (int32_t)std::nearbyint(s);
+
+    // Low-end saturation
+    intermediate = std::max(intermediate, (int32_t)std::numeric_limits<int8_t>::lowest());
+
+    // High-end saturation
+    intermediate = std::min(intermediate, (int32_t)std::numeric_limits<int8_t>::max());
+
+    return static_cast<result_type>(intermediate);
+  }
+
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<int8_t, float, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = int8_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style =  FloatRoundStyle::round_toward_zero;
+
+  static result_type convert(source_type const & s) {
+    std::fesetround(FE_TOWARDZERO);
+    int32_t intermediate = (int32_t)std::nearbyint(s);
+
+    // Low-end saturation
+    intermediate = std::max(intermediate, (int32_t)std::numeric_limits<int8_t>::lowest());
+
+    // High-end saturation
+    intermediate = std::min(intermediate, (int32_t)std::numeric_limits<int8_t>::max());
+
+    return static_cast<result_type>(intermediate);
+  }
+
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+#endif
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for float <= mutlass::half_t
+template <typename T, FloatRoundStyle Round>
+struct NumericConverter<T, T, Round> {
+
+  using result_type = T;
+  using source_type = T;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return s;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for float <=> mutlass::half_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for float <= mutlass::half_t
+template <FloatRoundStyle Round>
+struct NumericConverter<float, mutlass::half_t, Round> {
+
+  using result_type = float;
+  using source_type = mutlass::half_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<float>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-to-nearest
+template <>
+struct NumericConverter<mutlass::half_t, float, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = mutlass::half_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<mutlass::half_t>(s);
+
+    return result;
+
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-toward-zero
+template <>
+struct NumericConverter<mutlass::half_t, float, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = mutlass::half_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  /// Round toward zero
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & flt) {
+    return mutlass::half_t(__float2half_rz(flt));
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for int32_t <=> mutlass::half_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+#if defined(__MUSA_ARCH__)
+/// Partial specialization for int32_t <= mutlass::half_t
+template <FloatRoundStyle Round>
+struct NumericConverter<int32_t, mutlass::half_t, Round> {
+
+  using result_type = int32_t;
+  using source_type = mutlass::half_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<int32_t>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-to-nearest
+template <>
+struct NumericConverter<int32_t, mutlass::half_t, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = int32_t;
+  using source_type = mutlass::half_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    return __half2int_rn(s.to_half());
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-toward-zero
+template <>
+struct NumericConverter<int32_t, mutlass::half_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = int32_t;
+  using source_type = mutlass::half_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  /// Round toward zero
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    return __half2int_rz(s.to_half());
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Partial specialization for int32_t => mutlass::half_t
+template <FloatRoundStyle Round>
+struct NumericConverter<mutlass::half_t, int32_t, Round> {
+
+  using result_type = mutlass::half_t;
+  using source_type = int32_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<mutlass::half_t>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-to-nearest
+template <>
+struct NumericConverter<mutlass::half_t, int32_t, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = mutlass::half_t;
+  using source_type = int32_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    return mutlass::half_t(__int2half_rn(s));
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-toward-zero
+template <>
+struct NumericConverter<mutlass::half_t, int32_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = mutlass::half_t;
+  using source_type = int32_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  /// Round toward zero
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    return mutlass::half_t(__int2half_rz(s));
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for uint32_t <=> mutlass::half_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for uint32_t <= mutlass::half_t
+template <FloatRoundStyle Round>
+struct NumericConverter<uint32_t, mutlass::half_t, Round> {
+
+  using result_type = uint32_t;
+  using source_type = mutlass::half_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<uint32_t>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-to-nearest
+template <>
+struct NumericConverter<uint32_t, mutlass::half_t, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = uint32_t;
+  using source_type = mutlass::half_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    return __half2uint_rn(s.to_half());
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-toward-zero
+template <>
+struct NumericConverter<uint32_t, mutlass::half_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = uint32_t;
+  using source_type = mutlass::half_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  /// Round toward zero
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    return __half2uint_rz(s.to_half());
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Partial specialization for uint32_t => mutlass::half_t
+template <FloatRoundStyle Round>
+struct NumericConverter<mutlass::half_t, uint32_t, Round> {
+
+  using result_type = mutlass::half_t;
+  using source_type = uint32_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<mutlass::half_t>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-to-nearest
+template <>
+struct NumericConverter<mutlass::half_t, uint32_t , FloatRoundStyle::round_to_nearest> {
+
+  using result_type = mutlass::half_t;
+  using source_type = uint32_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    return mutlass::half_t(__uint2half_rn(s));
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-toward-zero
+template <>
+struct NumericConverter<mutlass::half_t, uint32_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = mutlass::half_t;
+  using source_type = uint32_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  /// Round toward zero
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    return mutlass::half_t(__uint2half_rz(s));
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+#endif
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for float <=> mutlass::bfloat16_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 220)
+/// Partial specialization for float <= mutlass::bfloat16_t
+template <FloatRoundStyle Round>
+struct NumericConverter<float, mutlass::bfloat16_t, Round> {
+
+  using result_type = float;
+  using source_type = mutlass::bfloat16_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return static_cast<float>(s);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<mutlass::bfloat16_t, float, FloatRoundStyle::round_to_nearest> {
+  using result_type = mutlass::bfloat16_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    return static_cast<mutlass::bfloat16_t>(s);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<mutlass::bfloat16_t, float, FloatRoundStyle::round_half_ulp_truncate> {
+  using result_type = mutlass::bfloat16_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_half_ulp_truncate;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    uint32_t x32 = reinterpret_cast<uint32_t const &>(s);
+
+    #if defined(__MUSA_ARCH__)
+    if (::isfinite(s)) {
+      x32 += 0x8000;
+    }
+    #else
+    if (std::isfinite(s)) {
+      x32 += 0x8000;
+    }
+    #endif
+
+    uint16_t x16 = uint16_t((x32 >> 16) & 0xffff);
+    return mutlass::bfloat16_t::bitcast(x16);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<mutlass::bfloat16_t, float, FloatRoundStyle::round_toward_zero> {
+  using result_type = mutlass::bfloat16_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return mutlass::bfloat16_t(__float2bfloat16_rz(s));
+
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for int32_t <=> mutlass::bfloat16_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for int32_t <= mutlass::bfloat16_t
+template <FloatRoundStyle Round>
+struct NumericConverter<int32_t, mutlass::bfloat16_t, Round> {
+
+  using result_type = int32_t;
+  using source_type = mutlass::bfloat16_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<int32_t>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-to-nearest
+template <>
+struct NumericConverter<int32_t, mutlass::bfloat16_t, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = int32_t;
+  using source_type = mutlass::bfloat16_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    return __bfloat162int_rn(__mt_bfloat16(s));
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-toward-zero
+template <>
+struct NumericConverter<int32_t, mutlass::bfloat16_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = int32_t;
+  using source_type = mutlass::bfloat16_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  /// Round toward zero
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    return __bfloat162int_rz(__mt_bfloat16(s));
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for uint32_t <=> mutlass::bfloat16_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for uint32_t <= mutlass::bfloat16_t
+template <FloatRoundStyle Round>
+struct NumericConverter<uint32_t, mutlass::bfloat16_t, Round> {
+
+  using result_type = uint32_t;
+  using source_type = mutlass::bfloat16_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<uint32_t>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-to-nearest
+template <>
+struct NumericConverter<uint32_t, mutlass::bfloat16_t, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = uint32_t;
+  using source_type = mutlass::bfloat16_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const & s) {
+    return __bfloat162uint_rn(__mt_bfloat16(s));
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-toward-zero
+template <>
+struct NumericConverter<uint32_t, mutlass::bfloat16_t, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = uint32_t;
+  using source_type = mutlass::bfloat16_t;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  /// Round toward zero
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    return __bfloat162uint_rz(__mt_bfloat16(s));
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Partial specialization for uint32_t => mutlass::bfloat16_t
+template <FloatRoundStyle Round>
+struct NumericConverter<mutlass::bfloat16_t, uint32_t, Round> {
+
+  using result_type = mutlass::bfloat16_t;
+  using source_type = uint32_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<mutlass::bfloat16_t>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Partial specialization for bfloat16_t <=> int8_t
+template <typename S, FloatRoundStyle Round>
+struct NumericConverter<mutlass::bfloat16_t, S, Round> {
+
+  using result_type = mutlass::bfloat16_t;
+  using source_type = S;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result = static_cast<mutlass::bfloat16_t>(s);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-to-nearest
+template <typename S>
+struct NumericConverter< mutlass::bfloat16_t, S, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = mutlass::bfloat16_t;
+  using source_type = S;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    NumericConverter<float, S, FloatRoundStyle::round_to_nearest> convert_op1;
+    NumericConverter<mutlass::bfloat16_t, float, FloatRoundStyle::round_to_nearest> convert_op2;
+
+    return convert_op2(convert_op1(s));
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Specialization for round-toward-zero
+template <typename S>
+struct NumericConverter<mutlass::bfloat16_t, S, FloatRoundStyle::round_toward_zero> {
+
+  using result_type = mutlass::bfloat16_t;
+  using source_type = S;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    NumericConverter<float, S, FloatRoundStyle::round_toward_zero> convert_op1;
+    NumericConverter<mutlass::bfloat16_t, float, FloatRoundStyle::round_toward_zero> convert_op2;
+
+    return convert_op2(convert_op1(s));
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations for float <=> mutlass::tfloat32_t
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for float <= mutlass::tfloat32_t
+template <FloatRoundStyle Round>
+struct NumericConverter<float, mutlass::tfloat32_t, Round> {
+
+  using result_type = float;
+  using source_type = mutlass::tfloat32_t;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    return static_cast<float>(s);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<mutlass::tfloat32_t, float, FloatRoundStyle::round_to_nearest> {
+  using result_type = mutlass::tfloat32_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    unsigned storage = reinterpret_cast<unsigned const &>(s);
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ >= 220
+    return __musa_float_to_tf32(s);
+#else
+    if ((storage & 0x7f800000) != 0x7f800000) {
+
+      bool mantissa_bit = ((storage & (1 << 13)) != 0);
+      bool round_bit = ((storage & (1 << 12)) != 0);
+      bool sticky_bit = ((storage & ((1 << 12) - 1)) != 0);
+
+      if ((round_bit && sticky_bit) || (round_bit && mantissa_bit)) {
+        storage += uint32_t(1 << 13);
+      }
+
+      // Note, the following is intentionally commented out. TF32
+      // does not define the low order bits, so they may be left in
+      // an undefined state.
+      //
+      // By not truncating these bit explicitly, we avoid an extra logical
+      // operation.
+      //
+      // TF32 may be implicitly converted to float by performing this
+      // operation as needed.
+      //
+      // storage = (storage & ~0x1fff);
+    }
+    else if (storage & ~0xff800000) {
+      storage = 0x7fffffff;
+    }
+#endif
+
+    return mutlass::tfloat32_t::bitcast(storage);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<mutlass::tfloat32_t, float, FloatRoundStyle::round_half_ulp_truncate> {
+  using result_type = mutlass::tfloat32_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_half_ulp_truncate;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    return mutlass::tfloat32_t::round_half_ulp_truncate(s);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// This rounding operation is similar to half_ulp_truncate except it rounds denorms toward zero.
+/// It avoids predicated code, though it requires a temporary register.
+template <>
+struct NumericConverter<mutlass::tfloat32_t, float, FloatRoundStyle::round_half_ulp_trunc_dntz> {
+  using result_type = mutlass::tfloat32_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_half_ulp_trunc_dntz;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    unsigned y = reinterpret_cast<unsigned const &>(s);
+    y = y & 0xff800000;
+    float d = reinterpret_cast<float const &>(y);
+    float z = d / float(1 << 11) + s;
+
+    return reinterpret_cast<result_type const &>(z);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <>
+struct NumericConverter<mutlass::tfloat32_t, float, FloatRoundStyle::round_toward_zero> {
+  using result_type = mutlass::tfloat32_t;
+  using source_type = float;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_toward_zero;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+    uint32_t x = reinterpret_cast<uint32_t const &>(s);
+    return mutlass::tfloat32_t::bitcast(x & 0xffffe000);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Conversion operator for float to mutlass::tfloat32_t big and small values
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+template <
+  FloatRoundStyle RoundBig = FloatRoundStyle::round_toward_zero,
+  FloatRoundStyle RoundSmall = FloatRoundStyle::round_half_ulp_truncate
+>
+struct NumericConverterFastF32 {
+
+  // result_type holds big mutlass::tfloat32_t at idx(0) and small mutlass::tfloat32_t at idx(1)
+  using result_type = Array<mutlass::tfloat32_t, 2>;
+
+  // source data type
+  using source_type = float;
+
+  // rounding styles for big and small part
+  static FloatRoundStyle const kRoundBig = RoundBig;
+  static FloatRoundStyle const kRoundSmall = RoundSmall;
+
+  MUTLASS_HOST_DEVICE
+    static result_type convert(source_type const & source) {
+
+    result_type result;
+    NumericConverter<mutlass::tfloat32_t, float, kRoundBig> convert_big_;
+    NumericConverter<mutlass::tfloat32_t, float, kRoundSmall> convert_small_;
+
+    // convert and fill mutlass::tfloat32_t big at idx 0
+    result[0] = convert_big_(source);
+
+    // convert and fill mutlass::tfloat32_t small at idx 1
+    result[1] = convert_small_(source - static_cast<float>(result[0]));
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+    result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+#endif //defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 220)
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Conversion and Clamp operator for Integers
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename T,
+  typename S
+>
+struct NumericConverterClamp {
+
+  using result_type = T;
+  using source_type = S;
+
+  MUTLASS_HOST_DEVICE
+    static result_type convert(source_type const & s) {
+    NumericConverter<result_type, source_type> convert_op;
+    result_type const kClamp_max = platform::numeric_limits<result_type>::max();
+    result_type const kClamp_min = platform::numeric_limits<result_type>::lowest();
+    if (s < (source_type)kClamp_min)
+      return kClamp_min;
+    if (s > (source_type)kClamp_max)
+      return kClamp_max;
+    return convert_op(s);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+// This converter is needed to enable mutlass::half_t output types when using int32_t accumulators.
+// Since floating-point types do not require a clamp, this converter simply casts from
+// the source type to mutlass::half_t.
+template <
+  typename S
+>
+struct NumericConverterClamp<mutlass::half_t, S> {
+
+  using result_type = mutlass::half_t;
+  using source_type = S;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const &source) {
+    return static_cast<mutlass::half_t>(source);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Conversion operator for Array
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Conversion operator for Array
+template <
+  typename T,
+  typename S,
+  int N,
+  FloatRoundStyle Round = FloatRoundStyle::round_to_nearest,
+  typename Transform = mutlass::transform::thread::UnaryTransform::Identity
+>
+struct NumericArrayConverter {
+
+  using result_type = Array<T, N>;
+  using source_type = Array<S, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  static_assert(platform::is_same<Transform, mutlass::transform::thread::UnaryTransform::Identity>::value ||
+                platform::is_same<Transform, mutlass::transform::thread::UnaryTransform::Conjugate>::value,
+                  "Unary Operator not supported.");
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & s) {
+
+    result_type result;
+    NumericConverter<T, S, Round> convert_;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      if (platform::is_same<Transform, mutlass::transform::thread::UnaryTransform::Identity>::value) {
+        result[i] = convert_(s[i]);
+      } else { // conjugate
+        result[i] = conj(convert_(s[i]));
+      }
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <
+  typename T,
+  int N,
+  FloatRoundStyle Round,
+  typename Transform
+>
+struct NumericArrayConverter<T, T, N, Round, Transform> {
+
+  using result_type = Array<T, N>;
+  using source_type = Array<T, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  static_assert(platform::is_same<Transform, mutlass::transform::thread::UnaryTransform::Identity>::value ||
+                platform::is_same<Transform, mutlass::transform::thread::UnaryTransform::Conjugate>::value,
+                  "Unary Operator not supported.");
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const &source) {
+    if (platform::is_same<Transform, mutlass::transform::thread::UnaryTransform::Identity>::value) {
+      return source;
+    } else {
+      result_type result;
+      for (int i = 0; i < N; ++i) {
+          result[i] = conj(source[i]);
+      }
+      return result;
+    }
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for Array<half, 2> <= Array<float, 2>, round to nearest
+template <>
+struct NumericArrayConverter<mutlass::half_t, float, 2, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = Array<mutlass::half_t, 2>;
+  using source_type = Array<float, 2>;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+    //Array<mutlass::half_t, 2> result;
+    //reinterpret_cast<__half2 &>(result) = __float22half2_rn(reinterpret_cast<float2 const &>(source));
+    using v2f16_t = _Float16 __attribute__((vector_size(4)));
+    using v2f32_t = float __attribute__((vector_size(8)));
+
+    Array<mutlass::half_t, 2> result;
+
+    reinterpret_cast<v2f16_t&>(result) = __builtin_convertvector(*reinterpret_cast<v2f32_t const*>(source.data()), v2f16_t);
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Partial specialization for Array<float, 2> <= Array<mutlass::half_t, 2>, round to nearest
+template <FloatRoundStyle Round>
+struct NumericArrayConverter<float, mutlass::half_t, 2, Round> {
+
+  using result_type = Array<float, 2>;
+  using source_type = Array<mutlass::half_t, 2>;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+
+    float2 result2 = __half22float2(reinterpret_cast<__half2 const &>(source));
+    return {
+      float{result2.x},
+      float{result2.y}
+    };
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for Array<half> <= Array<float>
+template <
+  int N,
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<mutlass::half_t, float, N, Round> {
+
+  using result_type = Array<mutlass::half_t, N>;
+  using source_type = Array<float, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+
+    NumericArrayConverter<mutlass::half_t, float, 2, Round> convert_vector_;
+    NumericConverter<mutlass::half_t, float, Round> convert_element_;
+
+    result_type result;
+
+    Array<mutlass::half_t, 2> *result_ptr = reinterpret_cast<Array<mutlass::half_t, 2> *>(&result);
+    Array<float, 2> const *source_ptr = reinterpret_cast<Array<float, 2> const *>(&source);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i]);
+    }
+
+    if (N % 2) {
+      result[N - 1] = convert_element_(source[N - 1]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+
+/// Partial specialization for Array<half> <= Array<float>
+template <
+  int N,
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<float, mutlass::half_t, N, Round> {
+
+  using result_type = Array<float, N>;
+  using source_type = Array<mutlass::half_t, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+
+    NumericArrayConverter<float, mutlass::half_t, 2, Round> convert_vector_;
+    NumericConverter<float, mutlass::half_t, Round> convert_element_;
+
+    result_type result;
+
+    Array<float, 2> *result_ptr = reinterpret_cast<Array<float, 2> *>(&result);
+    Array<mutlass::half_t, 2> const *source_ptr = reinterpret_cast<Array<mutlass::half_t, 2> const *>(&source);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i]);
+    }
+
+    if (N % 2) {
+      result[N - 1] = convert_element_(source[N - 1]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 220)
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for Array<mutlass::bfloat16_t, 2> <= Array<float, 2>, round to nearest
+template <>
+struct NumericArrayConverter<mutlass::bfloat16_t, float, 2, FloatRoundStyle::round_to_nearest> {
+
+  using result_type = Array<mutlass::bfloat16_t, 2>;
+  using source_type = Array<float, 2>;
+  static FloatRoundStyle const round_style = FloatRoundStyle::round_to_nearest;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+
+    Array<mutlass::bfloat16_t, 2> result;
+
+    reinterpret_cast<__mt_bfloat162 &>(result) = __float22bfloat162_rn(reinterpret_cast<float2 const &>(source));
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Partial specialization for Array<mutlass::bfloat16_t> <= Array<float>
+template <
+  int N,
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<mutlass::bfloat16_t, float, N, Round> {
+
+  using result_type = Array<mutlass::bfloat16_t, N>;
+  using source_type = Array<float, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+
+    NumericArrayConverter<mutlass::bfloat16_t, float, 2, Round> convert_vector_;
+    NumericConverter<mutlass::bfloat16_t, float, Round> convert_element_;
+
+    result_type result;
+
+    Array<mutlass::bfloat16_t, 2> *result_ptr = reinterpret_cast<Array<mutlass::bfloat16_t, 2> *>(&result);
+    Array<float, 2> const *source_ptr = reinterpret_cast<Array<float, 2> const *>(&source);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 2; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i]);
+    }
+
+    if (N % 2) {
+      result[N - 1] = convert_element_(source[N - 1]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+#endif // if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 220)
+
+/// Partial specialization for Array<uint8_t, 1> <= Array<int, 1>
+template <
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<uint8_t, int, 1, Round> {
+
+  using result_type = Array<uint8_t, 1>;
+  using source_type = Array<int, 1>;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+    NumericConverter<uint8_t, int, Round> convert_element_;
+
+    result_type result;
+
+    result[0] = convert_element_(source[0]);
+
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/// Partial specialization for Array<int8_t> <= Array<float>
+/// Conversion is performed with saturation regardless of setting of
+/// the `Round` template parameter.
+template <
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<int8_t, float, 1, Round> {
+
+  using result_type = Array<int8_t, 1>;
+  using source_type = Array<float, 1>;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+    // Convert to int to int8_t
+    NumericConverter<int8_t, float, Round> destination_converter;
+    result_type result;
+    result[0] = destination_converter(source[0]);
+    return result;
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+// To convert a FP32 to Int that has less than 32 bits, we need to convert it to int32 first.
+template <
+  typename T,
+  int N,
+  FloatRoundStyle Round
+>
+struct NumericArrayFP32ToIntConverter {
+
+  using result_type = Array<T, N>;
+  using source_type = Array<float, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  static_assert(platform::numeric_limits<T>::is_integer, "the dest type has to be int.");
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+    // Convert float to int
+    Array<int32_t, N> temporary;
+
+    NumericArrayConverter<int32_t, float, N, Round> compute_converter;
+    temporary = compute_converter(source);
+
+    // Convert to int to int8_t
+    NumericArrayConverter<T, int32_t, N, Round> destination_converter;
+    return destination_converter(temporary);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+
+template <
+  int N,
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<int8_t, float, N, Round> {
+
+  using result_type = Array<int8_t, N>;
+  using source_type = Array<float, N>;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+    NumericArrayFP32ToIntConverter<int8_t, N, Round> converter;
+    return converter(source);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <
+  int N,
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<uint8_t, float, N, Round> {
+
+  using result_type = Array<uint8_t, N>;
+  using source_type = Array<float, N>;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+    NumericArrayFP32ToIntConverter<uint8_t, N, Round> converter;
+    return converter(source);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <
+  int N,
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<int4b_t, float, N, Round> {
+
+  using result_type = Array<int4b_t, N>;
+  using source_type = Array<float, N>;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+    NumericArrayFP32ToIntConverter<int4b_t, N, Round> converter;
+    return converter(source);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+template <
+  int N,
+  FloatRoundStyle Round
+>
+struct NumericArrayConverter<uint4b_t, float, N, Round> {
+
+  using result_type = Array<uint4b_t, N>;
+  using source_type = Array<float, N>;
+
+  MUTLASS_HOST_DEVICE
+  static result_type convert(source_type const & source) {
+    NumericArrayFP32ToIntConverter<uint4b_t, N, Round> converter;
+    return converter(source);
+  }
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(source_type const &s) const {
+    return convert(s);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// FastNumericArrayConverter only works when the source is within center range.
+/// Conversion operator for Array.  See the comments before
+/// FastLinearCombinationClamp.
+template <typename T, typename S, int N,
+          FloatRoundStyle Round = FloatRoundStyle::round_to_nearest,
+          typename Enable = void>
+struct FastNumericArrayConverter {
+  using result_type = Array<T, N>;
+  using source_type = Array<S, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const &s) {
+    NumericArrayConverter<T, S, N, Round> convert_;
+
+    return convert_(s);
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const { return convert(s); }
+};
+
+/// Partial specialization for Array<float> <= Array<int>
+template <int N, FloatRoundStyle Round>
+struct FastNumericArrayConverter<float, int, N, Round> {
+  using result_type = Array<float, N>;
+  using source_type = Array<int, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const &source) {
+    result_type result;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      int tmp = source[i] + 1262485504 /*0x4B400000*/;
+      result[i] = reinterpret_cast<float const &>(tmp) - 12582912.0f;
+    }
+
+    return result;
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const { return convert(s); }
+};
+
+/// Partial specialization for Array<int8_t, 4> <= Array<float, 4>
+// template <FloatRoundStyle Round>
+// struct FastNumericArrayConverter<int8_t, float, 4, Round> {
+//   using result_type = Array<int8_t, 4>;
+//   using source_type = Array<float, 4>;
+//   static FloatRoundStyle const round_style = Round;
+
+//   MUTLASS_DEVICE
+//   static result_type convert(source_type const &source) {
+//     Array<int32_t, 4> result;
+
+//     MUTLASS_PRAGMA_UNROLL
+//     for (int i = 0; i < 4; ++i) {
+//       float tmp = source[i] + 12582912.0f;
+//       result[i] = reinterpret_cast<int32_t const &>(tmp);
+//     }
+
+//     result[0] = __byte_perm(result[0], result[1], 0x40);
+//     result[2] = __byte_perm(result[2], result[3], 0x40);
+//     result[0] = __byte_perm(result[0], result[2], 0x5410);
+
+//     return reinterpret_cast<result_type const &>(result[0]);
+//   }
+
+//   MUTLASS_DEVICE
+//   result_type operator()(source_type const &s) const { return convert(s); }
+// };
+
+/// Partial specialization for Array<int8_t> <= Array<float>
+template <int N, FloatRoundStyle Round>
+struct FastNumericArrayConverter<int8_t, float, N, Round> {
+  static_assert(!(N % 4), "N must be multiple of 4.");
+
+  using result_type = Array<int8_t, N>;
+  using source_type = Array<float, N>;
+  static FloatRoundStyle const round_style = Round;
+
+  MUTLASS_DEVICE
+  static result_type convert(source_type const &source) {
+    FastNumericArrayConverter<int8_t, float, 4, Round> convert_vector_;
+
+    result_type result;
+
+    Array<int8_t, 4> *result_ptr =
+        reinterpret_cast<Array<int8_t, 4> *>(&result);
+    Array<float, 4> const *source_ptr =
+        reinterpret_cast<Array<float, 4> const *>(&source);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N / 4; ++i) {
+      result_ptr[i] = convert_vector_(source_ptr[i]);
+    }
+
+    return result;
+  }
+
+  MUTLASS_DEVICE
+  result_type operator()(source_type const &s) const { return convert(s); }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines preferred rounding mode for a pair of types
+template <typename T, typename S>
+struct PreferredRoundingMode {
+  static FloatRoundStyle const kRound = FloatRoundStyle::round_to_nearest;
+};
+
+#if defined(__MUSA_ARCH__)
+/// Defines preferred rounding mode for a pair of types
+template <>
+struct PreferredRoundingMode<mutlass::tfloat32_t, float> {
+  static FloatRoundStyle const kRound = FloatRoundStyle::round_half_ulp_truncate;
+};
+#endif // defined(__MUSA_ARCH__)
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Packs predicates into an array.
+template <int N>
+struct PackPredicates {
+  using result_type = Array<uint1b_t, N>;
+
+  static_assert(!(N % 4), "Must pack predicates in a count that is a multiple of 4");
+
+  MUTLASS_HOST_DEVICE
+  result_type operator()(bool const predicates[]) {
+
+    result_type packed;
+    packed.clear();
+
+    int const kWordSize = 8;
+    uint8_t *bytes = reinterpret_cast<uint8_t *>(packed.data());
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      int word_idx = (i / kWordSize);
+      int bit_idx = (i % kWordSize);
+
+      uint8_t mask = static_cast<uint8_t>((predicates[i] ? 1u : 0u) << bit_idx);
+      bytes[word_idx] = (bytes[word_idx] | mask);
+    }
+    return packed;
+  }
+};
+
+/// Packs predicates into an array
+template <int N>
+struct UnpackPredicates {
+  using result_type = Array<uint1b_t, N>;
+
+  static_assert(!(N % 4), "Must unpack predicates in a count that is a multiple of 4");
+
+  MUTLASS_HOST_DEVICE
+  void operator()(bool predicates[], result_type const &packed) {
+
+    int const kWordSize = 8;
+    uint8_t const *bytes = reinterpret_cast<uint8_t const *>(packed.data());
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < N; ++i) {
+      int word_idx = (i / kWordSize);
+      int bit_idx = (i % kWordSize);
+
+      predicates[i] = bool((bytes[word_idx] >> bit_idx) & 0x1);
+    }
+
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/numeric_size.h b/vllm_musa/csrc/flash_attn/include/mutlass/numeric_size.h
new file mode 100644
index 000000000..e737ed0e8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/numeric_size.h
@@ -0,0 +1,92 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Top-level include for all MUTLASS numeric types.
+*/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines the size of an element in bits
+template <typename T>
+struct sizeof_bits {
+  static constexpr int value = int(sizeof(T) * 8);
+};
+
+template <typename T>
+struct sizeof_bits<T const>: sizeof_bits<T> {};
+
+template <>
+struct sizeof_bits<void> {
+  static constexpr int value = 0;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns the number of bytes required to hold a specified number of bits
+MUTLASS_HOST_DEVICE
+MUTLASS_CONSTEXPR_IF_CXX17
+int
+bits_to_bytes(int bits) {
+  return (bits + 7) / 8;
+}
+
+/// Returns the number of bits required to hold a specified number of bytes
+template <class R = int, class T>
+MUTLASS_HOST_DEVICE
+constexpr
+R
+bytes_to_bits(T bytes) {
+  return R(bytes) * R(8);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+struct is_subbyte {
+  static constexpr bool value = sizeof_bits<T>::value < 8;
+};
+
+template <class T>
+struct is_subbyte<T const> : is_subbyte<T> {};
+
+}  // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/numeric_types.h b/vllm_musa/csrc/flash_attn/include/mutlass/numeric_types.h
new file mode 100644
index 000000000..ec35eccfa
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/numeric_types.h
@@ -0,0 +1,86 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Top-level include for all MUTLASS numeric types.
+*/
+#pragma once
+
+#include "mutlass/numeric_size.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <size_t... Seq>
+struct index_sequence;
+
+template <size_t N, size_t... Next>
+struct index_sequence_helper : index_sequence_helper<N - 1, N - 1, Next...> {};
+
+template <size_t... Next>
+struct index_sequence_helper<0, 0, Next...> {
+  using type = index_sequence<0, Next...>;
+};
+
+template <size_t N>
+using make_index_sequence = typename index_sequence_helper<N>::type;
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Get the register type used in kernel
+//
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template<typename T>
+struct get_unpacked_element_type {
+  using type = T;
+};
+
+} // namespace detail
+
+}  // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "mutlass/integer_subbyte.h"
+#include "mutlass/half.h"
+#include "mutlass/bfloat16.h"
+#include "mutlass/tfloat32.h"
+#include "mutlass/float8.h"
+#include "mutlass/uint128.h"
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/pipeline/mp31_pipeline.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/pipeline/mp31_pipeline.hpp
new file mode 100644
index 000000000..8c8fd51a5
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/pipeline/mp31_pipeline.hpp
@@ -0,0 +1,439 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+ #pragma once
+
+ #include "mutlass/mutlass.h"
+ #include "mutlass/detail/dependent_false.hpp"
+ #include "mutlass/arch/barrier.hpp"
+
+ ////////////////////////////////////////////////////////////////////////////////////////////////////
+
+ namespace mutlass {
+
+ ////////////////////////////////////////////////////////////////////////////////////////////////////
+
+ // Circular Buffer Index + Associated Phase
+ // Assumes only one operation possible - i.e., ++
+ template <uint32_t Stages_>
+ struct PipelineState {
+
+   static constexpr uint32_t Stages = Stages_;
+
+   int index_ = 0;
+   uint32_t phase_ = 0;
+   uint32_t count_ = 0;
+
+   MUTLASS_DEVICE
+   PipelineState(): index_{}, phase_{}, count_{} {}
+
+   MUTLASS_DEVICE
+   PipelineState(int index, uint32_t phase, uint32_t count)
+     : index_(index)
+     , phase_(phase)
+     , count_(count) {}
+
+   MUTLASS_DEVICE
+   int index() const {
+     return index_;
+   }
+
+   MUTLASS_DEVICE
+   uint32_t phase() const {
+     return phase_;
+   }
+
+   MUTLASS_DEVICE
+   uint32_t count() const {
+     return count_;
+   }
+
+   MUTLASS_DEVICE
+   void operator++() {
+     if constexpr (Stages > 0) {
+       ++index_;
+       ++count_;
+       if (index_ == Stages) {
+         index_ = 0;
+         phase_ ^= 1;
+       }
+     }
+   }
+
+   MUTLASS_DEVICE
+   PipelineState& operator+=(uint32_t num_iterations) {
+     return advance(num_iterations);
+   }
+
+   MUTLASS_DEVICE
+   PipelineState& operator=(PipelineState const& other) {
+     index_ = other.index();
+     phase_ = other.phase();
+     count_ = other.count();
+     return *this;
+   }
+
+   MUTLASS_DEVICE
+   PipelineState& advance(uint32_t num_iterations) {
+     if constexpr (Stages > 0) {
+       // Number of iterations cross over the stage boundary => flipped phase
+       if ((num_iterations < Stages) && (index_ + num_iterations) >= Stages ) {
+         phase_ ^= 1;
+       }
+       // How many times number of iterations cross over the stage boundary and
+       // end up on a odd number => flipped phase
+       if ((num_iterations >= Stages) && (((index_ + num_iterations) / Stages) % 2) == 1) {
+         phase_ ^= 1;
+       }
+       index_ = (index_ + num_iterations) % Stages;
+       count_ += num_iterations;
+     }
+     return *this;
+   }
+
+   MUTLASS_DEVICE
+   static PipelineState make_pipeline_state(PipelineState start_state, uint32_t num_iterations) {
+     return start_state.advance(num_iterations);
+   }
+ };
+
+ template <class Pipeline>
+ MUTLASS_DEVICE
+ PipelineState<Pipeline::Stages> make_producer_start_state() {
+   // Producer starts with an opposite phase as the buffers are initially empty
+   constexpr int InitialProducerStage = 0;
+   constexpr uint32_t InitialProducerPhase = 1;
+   constexpr uint32_t InitialProducerCount = 0;
+   return {InitialProducerStage, InitialProducerPhase, InitialProducerCount};
+ }
+
+ ///////////////////////////////////////////////////////////////////////////////////////////////////
+ //
+ // Mp31 TME load (producer) Async Pipeline class
+ //
+ ///////////////////////////////////////////////////////////////////////////////////////////////////
+ template <int Stages_>
+ class Mp31PipelineTmeAsync {
+ public:
+   using FullBarrier = mutlass::arch::AsyncTransactionBarrier;
+   using EmptyBarrier = mutlass::arch::AsyncBarrier;
+   static constexpr uint32_t Stages = Stages_;
+   using PipelineState = mutlass::PipelineState<Stages>;
+   static_assert(FullBarrier::ReservedAsyncBarrierCount == EmptyBarrier::ReservedAsyncBarrierCount &&
+                 FullBarrier::ReservedAsyncBarrierCount == 1);
+
+   static constexpr uint32_t NumBarriers = 2 * Stages;
+
+   struct Params {
+     uint32_t transaction_bytes = 0;
+     uint32_t num_consumers = 0;
+     uint32_t num_producers = 1;
+   };
+
+   // Constructor
+   MUTLASS_DEVICE
+   Mp31PipelineTmeAsync(Params params, uint32_t barrier_base = 0)
+     : params_(params)
+     , barrier_base_(barrier_base + FullBarrier::ReservedAsyncBarrierCount)
+   {
+     int warp_idx = canonical_warp_idx();
+
+     if (warp_idx == 0) {
+       // Init full barriers
+       MUTLASS_PRAGMA_UNROLL
+       for (int i = 0; i < Stages; ++i) {
+         FullBarrier::init(barrier_base_ + i, params_.num_producers, 0);
+       }
+
+       // Init empty barriers
+       MUTLASS_PRAGMA_UNROLL
+       for (int i = 0; i < Stages; ++i) {
+         EmptyBarrier::init(barrier_base_ + i + Stages, params_.num_consumers, 0);
+       }
+     }
+   }
+
+   ////////////////////
+   // Producer APIs
+   ////////////////////
+   MUTLASS_DEVICE
+   void producer_acquire(PipelineState state) {
+     producer_acquire(state.index(), state.phase());
+   }
+
+   MUTLASS_DEVICE
+   uint32_t producer_get_barrier_id(PipelineState state) {
+     return producer_get_barrier_id(state.index());
+   }
+
+   ////////////////////
+   // Consumers APIs
+   ////////////////////
+   MUTLASS_DEVICE
+   void consumer_wait(PipelineState state) {
+     consumer_wait(state.index(), state.phase());
+   }
+
+   MUTLASS_DEVICE
+   void consumer_release(PipelineState state) {
+     consumer_release(state.index());
+   }
+
+ private:
+   Params params_;
+   uint32_t barrier_base_;
+
+   MUTLASS_DEVICE
+   void producer_acquire(uint32_t stage, uint32_t phase) {
+     uint32_t empty_barrier_id = barrier_base_ + stage + Stages;
+     EmptyBarrier::wait(empty_barrier_id, phase);
+
+     uint32_t full_barrier_id = barrier_base_ + stage;
+     FullBarrier::arrive_and_expect_tx(full_barrier_id, params_.transaction_bytes);
+   }
+
+   MUTLASS_DEVICE
+   uint32_t producer_get_barrier_id(uint32_t stage) {
+     return barrier_base_ + stage;
+   }
+
+   MUTLASS_DEVICE
+   void consumer_wait(uint32_t stage, uint32_t phase) {
+     uint32_t full_barrier_id = barrier_base_ + stage;
+     FullBarrier::wait(full_barrier_id, phase);
+   }
+
+   MUTLASS_DEVICE
+   void consumer_release(uint32_t stage) {
+     uint32_t empty_barrier_id = barrier_base_ + stage + Stages;
+     EmptyBarrier::arrive(empty_barrier_id);
+   }
+ };
+
+ ///////////////////////////////////////////////////////////////////////////////////////////////////
+ //
+ // Simple producer-consumer async Pipeline class
+ //
+ ///////////////////////////////////////////////////////////////////////////////////////////////////
+
+ template <int Stages_>
+ class Mp31PipelineAsync {
+ public:
+   using FullBarrier = mutlass::arch::AsyncBarrier;
+   using EmptyBarrier = mutlass::arch::AsyncBarrier;
+   static constexpr uint32_t Stages = Stages_;
+   using PipelineState = mutlass::PipelineState<Stages>;
+   static_assert(FullBarrier::ReservedAsyncBarrierCount == EmptyBarrier::ReservedAsyncBarrierCount &&
+                 FullBarrier::ReservedAsyncBarrierCount == 1);
+
+   static constexpr uint32_t NumBarriers = 2 * Stages;
+
+   struct Params {
+     uint32_t producer_arv_count = 1;
+     uint32_t consumer_arv_count = 1;
+   };
+
+   // Constructor
+   MUTLASS_DEVICE
+   Mp31PipelineAsync(Params params, uint32_t barrier_base = 0)
+     : params_(params)
+     , barrier_base_(barrier_base + FullBarrier::ReservedAsyncBarrierCount)
+   {
+     int warp_idx = canonical_warp_idx();
+
+     if (warp_idx == 0) {
+       // Init full barriers
+       MUTLASS_PRAGMA_UNROLL
+       for (int i = 0; i < Stages; ++i) {
+         FullBarrier::init(barrier_base_ + i, params_.producer_arv_count, 0);
+       }
+
+       // Init empty barriers
+       MUTLASS_PRAGMA_UNROLL
+       for (int i = 0; i < Stages; ++i) {
+         EmptyBarrier::init(barrier_base_ + i + Stages, params_.consumer_arv_count, 0);
+       }
+     }
+   }
+
+   ////////////////////
+   // Producer APIs
+   ////////////////////
+   MUTLASS_DEVICE
+   void producer_acquire(PipelineState state) {
+     producer_acquire(state.index(), state.phase());
+   }
+
+   MUTLASS_DEVICE
+   void producer_commit(PipelineState state) {
+     producer_commit(state.index());
+   }
+
+   MUTLASS_DEVICE
+   uint32_t producer_get_barrier_id(PipelineState state) {
+     return producer_get_barrier_id(state.index());
+   }
+
+   ////////////////////
+   // Consumers APIs
+   ////////////////////
+   MUTLASS_DEVICE
+   void consumer_wait(PipelineState state) {
+     consumer_wait(state.index(), state.phase());
+   }
+
+   MUTLASS_DEVICE
+   void consumer_release(PipelineState state) {
+     consumer_release(state.index());
+   }
+
+ private:
+   Params params_;
+   uint32_t barrier_base_;
+
+   MUTLASS_DEVICE
+   void producer_acquire(uint32_t stage, uint32_t phase) {
+     uint32_t empty_barrier_id = barrier_base_ + stage + Stages;
+     EmptyBarrier::wait(empty_barrier_id, phase);
+   }
+
+   MUTLASS_DEVICE
+   void producer_commit(uint32_t stage) {
+     uint32_t full_barrier_id = barrier_base_ + stage;
+     FullBarrier::arrive(full_barrier_id);
+   }
+
+   MUTLASS_DEVICE
+   uint32_t producer_get_barrier_id(uint32_t stage) {
+     return barrier_base_ + stage;
+   }
+
+   MUTLASS_DEVICE
+   void consumer_wait(uint32_t stage, uint32_t phase) {
+     uint32_t full_barrier_id = barrier_base_ + stage;
+     FullBarrier::wait(full_barrier_id, phase);
+   }
+
+   MUTLASS_DEVICE
+   void consumer_release(uint32_t stage) {
+     uint32_t empty_barrier_id = barrier_base_ + stage + Stages;
+     EmptyBarrier::arrive(empty_barrier_id);
+   }
+ };
+
+ ///////////////////////////////////////////////////////////////////////////////////////////////////
+ //
+ // Barrier to ensure an Ordered Sequence between
+ // SequenceLength number of groups (each with group_size participants) executing SequenceDepth Stages
+ // i.e., for all i < j - only after id "i" arrives at a particular stage "m"
+ // will the wait() for id "j" succeed for the same stage
+ //
+ ///////////////////////////////////////////////////////////////////////////////////////////////////
+
+ template <int SequenceDepth, int SequenceLength>
+ class OrderedSequenceBarrier {
+ public :
+   using Barrier = mutlass::arch::AsyncTransactionBarrier;
+   static constexpr int NumBarriers = SequenceDepth * SequenceLength;
+
+   struct Params {
+     uint32_t group_id;
+     uint32_t group_size;
+   };
+
+ private:
+   Params params_;
+   PipelineState<SequenceDepth> stage_;
+   uint32_t barrier_base_;
+
+   static constexpr int Depth = SequenceDepth;
+   static constexpr int Length = SequenceLength;
+
+ public:
+   OrderedSequenceBarrier() = delete;
+   OrderedSequenceBarrier(const OrderedSequenceBarrier&) = delete;
+   OrderedSequenceBarrier(OrderedSequenceBarrier&&) = delete;
+   OrderedSequenceBarrier& operator=(const OrderedSequenceBarrier&) = delete;
+   OrderedSequenceBarrier& operator=(OrderedSequenceBarrier&&) = delete;
+   ~OrderedSequenceBarrier() = default;
+
+   MUTLASS_DEVICE
+   OrderedSequenceBarrier(Params const& params, uint32_t barrier_base = 0)
+       : params_(params)
+       , stage_({0, params.group_id == 0, 0})
+       , barrier_base_(barrier_base + Barrier::ReservedAsyncBarrierCount)
+   {
+     int warp_idx = canonical_warp_idx();
+     if (warp_idx == 0) {
+       MUTLASS_PRAGMA_UNROLL
+       for (int d = 0; d < Depth; ++d) {
+         MUTLASS_PRAGMA_UNROLL
+         for (int l = 0; l < Length; ++l) {
+           uint32_t barrier_id = d * Length + l + barrier_base_;
+           Barrier::init(barrier_id, params.group_size, 0);
+         }
+       }
+     }
+   }
+
+   // Wait on a stage to be unlocked
+   MUTLASS_DEVICE
+   void wait() {
+     uint32_t barrier_id = get_barrier_for_current_stage(params_.group_id);
+     Barrier::wait(barrier_id, stage_.phase());
+   }
+
+   // Signal completion of Stage and move to the next stage
+   // (group_id) signals to (group_id+1)
+   MUTLASS_DEVICE
+   void arrive() {
+     int signalling_id = (params_.group_id + 1) % Length;
+     uint32_t barrier_id = get_barrier_for_current_stage(signalling_id);
+     Barrier::arrive</* return_phase = */ false>(barrier_id);
+     ++stage_;
+   }
+
+   MUTLASS_DEVICE
+   void advance() {
+     ++stage_;
+   }
+
+ private:
+   MUTLASS_DEVICE
+   uint32_t get_barrier_for_current_stage(int group_id) {
+     return stage_.index() * Length + group_id + barrier_base_;
+   }
+ };
+
+ ////////////////////////////////////////////////////////////////////////////////////////////////////
+
+ } // end namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/pipeline/pipeline.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/pipeline/pipeline.hpp
new file mode 100644
index 000000000..650898ae7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/pipeline/pipeline.hpp
@@ -0,0 +1,37 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "mutlass/pipeline/mp31_pipeline.hpp"
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/pitch_linear_coord.h b/vllm_musa/csrc/flash_attn/include/mutlass/pitch_linear_coord.h
new file mode 100644
index 000000000..89c543a0f
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/pitch_linear_coord.h
@@ -0,0 +1,181 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 layout functions used by TensorRef and derived classes for pitch-linear memory.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template defining a shape used by pitch-linear operators
+template <
+  int Contiguous,
+  int Strided
+>
+struct PitchLinearShape {
+  static int const kContiguous = Contiguous;
+  static int const kStrided = Strided;
+  static int const kCount = Contiguous * Strided;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Coordinate in pitch-linear space
+struct PitchLinearCoord : public Coord<2, int> {
+public:
+
+  /// Integer-valued index
+  using Index = int;
+
+  /// Base type is a Coord of rank=2
+  using Base = Coord<2, Index>;
+
+  /// Long integer type
+  using LongIndex = typename Base::LongIndex;
+
+private:
+
+  /// Rows dimension
+  static int const kContiguous = 0;
+
+  /// Columns dimension
+  static int const kStrided = 1;
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord() { }
+
+  /// Constructs from Coord<2>
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord(Coord<2, Index> const &coord): Base(coord) { }
+
+  /// Helper to construct from a row and column
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord(Index contiguous_, Index strided_): Base(make_Coord(contiguous_, strided_)) { }
+
+  /// Helper to construct from a row and column based on LongIndex
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord(LongIndex contiguous_, LongIndex strided_)
+    : Base(make_Coord(Index(contiguous_), Index(strided_))) { }
+
+  /// Returns the contiguous dimension
+  MUTLASS_HOST_DEVICE
+  Index const & contiguous() const { return this->at(kContiguous); }
+
+  /// Returns the contiguous dimension
+  MUTLASS_HOST_DEVICE
+  Index & contiguous() { return this->at(kContiguous); }
+
+  /// Returns the column of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & strided() const { return this->at(kStrided); }
+
+  /// Returns the column of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & strided() { return this->at(kStrided); }
+
+  //
+  // Coord operators
+  //
+
+  /// Element-wise addition
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord operator+(Base const& b) const {
+    return PitchLinearCoord(Base::operator+(b));
+  }
+
+  /// Element-wise subtraction
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord operator-(Base const& b) const {
+    return PitchLinearCoord(Base::operator-(b));
+  }
+
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord operator-() const {
+    return PitchLinearCoord(-at(0), -at(1));
+  }
+
+  /// Element-wise multiplication
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord operator*(Base const& b) const {
+    return PitchLinearCoord(Base::operator*(b));
+  }
+
+  /// Element-wise division
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord operator/(Base const& b) const {
+    return PitchLinearCoord(Base::operator/(b));
+  }
+
+  /// In-place addition
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord& operator+=(Base const& b) {
+    Base::operator+=(b);
+    return *this;
+  }
+
+  /// In-place subtraction
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord& operator-=(Base const& b) {
+    Base::operator-=(b);
+    return *this;
+  }
+
+  /// In-place multiplication
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord& operator*=(Base const& b) {
+    Base::operator*=(b);
+    return *this;
+  }
+
+  /// In-place division
+  MUTLASS_HOST_DEVICE
+  PitchLinearCoord& operator/=(Base const& b) {
+    Base::operator/=(b);
+    return *this;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/platform/platform.h b/vllm_musa/csrc/flash_attn/include/mutlass/platform/platform.h
new file mode 100644
index 000000000..5af22c38f
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/platform/platform.h
@@ -0,0 +1,958 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief C++ features that may be otherwise unimplemented for MUSA device functions.
+ *
+ * This file has three components:
+ *
+ *   (1) Macros:
+ *       - Empty macro defines for C++ keywords not supported by the current
+ *         version of C++. These simply allow compilation to proceed (but do
+ *         not provide the added semantics).
+ *           - \p noexcept
+ *           - \p constexpr
+ *           - \p nullptr
+ *           - \p static_assert
+ *
+ *       - Macro functions that we need in constant expressions because the
+ *         C++ equivalents require constexpr compiler support.  These are
+ *         prefixed with \p __NV_STD_*
+ *           - \p __NV_STD_MAX
+ *           - \p __NV_STD_MIN
+ *
+ *   (2) Re-implementations of STL functions and types:
+ *       - C++ features that need the \p __device__ annotation.  These are
+ *         placed into the \p platform namespace.
+ *           - \p abs
+ *           - \p plus
+ *           - \p less
+ *           - \p greater
+ *           - \p min
+ *           - \p max
+ *           - \p methods on std::pair (==, !=, <, <=, >, >=, and make_pair())
+ *
+ *   (3) Stop-gap implementations of unsupported STL functions and types:
+ *       - STL functions and types defined by C++ 11/14/17/etc. that are not
+ *         provided by the current version of C++. These are placed into the
+ *         \p platform namespace
+ *           - \p integral_constant
+ *           - \p nullptr_t
+ *           - \p true_type
+ *           - \p false_type
+ *           - \p bool_constant
+ *           - \p enable_if
+ *           - \p conditional
+ *           - \p is_same
+ *           - \p is_base_of
+ *           - \p remove_const
+ *           - \p remove_volatile
+ *           - \p remove_cv
+ *           - \p is_volatile
+ *           - \p is_pointer
+ *           - \p is_void
+ *           - \p is_integral
+ *           - \p is_floating_point
+ *           - \p is_arithmetic
+ *           - \p is_fundamental
+ *           - \p is_trivially_copyable
+ *           - \p alignment_of
+ *           - \p aligned_storage
+ *
+ * The idea is that, as we drop support for older compilers, we can simply #define
+ * the \p __NV_STD_XYZ macros and \p platform namespace to alias their C++
+ * counterparts (or trivially find-and-replace their occurrences in code text).
+ */
+
+//-----------------------------------------------------------------------------
+// Dependencies
+//-----------------------------------------------------------------------------
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/type_traits>
+#include <musa/std/utility>
+#include <musa/std/cstddef>
+#include <musa/std/cstdint>
+#include <musa/std/limits>
+#else
+#include <type_traits>
+#include <utility>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#endif
+
+#if !defined(__MUSACC_RTC__)
+//-----------------------------------------------------------------------------
+// Include STL files that platform provides functionality for
+//-----------------------------------------------------------------------------
+
+#include <algorithm>   // Minimum/maximum operations
+#include <cstddef>     // nullptr_t
+#include <functional>  // Arithmetic operations
+#include <utility>     // For methods on std::pair
+#include <limits>      // float_round_style, float_denorm_style
+#if (!defined(_MSC_VER) && (__cplusplus >= 201103L)) || (defined(_MSC_VER) && (_MS_VER >= 1500))
+#include <type_traits>  // For integral constants, conditional metaprogramming, and type traits
+#endif
+
+#include <vector_types.h>
+#include <mutlass/mutlass.h>
+
+#endif
+
+//-----------------------------------------------------------------------------
+// OS
+//-----------------------------------------------------------------------------
+#if defined(WIN32) || defined(_WIN32) || defined(__WIN32) && !defined(__CYGWIN__)
+#define MUTLASS_OS_WINDOWS
+#endif
+
+/******************************************************************************
+ * Macros
+ ******************************************************************************/
+/// std
+#if !defined(MUTLASS_STL_NAMESPACE)
+#if defined(__MUSACC_RTC__)
+#define MUTLASS_STL_NAMESPACE musa::std
+#else
+#define MUTLASS_STL_NAMESPACE std
+#endif
+#endif
+
+/// builtin_unreachable
+#if !defined(MUTLASS_GCC_UNREACHABLE)
+#  if defined(__GNUC__)
+#    define MUTLASS_GCC_UNREACHABLE __builtin_unreachable()
+#  else
+#    define MUTLASS_GCC_UNREACHABLE
+#  endif
+#endif
+
+//-----------------------------------------------------------------------------
+// Keywords
+//-----------------------------------------------------------------------------
+
+/// noexcept, constexpr
+#if (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1900))
+#ifndef noexcept
+#define noexcept
+#endif
+#ifndef constexpr
+#define constexpr
+#endif
+#endif
+
+/// nullptr
+#if (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1310))
+#ifndef nullptr
+#define nullptr 0
+#endif
+#endif
+
+/// static_assert
+#if (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1600))
+#ifndef static_assert
+#define __platform_cat_(a, b) a##b
+#define __platform_cat(a, b) __platform_cat_(a, b)
+#define static_assert(__e, __m) typedef int __platform_cat(AsSeRt, __LINE__)[(__e) ? 1 : -1]
+#endif
+#endif
+
+//-----------------------------------------------------------------------------
+// Functions
+//-----------------------------------------------------------------------------
+
+/// Select maximum(a, b)
+#ifndef __NV_STD_MAX
+#define __NV_STD_MAX(a, b) (((b) > (a)) ? (b) : (a))
+#endif
+
+/// Select minimum(a, b)
+#ifndef __NV_STD_MIN
+#define __NV_STD_MIN(a, b) (((b) < (a)) ? (b) : (a))
+#endif
+
+/******************************************************************************
+ * Re-implementations
+ ******************************************************************************/
+namespace mutlass {
+namespace platform {
+
+//-----------------------------------------------------------------------------
+// Abs operations <algorithm>
+//-----------------------------------------------------------------------------
+
+#if defined(__MUSACC_RTC__)
+/// std::abs
+MUTLASS_HOST_DEVICE constexpr int abs(int a) {
+    return (a < 0) ? -a : a;
+}
+MUTLASS_HOST_DEVICE constexpr long long abs(long long a) {
+    return (a < 0) ? -a : a;
+}
+#else
+using std::abs;
+#endif
+
+//-----------------------------------------------------------------------------
+// Minimum/maximum operations <algorithm>
+//-----------------------------------------------------------------------------
+
+/// std::min
+template <typename T>
+MUTLASS_HOST_DEVICE constexpr const T& min(const T& a, const T& b) {
+  return (b < a) ? b : a;
+}
+
+/// std::max
+template <typename T>
+MUTLASS_HOST_DEVICE constexpr const T& max(const T& a, const T& b) {
+  return (a < b) ? b : a;
+}
+
+#if !defined(__MUSACC_RTC__)
+//-----------------------------------------------------------------------------
+// Methods on std::pair
+//-----------------------------------------------------------------------------
+
+using std::pair;
+
+template <class T1, class T2>
+MUTLASS_HOST_DEVICE constexpr bool operator==(const pair<T1, T2>& lhs, const pair<T1, T2>& rhs) {
+  return (lhs.first == rhs.first) && (lhs.second == rhs.second);
+}
+
+template <class T1, class T2>
+MUTLASS_HOST_DEVICE constexpr bool operator!=(const pair<T1, T2>& lhs, const pair<T1, T2>& rhs) {
+  return (lhs.first != rhs.first) && (lhs.second != rhs.second);
+}
+
+template <class T1, class T2>
+MUTLASS_HOST_DEVICE constexpr bool operator<(const pair<T1, T2>& lhs, const pair<T1, T2>& rhs) {
+  return (lhs.first < rhs.first) ? true : (rhs.first < lhs.first) ? false
+                                                                  : (lhs.second < rhs.second);
+}
+
+template <class T1, class T2>
+MUTLASS_HOST_DEVICE constexpr bool operator<=(const pair<T1, T2>& lhs, const pair<T1, T2>& rhs) {
+  return !(rhs < lhs);
+}
+
+template <class T1, class T2>
+MUTLASS_HOST_DEVICE constexpr bool operator>(const pair<T1, T2>& lhs, const pair<T1, T2>& rhs) {
+  return (rhs < lhs);
+}
+
+template <class T1, class T2>
+MUTLASS_HOST_DEVICE constexpr bool operator>=(const pair<T1, T2>& lhs, const pair<T1, T2>& rhs) {
+  return !(lhs < rhs);
+}
+
+template <class T1, class T2>
+MUTLASS_HOST_DEVICE std::pair<T1, T2> make_pair(T1 t, T2 u) {
+  std::pair<T1, T2> retval;
+  retval.first = t;
+  retval.second = u;
+  return retval;
+}
+#endif
+
+}  // namespace platform
+
+/******************************************************************************
+ * Implementations of C++ 11/14/17/... STL features
+ ******************************************************************************/
+
+namespace platform {
+
+//-----------------------------------------------------------------------------
+// Integral constant helper types <type_traits>
+//-----------------------------------------------------------------------------
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500))
+
+#else
+
+using std::pair;
+
+#endif
+
+using MUTLASS_STL_NAMESPACE::integral_constant;
+using MUTLASS_STL_NAMESPACE::bool_constant;
+using MUTLASS_STL_NAMESPACE::true_type;
+using MUTLASS_STL_NAMESPACE::false_type;
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus <= 201402L)) || (defined(_MSC_VER) && (_MSC_VER < 1900))
+
+/// std::bool_constant
+template <bool V>
+struct bool_constant : platform::integral_constant<bool, V> {};
+
+#else
+
+using std::bool_constant;
+
+#endif
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1700))
+
+/// std::nullptr_t
+struct nullptr_t {};
+
+#else
+
+using std::nullptr_t;
+
+#endif
+
+//-----------------------------------------------------------------------------
+// Conditional metaprogramming <type_traits>
+//-----------------------------------------------------------------------------
+
+using MUTLASS_STL_NAMESPACE::conditional;
+using MUTLASS_STL_NAMESPACE::conditional_t;
+using MUTLASS_STL_NAMESPACE::enable_if;
+using MUTLASS_STL_NAMESPACE::enable_if_t;
+using MUTLASS_STL_NAMESPACE::void_t;
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201700L)) || (defined(_MSC_VER) && (_MSC_VER < 1600))
+
+/// std::enable_if (true specialization)
+template <bool C, typename T = void>
+struct enable_if {
+  typedef T type;
+};
+
+/// std::enable_if (false specialization)
+template <typename T>
+struct enable_if<false, T> {};
+
+/// std::conditional (true specialization)
+template <bool B, class T, class F>
+struct conditional {
+  typedef T type;
+};
+
+/// std::conditional (false specialization)
+template <class T, class F>
+struct conditional<false, T, F> {
+  typedef F type;
+};
+
+template <class...>
+using void_t = void;
+
+#else
+
+using std::enable_if;
+using std::conditional;
+using std::void_t;
+
+#endif
+
+#if (201703L <=__cplusplus)
+/// std::conditional_t
+using MUTLASS_STL_NAMESPACE::conditional_t;
+#endif
+
+//-----------------------------------------------------------------------------
+// Const/volatility specifiers <type_traits>
+//-----------------------------------------------------------------------------
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201703L)) || (defined(_MSC_VER) && (_MSC_VER < 1500))
+
+/// std::remove_const (non-const specialization)
+template <typename T>
+struct remove_const {
+  typedef T type;
+};
+
+/// std::remove_const (const specialization)
+template <typename T>
+struct remove_const<const T> {
+  typedef T type;
+};
+
+/// std::remove_volatile (non-volatile specialization)
+template <typename T>
+struct remove_volatile {
+  typedef T type;
+};
+
+/// std::remove_volatile (volatile specialization)
+template <typename T>
+struct remove_volatile<volatile T> {
+  typedef T type;
+};
+
+/// std::remove_cv
+template <typename T>
+struct remove_cv {
+  typedef typename remove_volatile<typename remove_const<T>::type>::type type;
+};
+
+#else
+
+using std::remove_const;
+using std::remove_volatile;
+using std::remove_cv;
+
+#endif
+
+#if (201703L <=__cplusplus)
+
+/// std::remove_cv_t
+using MUTLASS_STL_NAMESPACE::remove_cv_t;
+/// std::remove_reference_t
+using MUTLASS_STL_NAMESPACE::remove_reference_t;
+
+// C++20
+// using std::remove_cvref;
+template <class T>
+struct remove_cvref {
+  using type = remove_cv_t<remove_reference_t<T>>;
+};
+
+// C++20
+// using std::remove_cvref_t;
+template <class T>
+using remove_cvref_t = typename remove_cvref<T>::type;
+
+#endif
+
+
+//-----------------------------------------------------------------------------
+// Type relationships <type_traits>
+//-----------------------------------------------------------------------------
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500))
+
+/// std::is_same (false specialization)
+template <typename A, typename B>
+struct is_same : false_type {};
+
+/// std::is_same (true specialization)
+template <typename A>
+struct is_same<A, A> : true_type {};
+
+/// Helper for std::is_base_of
+template <typename BaseT, typename DerivedT>
+struct is_base_of_helper {
+  typedef char (&yes)[1];
+  typedef char (&no)[2];
+
+  template <typename B, typename D>
+  struct dummy {
+    MUTLASS_HOST_DEVICE operator B*() const;
+    MUTLASS_HOST_DEVICE operator D*();
+  };
+
+  template <typename T>
+  MUTLASS_HOST_DEVICE static yes check(DerivedT*, T);
+
+  MUTLASS_HOST_DEVICE static no check(BaseT*, int);
+
+  static const bool value = sizeof(check(dummy<BaseT, DerivedT>(), int())) == sizeof(yes);
+};
+
+/// std::is_base_of
+template <typename BaseT, typename DerivedT>
+struct is_base_of
+    : integral_constant<bool,
+                        (is_base_of_helper<typename remove_cv<BaseT>::type,
+                                           typename remove_cv<DerivedT>::type>::value) ||
+                            (is_same<typename remove_cv<BaseT>::type,
+                                     typename remove_cv<DerivedT>::type>::value)> {};
+
+#else
+
+using std::is_same;
+using std::is_base_of;
+
+#endif
+
+//-----------------------------------------------------------------------------
+// Type properties <type_traits>
+//-----------------------------------------------------------------------------
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500))
+
+/// std::is_volatile
+template <typename T>
+struct is_volatile : false_type {};
+template <typename T>
+struct is_volatile<volatile T> : true_type {};
+
+/// Helper for std::is_pointer (false specialization)
+template <typename T>
+struct is_pointer_helper : false_type {};
+
+/// Helper for std::is_pointer (true specialization)
+template <typename T>
+struct is_pointer_helper<T*> : true_type {};
+
+/// std::is_pointer
+template <typename T>
+struct is_pointer : is_pointer_helper<typename remove_cv<T>::type> {};
+
+/// std::is_void
+template <typename T>
+struct is_void : is_same<void, typename remove_cv<T>::type> {};
+
+/// std::is_integral
+template <typename T>
+struct is_integral : false_type {};
+template <>
+struct is_integral<char> : true_type {};
+template <>
+struct is_integral<signed char> : true_type {};
+template <>
+struct is_integral<unsigned char> : true_type {};
+template <>
+struct is_integral<short> : true_type {};
+template <>
+struct is_integral<unsigned short> : true_type {};
+template <>
+struct is_integral<int> : true_type {};
+template <>
+struct is_integral<unsigned int> : true_type {};
+template <>
+struct is_integral<long> : true_type {};
+template <>
+struct is_integral<unsigned long> : true_type {};
+template <>
+struct is_integral<long long> : true_type {};
+template <>
+struct is_integral<unsigned long long> : true_type {};
+template <typename T>
+struct is_integral<volatile T> : is_integral<T> {};
+template <typename T>
+struct is_integral<const T> : is_integral<T> {};
+template <typename T>
+struct is_integral<const volatile T> : is_integral<T> {};
+
+/// std::is_floating_point
+template <typename T>
+struct is_floating_point
+    : integral_constant<bool,
+                        (is_same<float, typename remove_cv<T>::type>::value ||
+                         is_same<double, typename remove_cv<T>::type>::value)> {};
+
+/// std::is_arithmetic
+template <typename T>
+struct is_arithmetic
+    : integral_constant<bool, (is_integral<T>::value || is_floating_point<T>::value)> {};
+
+/// std::is_fundamental
+template <typename T>
+struct is_fundamental
+    : integral_constant<bool,
+                        (is_arithmetic<T>::value || is_void<T>::value ||
+                         is_same<nullptr_t, typename remove_cv<T>::type>::value)> {};
+
+#else
+
+using std::is_volatile;
+using std::is_pointer;
+using std::is_void;
+using std::is_integral;
+using std::is_floating_point;
+using std::is_arithmetic;
+using std::is_fundamental;
+
+#endif
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1800)) || \
+    (defined(__GNUG__) && (__GNUC__ < 5))
+
+/**
+     * std::is_trivially_copyable
+     *
+     * This implementation only evaluates true if T is fundamental or pointer
+     *
+     * Without help from partial template specializations provided by the user for
+     * a specific class or struct, this trait will never report that the specified
+     * class or struct  is trivially-copyable ; this is always safe,
+     * if possibly sub-optimal.
+     */
+template <typename T>
+struct is_trivially_copyable
+    : integral_constant<bool, (is_fundamental<T>::value || is_pointer<T>::value)> {};
+
+#else
+
+using std::is_trivially_copyable;
+
+#endif
+
+#if (201703L <=__cplusplus)
+
+/// std::is_unsigned_v
+using MUTLASS_STL_NAMESPACE::is_integral_v;
+/// std::is_unsigned_v
+using MUTLASS_STL_NAMESPACE::is_unsigned_v;
+
+#endif
+
+//-----------------------------------------------------------------------------
+// bit_cast <bit>
+//-----------------------------------------------------------------------------
+
+template< class To, class From >
+constexpr To MUTLASS_HOST_DEVICE bit_cast(const From& from ) noexcept;
+
+template <class To, class From>
+constexpr To MUTLASS_HOST_DEVICE bit_cast(const From& src) noexcept
+{
+  static_assert(sizeof(To) == sizeof(From), "sizes must match");
+  return reinterpret_cast<To const &>(src);
+}
+
+//-----------------------------------------------------------------------------
+// Alignment and layout utilities
+//-----------------------------------------------------------------------------
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1500))
+
+/// std::alignment_of
+template <typename value_t>
+struct alignment_of {
+  struct pad {
+    value_t val;
+    char byte;
+  };
+
+  enum { value = sizeof(pad) - sizeof(value_t) };
+};
+
+#else
+
+template <typename value_t>
+struct alignment_of : std::alignment_of<value_t> {};
+
+#endif
+
+/* 16B specializations where 32-bit Win32 host compiler disagrees with device compiler */
+template <>
+struct alignment_of<int4> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<uint4> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<float4> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<long4> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<ulong4> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<longlong2> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<ulonglong2> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<double2> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<longlong4> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<ulonglong4> {
+  enum { value = 16 };
+};
+template <>
+struct alignment_of<double4> {
+  enum { value = 16 };
+};
+
+// Specializations for volatile/const qualified types
+template <typename value_t>
+struct alignment_of<volatile value_t> : alignment_of<value_t> {};
+template <typename value_t>
+struct alignment_of<const value_t> : alignment_of<value_t> {};
+template <typename value_t>
+struct alignment_of<const volatile value_t> : alignment_of<value_t> {};
+
+#if defined(__MUSACC_RTC__) || (!defined(_MSC_VER) && (__cplusplus < 201103L)) || (defined(_MSC_VER) && (_MSC_VER < 1800))
+
+template <size_t Align>
+struct aligned_chunk;
+template <>
+struct __align__(1) aligned_chunk<1> {
+  uint8_t buff;
+};
+template <>
+struct __align__(2) aligned_chunk<2> {
+  uint16_t buff;
+};
+template <>
+struct __align__(4) aligned_chunk<4> {
+  uint32_t buff;
+};
+template <>
+struct __align__(8) aligned_chunk<8> {
+  uint32_t buff[2];
+};
+template <>
+struct __align__(16) aligned_chunk<16> {
+  uint32_t buff[4];
+};
+template <>
+struct __align__(32) aligned_chunk<32> {
+  uint32_t buff[8];
+};
+template <>
+struct __align__(64) aligned_chunk<64> {
+  uint32_t buff[16];
+};
+template <>
+struct __align__(128) aligned_chunk<128> {
+  uint32_t buff[32];
+};
+template <>
+struct __align__(256) aligned_chunk<256> {
+  uint32_t buff[64];
+};
+template <>
+struct __align__(512) aligned_chunk<512> {
+  uint32_t buff[128];
+};
+template <>
+struct __align__(1024) aligned_chunk<1024> {
+  uint32_t buff[256];
+};
+template <>
+struct __align__(2048) aligned_chunk<2048> {
+  uint32_t buff[512];
+};
+template <>
+struct __align__(4096) aligned_chunk<4096> {
+  uint32_t buff[1024];
+};
+
+/// std::aligned_storage
+template <size_t Len, size_t Align>
+struct aligned_storage {
+  typedef aligned_chunk<Align> type[Len / sizeof(aligned_chunk<Align>)];
+};
+
+#else
+
+using std::aligned_storage;
+
+#endif
+
+#if !defined(__MUSACC_RTC__)
+/// Default deleter
+template <typename T>
+struct default_delete {
+  void operator()(T* ptr) const { delete ptr; }
+};
+
+/// Partial specialization for deleting array types
+template <typename T>
+struct default_delete<T[]> {
+  void operator()(T* ptr) const { delete[] ptr; }
+};
+
+/// std::unique_ptr
+template <class T, class Deleter = default_delete<T> >
+class unique_ptr {
+ public:
+  typedef T* pointer;
+  typedef T element_type;
+  typedef Deleter deleter_type;
+
+ private:
+  /// Pointer to memory
+  pointer _ptr;
+
+  /// Deleter
+  deleter_type _deleter;
+
+ public:
+  unique_ptr() : _ptr(nullptr) {}
+  unique_ptr(pointer p) : _ptr(p) {}
+
+  ~unique_ptr() {
+    if (_ptr) {
+      _deleter(_ptr);
+    }
+  }
+  /// Returns a pointer to the managed object or nullptr if no object is owned.
+  pointer get() const noexcept { return _ptr; }
+
+  /// Releases ownership of the managed object, if any
+  pointer release() noexcept {
+    pointer p(_ptr);
+    _ptr = nullptr;
+    return p;
+  }
+
+  /// Replaces the managed object, deleting the old object.
+  void reset(pointer p = pointer()) noexcept {
+    pointer old_ptr = _ptr;
+    _ptr = p;
+    if (old_ptr != nullptr) {
+      get_deleter()(old_ptr);
+    }
+  }
+
+  /// Swaps the managed objects with *this and another unique_ptr
+  void swap(unique_ptr& other) noexcept { std::swap(_ptr, other._ptr); }
+
+  /// Returns the deleter object
+  Deleter& get_deleter() noexcept { return _deleter; }
+
+  /// Returns the deleter object
+  Deleter const& get_deleter() const noexcept { return _deleter; }
+
+  /// Checks whether an object is owned
+  operator bool() const noexcept { return _ptr != nullptr; }
+
+  /// Dereferences the unique_ptr
+  T& operator*() const { return *_ptr; }
+
+  /// Returns a pointer to the managed object
+  pointer operator->() const noexcept { return _ptr; }
+
+  /// Array access to managed object
+  T& operator[](size_t i) const { return _ptr[i]; }
+};
+
+/// Specializes the swap algorithm
+template <typename T, typename Deleter>
+void swap(unique_ptr<T, Deleter>& lhs, unique_ptr<T, Deleter>& rhs) noexcept {
+  lhs.swap(rhs);
+}
+#endif
+
+/// std::numeric_limits
+template <class T>
+struct numeric_limits;
+
+template <>
+struct numeric_limits<int32_t> {
+  MUTLASS_HOST_DEVICE
+  static constexpr int32_t lowest() noexcept { return -2147483647 - 1;}
+  MUTLASS_HOST_DEVICE
+  static constexpr int32_t max() noexcept { return 2147483647;}
+  static constexpr bool is_integer = true;
+};
+
+template <>
+struct numeric_limits<int16_t> {
+  MUTLASS_HOST_DEVICE
+  static constexpr int16_t lowest() noexcept { return -32768;}
+  MUTLASS_HOST_DEVICE
+  static constexpr int16_t max() noexcept { return 32767;}
+  static constexpr bool is_integer = true;
+};
+
+template <>
+struct numeric_limits<int8_t> {
+  MUTLASS_HOST_DEVICE
+  static constexpr int8_t lowest() noexcept { return -128;}
+  MUTLASS_HOST_DEVICE
+  static constexpr int8_t max() noexcept { return 127;}
+  static constexpr bool is_integer = true;
+};
+
+
+template <>
+struct numeric_limits<uint32_t> {
+  MUTLASS_HOST_DEVICE
+  static constexpr uint32_t lowest() noexcept { return 0;}
+  MUTLASS_HOST_DEVICE
+  static constexpr uint32_t max() noexcept { return 4294967295U;}
+  static constexpr bool is_integer = true;
+};
+
+template <>
+struct numeric_limits<uint16_t> {
+  MUTLASS_HOST_DEVICE
+  static constexpr uint16_t lowest() noexcept { return 0;}
+  MUTLASS_HOST_DEVICE
+  static constexpr uint16_t max() noexcept { return 65535U;}
+  static constexpr bool is_integer = true;
+};
+
+template <>
+struct numeric_limits<uint8_t> {
+  MUTLASS_HOST_DEVICE
+  static constexpr uint8_t lowest() noexcept { return 0;}
+  MUTLASS_HOST_DEVICE
+  static constexpr uint8_t max() noexcept { return 255U;}
+  static constexpr bool is_integer = true;
+};
+
+template <>
+struct numeric_limits<float> {
+  MUTLASS_HOST_DEVICE
+  static constexpr float infinity() noexcept { return bit_cast<float, int32_t>(0x7f800000);}
+  static constexpr bool is_integer = false;
+  static constexpr bool has_infinity = true;
+};
+
+/// std::float_round_style
+using MUTLASS_STL_NAMESPACE::float_round_style;
+using MUTLASS_STL_NAMESPACE::round_indeterminate;
+using MUTLASS_STL_NAMESPACE::round_toward_zero;
+using MUTLASS_STL_NAMESPACE::round_to_nearest;
+using MUTLASS_STL_NAMESPACE::round_toward_infinity;
+using MUTLASS_STL_NAMESPACE::round_toward_neg_infinity;
+
+/// std::float_denorm_style
+using MUTLASS_STL_NAMESPACE::float_denorm_style;
+using MUTLASS_STL_NAMESPACE::denorm_indeterminate;
+using MUTLASS_STL_NAMESPACE::denorm_absent;
+using MUTLASS_STL_NAMESPACE::denorm_present;
+
+}  // namespace platform
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/predicate_vector.h b/vllm_musa/csrc/flash_attn/include/mutlass/predicate_vector.h
new file mode 100644
index 000000000..8accc26fe
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/predicate_vector.h
@@ -0,0 +1,549 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 container classes and iterators for managing a statically sized vector
+      of boolean predicates.
+*/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <assert.h>
+#include <stdint.h>
+#endif
+
+#include <musa/std/cassert>
+
+#include "mutlass/mutlass.h"
+
+#include "mutlass/platform/platform.h"
+
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/*!@defgroup predicate_vector_concept Predicate Vector Concept
+@{
+
+Implementations of \ref predicate_vector_concept contain an ordered set of boolean predicates which
+may be used as conditionals in other device-side operations. Both random access and iterators
+offering sequential access are provided.
+
+@par Predicate Vector
+   A \ref predicate_vector_concept satisfies the following expressions
+  - <b>at(int idx)</b> - returns the value of the indexed predicate
+  - <b>set(int idx, bool value)</b> - sets the value of the indexed predicate
+  - <b>begin()</b> - returns a \ref predicate_iterator_concept pointing to the first predicate
+
+@}
+*/
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/*!@defgroup predicate_iterator_concept Predicate Iterator Concept
+@{
+
+Implementations of \ref predicate_iterator_concept enables accessing and traversing elements of a
+bit vector.
+
+@par Const Predicate Iterator
+  A const \ref predicate_iterator_concept satisfies the following expressions
+ - <b>++it</b> increments the iterator to the next predicate
+ - <b>*it</b> returns the value of the currently pointed-to predicate
+
+@par Mutable Predicate Iterator
+ A \ref predicate_iterator_concept that is non-const <b>also</b> satisfies the following expressions
+ - <b>it.set(bool value)</b> sets the value of the currently pointed-to predicate
+
+@}
+*/
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/*!@defgroup predicate_tile_adapter Predicate Tile Adapter Concept
+@{
+
+Implementations of \ref predicate_tile_adapter provide a mapping between a the elements of a \ref
+tile_traits_concept and a \ref predicate_vector_concept.
+
+@par Predicate Tile Adapter
+  A \ref predicate_tile_adapter satisfies the following expressions
+ - <b>at(int d, int h, int w, int c)</b> - returns the value of a predicate corresponding to the
+   access (d, h, w, c) within the tile.
+
+@}
+*/
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Statically sized array of bits implementing @concept{predicate_vector_concept}.
+template <
+    /// Number of predicates contained in predicate vector
+    int kPredicates_,
+    /// Number of predicates contained in each byte of internal storage
+    int kPredicatesPerByte_ = 4,
+    /// Location of first predicate within byte of internal storage
+    int kPredicateStart_ = 0>
+struct PredicateVector {
+  /// Number of bits stored by the PredicateVector
+  static constexpr int kPredicates = kPredicates_;
+
+  /// Number of bits stored within each byte of the predicate bit vector
+  static constexpr int kPredicatesPerByte = kPredicatesPerByte_;
+
+  /// First bit within each byte containing predicates
+  static constexpr int kPredicateStart = kPredicateStart_;
+
+  // Make sure no one tries to put more than 8 bits in a byte :)
+  static_assert(kPredicatesPerByte <= 8, "kPredicatesPerByte must fit within an actual byte");
+  // Make sure the "offsetted" bits fit in one byte.
+  static_assert(kPredicateStart + kPredicatesPerByte <= 8,
+                "The offsetted predicates must fit within an actual byte.");
+
+  /// Storage type of individual elements
+  typedef uint32_t Storage;
+
+  /// Number of bytes needed
+  static constexpr int kBytes = (kPredicates + kPredicatesPerByte - 1) / kPredicatesPerByte;
+
+  /// Number of storage elements needed
+  static constexpr int kWordCount = (kBytes + int(sizeof(Storage)) - 1) / int(sizeof(Storage));
+
+  /// The byte mask corresponding to predicates
+  static constexpr Storage kByteMask = (((1 << kPredicatesPerByte) - 1) << kPredicateStart);
+
+ private:
+  //
+  // Data members
+  //
+
+  /// Words of bit vector
+  Storage storageData[kWordCount];
+
+  //
+  // Methods
+  //
+
+  /// Computes the word and bit corresponding to a logical predicate index
+  MUTLASS_HOST_DEVICE void computeStorageOffset(int &word, int &bit, int idx) const {
+    MUTLASS_ASSERT(idx < kPredicates);
+
+    int byte = (idx / kPredicatesPerByte);
+    int bit_offset = (idx % kPredicatesPerByte);
+
+    word = byte / sizeof(Storage);
+    int byte_offset = (byte % sizeof(Storage));
+
+    bit = byte_offset * 8 + bit_offset + kPredicateStart;
+  }
+
+  /// Returns word mask.
+  MUTLASS_HOST_DEVICE static constexpr bool computeWordMask() {
+    Storage mask(0);
+    MUTLASS_PRAGMA_UNROLL
+    for (size_t byte = 0; byte < sizeof(Storage); ++byte) {
+      mask |= (kByteMask << (byte * 8));
+    }
+    return mask;
+  }
+
+  /// Returns mask of last word.
+  MUTLASS_HOST_DEVICE static constexpr bool computeLastWordMask() {
+    Storage mask(0);
+    MUTLASS_PRAGMA_UNROLL
+    for (int byte = 0; byte < kBytes % sizeof(Storage); ++byte) {
+      mask |= (kByteMask << (byte * 8));
+    }
+    return mask;
+  }
+
+  /// Accesses a given word with optional assertions
+  MUTLASS_HOST_DEVICE Storage &storage(int word) {
+    MUTLASS_ASSERT(word < kWordCount);
+    return storageData[word];
+  }
+
+  /// Accesses a given word with optional assertions
+  MUTLASS_HOST_DEVICE Storage const &storage(int word) const {
+    MUTLASS_ASSERT(word < kWordCount);
+    return storageData[word];
+  }
+
+ public:
+  //
+  // Iterator
+  //
+
+  /**
+  * @brief An iterator implementing \ref predicate_iterator_concept enabling sequential
+  * read and write access to predicates.
+  * @concept{predicate_iterator_concept}
+  */
+  class Iterator {
+    /// Reference to PredicateVector instance
+    PredicateVector &vec_;
+
+    /// Index into PredicateVector
+    int bit_;
+
+   public:
+    /// Copy constructor
+    MUTLASS_HOST_DEVICE
+    Iterator(Iterator const &it) : vec_(it.vec_), bit_(it.bit_) {}
+
+    /// Constructs an iterator from a PredicateVector
+    MUTLASS_HOST_DEVICE
+    Iterator(PredicateVector &vec, int _start = 0) : vec_(vec), bit_(_start) {}
+
+    /// Pre-increment
+    MUTLASS_HOST_DEVICE
+    Iterator &operator++() {
+      ++bit_;
+      return *this;
+    }
+
+    /// Increment
+    MUTLASS_HOST_DEVICE
+    Iterator &operator+=(int offset) {
+      bit_ += offset;
+      return *this;
+    }
+
+    /// Pre-decrement
+    MUTLASS_HOST_DEVICE
+    Iterator &operator--() {
+      --bit_;
+      return *this;
+    }
+
+    /// Decrement
+    MUTLASS_HOST_DEVICE
+    Iterator &operator-=(int offset) {
+      bit_ -= offset;
+      return *this;
+    }
+
+    /// Post-increment
+    MUTLASS_HOST_DEVICE
+    Iterator operator++(int) {
+      Iterator ret(*this);
+      ret.bit_++;
+      return ret;
+    }
+
+    /// Post-decrement
+    MUTLASS_HOST_DEVICE
+    Iterator operator--(int) {
+      Iterator ret(*this);
+      ret.bit_--;
+      return ret;
+    }
+
+    /// Iterator advances by some amount
+    MUTLASS_HOST_DEVICE
+    Iterator operator+(int offset) {
+      Iterator ret(*this);
+      ret.bit_ += offset;
+      return ret;
+    }
+
+    /// Iterator recedes by some amount
+    MUTLASS_HOST_DEVICE
+    Iterator operator-(int offset) {
+      ConstIterator ret(*this);
+      ret.bit_ -= offset;
+      return ret;
+    }
+
+    /// Returns true if iterators point to the same bit
+    MUTLASS_HOST_DEVICE
+    bool operator==(Iterator const &it) const { return bit_ == it.bit_; }
+
+    /// Returns false if iterators point to the same bit
+    MUTLASS_HOST_DEVICE
+    bool operator!=(Iterator const &it) const { return bit_ != it.bit_; }
+
+    /// Gets the bit at the pointed to location
+    MUTLASS_HOST_DEVICE
+    bool get() { return vec_.at(bit_); }
+
+    /// Gets the bit at the pointed to location
+    MUTLASS_HOST_DEVICE
+    bool at() const { return vec_.at(bit_); }
+
+    /// Dereferences iterator
+    MUTLASS_HOST_DEVICE
+    bool operator*() const { return at(); }
+
+    /// Sets the bit at the pointed to location
+    MUTLASS_HOST_DEVICE
+    void set(bool value = true) { vec_.set(bit_, value); }
+  };
+
+  /**
+  * @brief An iterator implementing \ref predicate_iterator_concept enabling sequential
+  * read and write access to predicates.
+  * @concept{predicate_iterator_concept}
+  */
+  class ConstIterator {
+    /// Reference to PredicateVector instance
+    PredicateVector const &vec_;
+
+    /// Index into PredicateVector
+    int bit_;
+
+   public:
+    /// Copy constructor
+    MUTLASS_HOST_DEVICE
+    ConstIterator(ConstIterator const &it) : vec_(it.vec_), bit_(it.bit_) {}
+
+    /// Constructs an iterator from a PredicateVector
+    MUTLASS_HOST_DEVICE
+    ConstIterator(PredicateVector const &vec, int _start = 0) : vec_(vec), bit_(_start) {}
+
+    /// Pre-increment
+    MUTLASS_HOST_DEVICE
+    ConstIterator &operator++() {
+      ++bit_;
+      return *this;
+    }
+
+    /// Increment
+    MUTLASS_HOST_DEVICE
+    ConstIterator &operator+=(int offset) {
+      bit_ += offset;
+      return *this;
+    }
+
+    /// Pre-decrement
+    MUTLASS_HOST_DEVICE
+    ConstIterator &operator--() {
+      --bit_;
+      return *this;
+    }
+
+    /// Decrement
+    MUTLASS_HOST_DEVICE
+    ConstIterator &operator-=(int offset) {
+      bit_ -= offset;
+      return *this;
+    }
+
+    /// Post-increment
+    MUTLASS_HOST_DEVICE
+    ConstIterator operator++(int) {
+      ConstIterator ret(*this);
+      ret.bit_++;
+      return ret;
+    }
+
+    /// Post-decrement
+    MUTLASS_HOST_DEVICE
+    ConstIterator operator--(int) {
+      ConstIterator ret(*this);
+      ret.bit_--;
+      return ret;
+    }
+
+    /// Iterator advances by some amount
+    MUTLASS_HOST_DEVICE
+    ConstIterator operator+(int offset) {
+      ConstIterator ret(*this);
+      ret.bit_ += offset;
+      return ret;
+    }
+
+    /// Iterator recedes by some amount
+    MUTLASS_HOST_DEVICE
+    ConstIterator operator-(int offset) {
+      ConstIterator ret(*this);
+      ret.bit_ -= offset;
+      return ret;
+    }
+
+    /// Returns true if iterators point to the same bit
+    MUTLASS_HOST_DEVICE
+    bool operator==(ConstIterator const &it) const { return bit_ == it.bit_; }
+
+    /// Returns false if iterators point to the same bit
+    MUTLASS_HOST_DEVICE
+    bool operator!=(ConstIterator const &it) const { return bit_ != it.bit_; }
+
+    /// Gets the bit at the pointed to location
+    MUTLASS_HOST_DEVICE
+    bool get() { return vec_.at(bit_); }
+
+    /// Gets the bit at the pointed to location
+    MUTLASS_HOST_DEVICE
+    bool at() const { return vec_.at(bit_); }
+
+    /// Dereferences iterator
+    MUTLASS_HOST_DEVICE
+    bool operator*() const { return at(); }
+  };
+
+  /// Iterator that always returns true
+  struct TrivialIterator {
+    /// Constructor
+    MUTLASS_HOST_DEVICE
+    TrivialIterator() {}
+
+    /// Copy constructor
+    MUTLASS_HOST_DEVICE
+    TrivialIterator(Iterator const &it) {}
+
+    /// Constructs an iterator from a PredicateVector
+    MUTLASS_HOST_DEVICE
+    TrivialIterator(PredicateVector const &_vec) {}
+
+    /// Pre-increment
+    MUTLASS_HOST_DEVICE
+    TrivialIterator &operator++() { return *this; }
+
+    /// Post-increment
+    MUTLASS_HOST_DEVICE
+    TrivialIterator operator++(int) { return *this; }
+
+    /// Dereferences iterator
+    MUTLASS_HOST_DEVICE
+    bool operator*() const { return true; }
+  };
+
+ public:
+  //
+  // Methods
+  //
+
+  /// Initialize the predicate vector
+  MUTLASS_HOST_DEVICE PredicateVector(bool value = true) { fill(value); }
+
+  /// Fills all predicates with a given value
+  MUTLASS_HOST_DEVICE void fill(bool value = true) {
+    Storage item = (value ? ~Storage(0) : Storage(0));
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kWordCount; ++i) {
+      storage(i) = item;
+    }
+  }
+
+  /// Clears all predicates
+  MUTLASS_HOST_DEVICE void clear() {
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kWordCount; ++i) {
+      storage(i) = 0;
+    }
+  }
+
+  /// Sets all predicates to true
+  MUTLASS_HOST_DEVICE void enable() {
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kWordCount; ++i) {
+      storage(i) = ~Storage(0);
+    }
+  }
+
+  /// Accesses a bit within the predicate vector.
+  MUTLASS_HOST_DEVICE bool operator[](int idx) const { return at(idx); }
+
+  /// Accesses a bit within the predicate vector.
+  MUTLASS_HOST_DEVICE bool at(int idx) const {
+    int bit, word;
+    computeStorageOffset(word, bit, idx);
+
+    return ((storage(word) >> bit) & 1);
+  }
+
+  /// Set a bit within the predicate vector.
+  MUTLASS_HOST_DEVICE void set(int idx, bool value = true) {
+    int bit, word;
+    computeStorageOffset(word, bit, idx);
+
+    Storage disable_mask = (~(Storage(1) << bit));
+    Storage enable_mask = (Storage(value) << bit);
+
+    storage(word) = ((storage(word) & disable_mask) | enable_mask);
+  }
+
+  /// Computes the intersection of two identical predicate vectors.
+  MUTLASS_HOST_DEVICE PredicateVector &operator&=(PredicateVector const &predicates) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kWordCount; ++i) {
+      storage(i) = (storage(i) & predicates.storage(i));
+    }
+    return *this;
+  }
+
+  /// Computes the union of two identical predicate vectors.
+  MUTLASS_HOST_DEVICE PredicateVector &operator|=(PredicateVector const &predicates) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kWordCount; ++i) {
+      storage(i) = (storage(i) | predicates.storage(i));
+    }
+    return *this;
+  }
+
+  /// Returns true if entire predicate array is zero.
+  MUTLASS_HOST_DEVICE bool is_zero() const {
+   constexpr Storage mask = computeWordMask();
+    Storage result = 0;
+    MUTLASS_PRAGMA_UNROLL
+    for (int word = 0; word < kWordCount - 1; ++word) {
+      result |= (storage(word) & mask);
+    }
+    constexpr Storage last_word_mask = computeLastWordMask();
+    result |= (storage(kWordCount - 1) & last_word_mask);
+
+    return result == 0;
+  }
+
+  /// Returns an iterator to the start of the bit vector
+  MUTLASS_DEVICE
+  Iterator begin() { return Iterator(*this); }
+
+  /// Returns an iterator
+  MUTLASS_DEVICE
+  Iterator end() { return Iterator(*this, kPredicates); }
+
+  /// Returns a ConstIterator
+  MUTLASS_DEVICE
+  ConstIterator const_begin() const { return ConstIterator(*this); }
+
+  /// Returns a ConstIterator
+  MUTLASS_DEVICE
+  ConstIterator const_end() const { return ConstIterator(*this, kPredicates); }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/quaternion.h b/vllm_musa/csrc/flash_attn/include/mutlass/quaternion.h
new file mode 100644
index 000000000..4564de52b
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/quaternion.h
@@ -0,0 +1,752 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a densely packed quaternion object intended for storing data in registers and
+    executing quaternion operations within a MUSA or host thread.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/functional.h"
+#include "mutlass/array.h"
+#include "mutlass/real.h"
+#include "mutlass/coord.h"
+#include "mutlass/matrix.h"
+#include "mutlass/fast_math.h"
+#include "mutlass/layout/vector.h"
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Quaternion: xi + yj + zk + w
+template <
+  typename Element_ = float      ///< element type
+>
+class Quaternion : public Array<Element_, 4> {
+public:
+
+  /// Logical rank of tensor index space
+  static int const kRank = 1;
+
+  /// Number of elements
+  static int const kExtent = 4;
+
+  /// Base class is a four-element array
+  using Base = Array<Element_, kExtent>;
+
+  /// Element type
+  using Element = typename Base::Element;
+
+  /// Reference type to an element
+  using Reference = typename Base::reference;
+
+  /// Index type
+  using Index = int;
+
+  /// Quaternion storage - imaginary part
+  static int const kX = 0;
+
+  /// Quaternion storage - imaginary part
+  static int const kY = 1;
+
+  /// Quaternion storage - imaginary part
+  static int const kZ = 2;
+
+  /// Quaternion storage - real part
+  static int const kW = 3;
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructs a quaternion q = 0
+  MUTLASS_HOST_DEVICE
+  Quaternion() {
+    Base::at(kX) = Element();
+    Base::at(kY) = Element();
+    Base::at(kZ) = Element();
+    Base::at(kW) = Element();
+  }
+
+  /// Constructs a quaternion q = w + 0*i + 0*j + 0*k
+  MUTLASS_HOST_DEVICE
+  Quaternion(
+    Element w_
+  ) {
+    Base::at(kX) = Element();
+    Base::at(kY) = Element();
+    Base::at(kZ) = Element();
+    Base::at(kW) = w_;
+  }
+
+  /// Constructs a quaternion q = w + x*i + y*j + z*k
+  MUTLASS_HOST_DEVICE
+  Quaternion(
+    Element x_,
+    Element y_,
+    Element z_,
+    Element w_
+  ) {
+    Base::at(kX) = x_;
+    Base::at(kY) = y_;
+    Base::at(kZ) = z_;
+    Base::at(kW) = w_;
+  }
+
+  /// Constructs a quaternion from a vector representing the imaginary part and a real number
+  MUTLASS_HOST_DEVICE
+  Quaternion(
+    Matrix3x1<Element> const &imag_,
+    Element w_ = Element()
+  ) {
+    Base::at(kX) = imag_[0];
+    Base::at(kY) = imag_[1];
+    Base::at(kZ) = imag_[2];
+    Base::at(kW) = w_;
+  }
+
+  /// Returns a reference to the element at a given Coord
+  MUTLASS_HOST_DEVICE
+  Reference at(Index idx) const {
+    return Base::at(idx);
+  }
+
+  /// Returns a reference to the element at a given Coord
+  MUTLASS_HOST_DEVICE
+  Reference at(Index idx) {
+    return Base::at(idx);
+  }
+
+  /// Accesses the x element of the imaginary part of the quaternion
+  MUTLASS_HOST_DEVICE
+  Element x() const {
+    return Base::at(kX);
+  }
+
+  /// Accesses the x element of the imaginary part of the quaternion
+  MUTLASS_HOST_DEVICE
+  Reference x() {
+    return Base::at(kX);
+  }
+
+  /// Accesses the y element of the imaginary part of the quaternion
+  MUTLASS_HOST_DEVICE
+  Element y() const {
+    return Base::at(kY);
+  }
+
+  /// Accesses the y element of the imaginary part of the quaternion
+  MUTLASS_HOST_DEVICE
+  Reference y() {
+    return Base::at(kY);
+  }
+
+  /// Accesses the z element of the imaginary part of the quaternion
+  MUTLASS_HOST_DEVICE
+  Element z() const {
+    return Base::at(kZ);
+  }
+
+  /// Accesses the z element of the imaginary part of the quaternion
+  MUTLASS_HOST_DEVICE
+  Reference z() {
+    return Base::at(kZ);
+  }
+
+  /// Accesses the real part of the quaternion
+  MUTLASS_HOST_DEVICE
+  Element w() const {
+    return Base::at(kW);
+  }
+
+  /// Accesses the real part of the quaternion
+  MUTLASS_HOST_DEVICE
+  Reference w() {
+    return Base::at(kW);
+  }
+
+  /// Returns the pure imaginary part of the quaternion as a 3-vector
+  MUTLASS_HOST_DEVICE
+  Matrix3x1<Element> pure() const {
+    return Matrix3x1<Element>(x(), y(), z());
+  }
+
+  /// Returns a quaternion representation of a spatial rotation given a unit-length axis and
+  /// a rotation in radians.
+  MUTLASS_HOST_DEVICE
+  static Quaternion<Element> rotation(
+    Matrix3x1<Element> const &axis_unit,    ///< axis of rotation (assumed to be unit length)
+    Element theta) {                        ///< angular rotation in radians
+
+    Element s = fast_sin(theta / Element(2));
+
+    return Quaternion(
+      s * axis_unit[0],
+      s * axis_unit[1],
+      s * axis_unit[2],
+      fast_cos(theta / Element(2))
+    );
+  }
+
+  /// Returns a quaternion representation of a spatial rotation represented as a
+  /// unit-length rotation axis (r_x, r_y, r_z) and an angular rotation in radians
+  MUTLASS_HOST_DEVICE
+  static Quaternion<Element> rotation(
+    Element r_x,
+    Element r_y,
+    Element r_z,
+    Element theta) {                      ///< angular rotation in radians
+
+    return rotation({r_x, r_y, r_z}, theta);
+  }
+
+  /// Geometric rotation of a 3-element vector
+  MUTLASS_HOST_DEVICE
+  Matrix3x1<Element> rotate(Matrix3x1<Element> const &rhs) const {
+    return (*this * Quaternion<Element>(rhs, 0) * reciprocal(*this)).pure();
+  }
+
+  /// Inverse rotation operation
+  MUTLASS_HOST_DEVICE
+  Matrix3x1<Element> rotate_inv(Matrix3x1<Element> const &rhs) const {
+    return (reciprocal(*this) * Quaternion<Element>(rhs, 0) * *this).pure();
+  }
+
+  /// Rotates a 3-vector assuming this is a unit quaternion (a spinor)
+  MUTLASS_HOST_DEVICE
+  Matrix3x1<Element> spinor(Matrix3x1<Element> const &rhs) const {
+    return (*this * Quaternion<Element>(rhs, 0) * conj(*this)).pure();
+  }
+
+  /// Inverse rotation of 3-vector assuming this is a unit quaternion (a spinor)
+  MUTLASS_HOST_DEVICE
+  Matrix3x1<Element> spinor_inv(Matrix3x1<Element> const &rhs) const {
+    return (conj(*this) * Quaternion<Element>(rhs, 0) * *this).pure();
+  }
+
+  /// In-place addition
+  template <typename Element>
+  MUTLASS_HOST_DEVICE
+  Quaternion<Element> &operator+=(Quaternion<Element> const &rhs) {
+    *this = (*this + rhs);
+    return *this;
+  }
+
+  /// In-place subtraction
+  template <typename Element>
+  MUTLASS_HOST_DEVICE
+  Quaternion<Element> &operator-=(Quaternion<Element> const &rhs) {
+    *this = (*this - rhs);
+    return *this;
+  }
+
+  /// In-place multiplication
+  template <typename T>
+  MUTLASS_HOST_DEVICE
+  Quaternion<Element> &operator*=(Quaternion<Element> const &rhs) {
+    *this = (*this * rhs);
+    return *this;
+  }
+
+  /// Scalar multiplication
+  template <typename T>
+  MUTLASS_HOST_DEVICE
+  Quaternion<Element> &operator*=(Element s) {
+    *this = (*this * s);
+    return *this;
+  }
+
+  /// In-place Division
+  template <typename T>
+  MUTLASS_HOST_DEVICE
+  Quaternion<Element> &operator/=(Quaternion<Element> const &rhs) {
+    *this = (*this / rhs);
+    return *this;
+  }
+
+  /// In-place Division
+  template <typename T>
+  MUTLASS_HOST_DEVICE
+  Quaternion<Element> &operator/=(Element s) {
+    *this = (*this / s);
+    return *this;
+  }
+
+  /// Computes a 3x3 rotation matrix (row-major representation)
+  MUTLASS_HOST_DEVICE
+  Matrix3x3<Element> as_rotation_matrix_3x3() const {
+    Matrix3x3<Element> m(
+      w() * w() + x() * x() - y() * y() - z() * z(),
+      2 * x() * y() - 2 * w() * z(),
+      2 * x() * z() + 2 * w() * y(),
+
+      2 * x() * y() + 2 * w() * z(),
+      w() * w() - x() * x() + y() * y() - z() * z(),
+      2 * y() * z() - 2 * w() * x(),
+
+      2 * x() * z() - 2 * w() * y(),
+      2 * y() * z() + 2 * w() * x(),
+      w() * w() - x() * x() - y() * y() + z() * z()
+    );
+    return m;
+  }
+
+  /// Computes a 4x4 rotation matrix (row-major representation)
+  MUTLASS_HOST_DEVICE
+  Matrix4x4<Element> as_rotation_matrix_4x4() const {
+    Matrix4x4<Element> m = Matrix4x4<Element>::identity();
+    m.set_slice_3x3(as_rotation_matrix_3x3());
+    return m;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Constructs a quaternion that is non-zero only in its real element.
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> make_Quaternion(
+  Element w) {                                ///< real part
+
+  return Quaternion<Element>(w);
+}
+
+/// Constructs a quaternion from a vector and real
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> make_Quaternion(
+  Matrix3x1<Element> const &imag,             ///< imaginary party as a vector
+  Element w) {                                ///< real part
+
+  return Quaternion<Element>(imag, w);
+}
+
+/// Constructs a quaternion from a unit-length rotation axis and a rotation
+/// angle in radians
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> make_QuaternionRotation(
+  Matrix3x1<Element> const &axis_unit,        ///< rotation axis (unit-length)
+  Element w) {                                ///< rotation angle in radians
+
+  return Quaternion<Element>::rotation(axis_unit, w);
+}
+
+/// Constructs a quaternion q = xi + yj + zk + w
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> make_Quaternion(Element x, Element y, Element z, Element w) {
+  return Quaternion<Element>(x, y, z, w);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns the real part of the quaternion number
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Element const &real(Quaternion<Element> const &q) {
+  return q.w();
+}
+
+/// Returns the real part of the quaternion number
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Element &real(Quaternion<Element> &q) {
+  return q.w();
+}
+
+/// Returns the magnitude of the quaternion number
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Element abs(Quaternion<Element> const &q) {
+  return fast_sqrt(norm(q));
+}
+
+/// Quaternion conjugate
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> conj(Quaternion<Element> const &q) {
+  return make_Quaternion(
+    -q.x(),
+    -q.y(),
+    -q.z(),
+    q.w()
+  );
+}
+
+/// Computes the squared magnitude of the quaternion
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Element norm(Quaternion<Element> const &q) {
+  return q.x() * q.x() + q.y() * q.y() + q.z() * q.z() + q.w() * q.w();
+}
+
+/// Quaternion reciprocal
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> reciprocal(Quaternion<Element> const &q) {
+
+  Element nsq = norm(q);
+
+  return make_Quaternion(
+    -q.x() / nsq,
+    -q.y() / nsq,
+    -q.z() / nsq,
+    q.w() / nsq
+  );
+}
+
+/// Returns a unit-length quaternion
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> unit(Quaternion<Element> const &q) {
+
+  Element rcp_mag = Element(1) / abs(q);
+
+  return make_Quaternion(
+    q.x() * rcp_mag,
+    q.y() * rcp_mag,
+    q.z() * rcp_mag,
+    q.w() * rcp_mag
+  );
+}
+
+/// Quaternion exponential
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> exp(Quaternion<Element> const &q) {
+
+  Element exp_ = fast_exp(q.w());
+  Element imag_norm = fast_sqrt(q.x() * q.x() + q.y() * q.y() + q.z() * q.z());
+  Element sin_norm = fast_sin(imag_norm);
+
+  return make_Quaternion(
+    exp_ * q.x() * sin_norm / imag_norm,
+    exp_ * q.y() * sin_norm / imag_norm,
+    exp_ * q.z() * sin_norm / imag_norm,
+    exp_ * fast_cos(imag_norm)
+  );
+}
+
+/// Quaternion natural logarithm
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> log(Quaternion<Element> const &q) {
+
+  Element v = fast_sqrt(q.x() * q.x() + q.y() * q.y() + q.z() * q.z());
+  Element s = fast_acos(q.w() / abs(q)) / v;
+
+  return make_Quaternion(
+    q.x() * s,
+    q.y() * s,
+    q.z() * s,
+    fast_log(q.w())
+  );
+}
+
+/// Gets the rotation angle from a unit-length quaternion
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Element get_rotation_angle(Quaternion<Element> const &q_unit) {
+  return fast_acos(q_unit.w()) * Element(2);
+}
+
+/// Gets the rotation axis from a unit-length quaternion
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Matrix3x1<Element> get_rotation_axis(Quaternion<Element> const &q_unit) {
+  return q_unit.pure().unit();
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Equality operator
+template <typename Element>
+MUTLASS_HOST_DEVICE
+bool operator==(Quaternion<Element> const &lhs, Quaternion<Element> const &rhs) {
+  return lhs.x() == rhs.x() &&
+    lhs.y() == rhs.y() &&
+    lhs.z() == rhs.z() &&
+    lhs.w() == rhs.w();
+}
+
+/// Inequality operator
+template <typename Element>
+MUTLASS_HOST_DEVICE
+bool operator!=(Quaternion<Element> const &lhs, Quaternion<Element> const &rhs) {
+  return !(lhs == rhs);
+}
+
+/// Quaternion scalar multiplication
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator*(Quaternion<Element> q, Element s) {
+  return make_Quaternion(
+    q.x() * s,
+    q.y() * s,
+    q.z() * s,
+    q.w() * s
+  );
+}
+
+/// Quaternion scalar multiplication
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator*(Element s, Quaternion<Element> const &q) {
+  return make_Quaternion(
+    s * q.x(),
+    s * q.y(),
+    s * q.z(),
+    s * q.w()
+  );
+}
+
+/// Quaternion scalar division
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator/(Quaternion<Element> const &q, Element s) {
+  return make_Quaternion(
+    q.x() / s,
+    q.y() / s,
+    q.z() / s,
+    q.w() / s
+  );
+}
+
+/// Quaternion unary negation
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator-(Quaternion<Element> const &q) {
+  return make_Quaternion(
+    -q.x(),
+    -q.y(),
+    -q.z(),
+    -q.w()
+  );
+}
+
+/// Quaternion addition
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator+(Quaternion<Element> const &lhs, Quaternion<Element> const &rhs) {
+  return make_Quaternion(
+    lhs.x() + rhs.x(),
+    lhs.y() + rhs.y(),
+    lhs.z() + rhs.z(),
+    lhs.w() + rhs.w()
+  );
+}
+
+/// Quaternion subtraction
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator-(Quaternion<Element> const &lhs, Quaternion<Element> const &rhs) {
+  return make_Quaternion(
+    lhs.x() - rhs.x(),
+    lhs.y() - rhs.y(),
+    lhs.z() - rhs.z(),
+    lhs.w() - rhs.w()
+  );
+}
+
+/// Quaternion product
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator*(Quaternion<Element> const &lhs, Quaternion<Element> const &rhs) {
+  return make_Quaternion(
+    lhs.w() * rhs.x() + rhs.w() * lhs.x() + lhs.y() * rhs.z() - lhs.z() * rhs.y(),
+    lhs.w() * rhs.y() + rhs.w() * lhs.y() + lhs.z() * rhs.x() - lhs.x() * rhs.z(),
+    lhs.w() * rhs.z() + rhs.w() * lhs.z() + lhs.x() * rhs.y() - lhs.y() * rhs.x(),
+    lhs.w() * rhs.w() - lhs.x() * rhs.x() - lhs.y() * rhs.y() - lhs.z() * rhs.z()
+  );
+}
+
+/// Quaternion division
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator/(Quaternion<Element> const &lhs, Quaternion<Element> const &rhs) {
+  return lhs * reciprocal(rhs);
+}
+
+/// Quaternion scalar division
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Quaternion<Element> operator/(Element s, Quaternion<Element> const &q) {
+  return s * reciprocal(q);
+}
+
+/// Comparison
+template <typename Element>
+MUTLASS_HOST_DEVICE
+bool operator<(Quaternion<Element> const &lhs, Quaternion<Element> const &rhs) {
+  return true;
+}
+
+/// Rotates a 3-vector assuming this is a unit quaternion (a spinor). This avoids computing
+/// a reciprocal.
+template <typename Element>
+MUTLASS_HOST_DEVICE
+Matrix3x1<Element> spinor_rotation(
+  Quaternion<Element> const &spinor,        /// unit-length quaternion
+  Matrix3x1<Element> const &rhs) {          /// arbitrary 3-vector
+
+  return (spinor * Quaternion<Element>(rhs, 0) * conj(spinor)).pure();
+}
+
+/// Inverse rotation of 3-vector assuming this is a unit quaternion (a spinor). This avoids computing
+/// a reciprocal.
+template <typename  Element>
+MUTLASS_HOST_DEVICE
+Matrix3x1<Element> spinor_rotation_inv(
+  Quaternion<Element> const &spinor,        /// unit-length quaternion
+  Matrix3x1<Element> const &rhs) {          /// arbitrary 3-vector
+
+  return (conj(spinor) * Quaternion<Element>(rhs, 0) * spinor).pure();
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for Quaternion-valued type.
+template <typename T>
+struct RealType< Quaternion<T> > {
+  using Type = T;
+
+  /// Number of elements
+  static int const kExtent = Quaternion<T>::kExtent;
+
+MUTLASS_HOST_DEVICE
+  static Quaternion<T> from_real(double x) {
+    return Quaternion<T>(static_cast<T>(x));
+  }
+};
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+// Factories
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <>
+MUTLASS_HOST_DEVICE
+mutlass::Quaternion<half_t> from_real<mutlass::Quaternion<half_t> >(double r) {
+  return mutlass::Quaternion<half_t>(half_t(r));
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+mutlass::Quaternion<float> from_real<mutlass::Quaternion<float> >(double r) {
+  return mutlass::Quaternion<float>(float(r));
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+mutlass::Quaternion<double> from_real<mutlass::Quaternion<double> >(double r) {
+  return mutlass::Quaternion<double>(r);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+// functional.h numeric specializations
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+struct multiplies<Quaternion<T>> {
+  MUTLASS_HOST_DEVICE
+  Quaternion<T> operator()(Quaternion<T> lhs, Quaternion<T> const &rhs) const {
+    lhs = lhs * rhs;
+    return lhs;
+  }
+};
+
+/// Squares with optional conversion
+template <typename T, typename Output>
+struct magnitude_squared<Quaternion<T>, Output> {
+  MUTLASS_HOST_DEVICE
+  Output operator()(Quaternion<T> lhs) const {
+    multiplies<Output> mul_op;
+
+    Output y_w = Output(lhs.w());
+    Output y_x = Output(lhs.x());
+    Output y_y = Output(lhs.y());
+    Output y_z = Output(lhs.z());
+
+    return mul_op(y_w, y_w) + mul_op(y_x, y_x) + mul_op(y_y, y_y) + \
+           mul_op(y_z, y_z);
+  }
+};
+
+template <typename T>
+struct multiply_add<Quaternion<T>, Quaternion<T>, Quaternion<T>> {
+  MUTLASS_HOST_DEVICE
+  Quaternion<T> operator()(
+    Quaternion<T> const &a,
+    Quaternion<T> const &b,
+    Quaternion<T> const &c) const {
+
+    T x = c.x();
+    T y = c.y();
+    T z = c.z();
+    T w = c.w();
+
+    x += a.w() * b.x();
+    x += b.w() * a.x();
+    x += a.y() * b.z();
+    x += -a.z() * b.y(),
+
+    y += a.w() * b.y();
+    y += b.w() * a.y();
+    y += a.z() * b.x();
+    y += -a.x() * b.z();
+
+    z += a.w() * b.z();
+    z += b.w() * a.z();
+    z += a.x() * b.y();
+    z += -a.y() * b.x();
+
+    w += a.w() * b.w();
+    w += -a.x() * b.x();
+    w += -a.y() * b.y();
+    w += -a.z() * b.z();
+
+    return mutlass::make_Quaternion(x, y, z, w);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/real.h b/vllm_musa/csrc/flash_attn/include/mutlass/real.h
new file mode 100644
index 000000000..07e9db95a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/real.h
@@ -0,0 +1,65 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 This class provides helpers to support real<> and complex<> types in generic code.
+*/
+
+#pragma once
+
+
+#include <mutlass/detail/helper_macros.hpp> // MUTLASS_DEVICE
+
+namespace mutlass {
+
+/// Used to determine the real-valued underlying type of a numeric type T.
+template <typename T>
+struct RealType {
+  using Type = T;
+
+  /// Number of elements
+  static int const kExtent = 1;
+
+MUTLASS_HOST_DEVICE
+  static T from_real(double x) {
+    return static_cast<T>(x);
+  }
+};
+
+template <typename T>
+MUTLASS_HOST_DEVICE
+static T from_real(double r) {
+  return T(r);
+}
+
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/relatively_equal.h b/vllm_musa/csrc/flash_attn/include/mutlass/relatively_equal.h
new file mode 100644
index 000000000..0cf317d71
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/relatively_equal.h
@@ -0,0 +1,226 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Performs comparison between two elements with support for floating-point comparisons.
+*/
+
+#pragma once
+
+#include "numeric_types.h"
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+MUTLASS_HOST_DEVICE
+bool relatively_equal(T a, T b, T epsilon, T nonzero_floor);
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+// This floating-point comparison function implements the method described in
+//
+// https://floating-point-gui.de/errors/comparison/
+//
+template <typename T>
+MUTLASS_HOST_DEVICE
+bool relatively_equal_float(T a, T b, T epsilon, T nonzero_floor) {
+
+#if defined(__MUSACC_RTC__)
+  using musa::std::abs;
+#else
+  using std::abs;
+#endif
+
+  T abs_A = abs(a);
+  T abs_B = abs(b);
+  T diff = abs(a - b);
+  T zero = T(0);
+
+  if (a == b) {
+    return true;
+  }
+  else if (a == zero || b == zero || diff < nonzero_floor) {
+    return diff < epsilon * nonzero_floor;
+  }
+
+  return diff < epsilon * (abs_A + abs_B);
+}
+
+} // namespace detail
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<bool>(bool a, bool b, bool, bool) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<uint1b_t>(uint1b_t a, uint1b_t b, uint1b_t, uint1b_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<int2b_t>(int2b_t a, int2b_t b, int2b_t, int2b_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<uint2b_t>(uint2b_t a, uint2b_t b, uint2b_t, uint2b_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<int4b_t>(int4b_t a, int4b_t b, int4b_t, int4b_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<uint4b_t>(uint4b_t a, uint4b_t b, uint4b_t, uint4b_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<int8_t>(int8_t a, int8_t b, int8_t, int8_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<uint8_t>(uint8_t a, uint8_t b, uint8_t, uint8_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<int16_t>(int16_t a, int16_t b, int16_t, int16_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<uint16_t>(uint16_t a, uint16_t b, uint16_t, uint16_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<int32_t>(int32_t a, int32_t b, int32_t, int32_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<uint32_t>(uint32_t a, uint32_t b, uint32_t, uint32_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<int64_t>(int64_t a, int64_t b, int64_t, int64_t) {
+  return (a == b);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<uint64_t>(uint64_t a, uint64_t b, uint64_t, uint64_t) {
+  return (a == b);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<float_e4m3_t>(float_e4m3_t a, float_e4m3_t b, float_e4m3_t epsilon, float_e4m3_t nonzero_floor) {
+  return detail::relatively_equal_float<float>(a, b, epsilon, nonzero_floor);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<float_e5m2_t>(float_e5m2_t a, float_e5m2_t b, float_e5m2_t epsilon, float_e5m2_t nonzero_floor) {
+  return detail::relatively_equal_float<float>(a, b, epsilon, nonzero_floor);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<half_t>(half_t a, half_t b, half_t epsilon, half_t nonzero_floor) {
+  return detail::relatively_equal_float(a, b, epsilon, nonzero_floor);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<bfloat16_t>(
+  bfloat16_t a,
+  bfloat16_t b,
+  bfloat16_t epsilon,
+  bfloat16_t nonzero_floor) {
+
+  return detail::relatively_equal_float(a, b, epsilon, nonzero_floor);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<tfloat32_t>(
+  tfloat32_t a,
+  tfloat32_t b,
+  tfloat32_t epsilon,
+  tfloat32_t nonzero_floor) {
+
+  return detail::relatively_equal_float(a, b, epsilon, nonzero_floor);
+}
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<float>(float a, float b, float epsilon, float nonzero_floor) {
+  return detail::relatively_equal_float(a, b, epsilon, nonzero_floor);
+}
+
+
+template <>
+MUTLASS_HOST_DEVICE
+bool relatively_equal<double>(double a, double b, double epsilon, double nonzero_floor) {
+  return detail::relatively_equal_float(a, b, epsilon, nonzero_floor);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/subbyte_reference.h b/vllm_musa/csrc/flash_attn/include/mutlass/subbyte_reference.h
new file mode 100644
index 000000000..908956ef7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/subbyte_reference.h
@@ -0,0 +1,1350 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Provides a mechanism for packing and unpacking elements smaller than one byte
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/integer_subbyte.h"
+#include "mutlass/fast_math.h"
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// This class provides a mechanism for packing and unpacking elements smaller than one byte. It
+/// assumes these sub-byte elements are packed in a traditional C++ numeric type.
+///
+/// The intended application is to provide a mechanism to indirectly reference elements in
+/// memory or Array<> objects whose addresses cannot otherwise be taken since they are smaller
+/// than one byte.
+///
+/// Supports basic pointer arithmetic:
+///
+/// Example:
+///
+///   int4b_t *ptr = ...;
+///
+///   SubbyteReference<int4b_t> ref = ptr;
+///   ref += 15;
+///
+///   int4b_t x = ref;      // load an int4b_t
+///   ref = x + 2_s4;      // perform arithmetic on int4b_t and then store
+///
+template <
+  typename Element_,              /// MUTLASS numeric element type.
+  typename Storage_ = uint8_t,    /// Underlying storage type. Must be able to hold an integer
+                                  ///   number of objects of type Element.
+  class = void
+>
+class ConstSubbyteReference {
+public:
+
+  using Element = Element_;
+  using Storage = Storage_;
+  using StoragePointer = Storage const *;
+
+  static_assert(sizeof_bits<Element>::value <= sizeof_bits<Storage>::value,
+    "Size of Element must not be greater than Storage.");
+
+  static_assert(!(sizeof_bits<Storage>::value % sizeof_bits<Element>::value),
+    "Storage must be divisible by Element");
+
+private:
+
+  ///! Number of elements per storage vector
+  int const kElementsPerVector = sizeof_bits<Storage>::value / sizeof_bits<Element>::value;
+
+  ///! Bit mask
+  Storage const kMask =
+    ((sizeof_bits<Element>::value < sizeof_bits<Storage>::value) ?
+      (Storage(1) << sizeof_bits<Element>::value) - Storage(1) :
+      ~Storage(0));
+
+private:
+
+  /// Pointer to array containing element
+  StoragePointer ptr_;
+
+  /// Offset (in units of elements) from pointer.
+  ///
+  /// Invariant: must always be in range [0, kElementsPerVector)
+  int offset_;
+
+public:
+
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference(): ptr_(nullptr), offset_(0) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference(
+    Element const *ptr,           /// pointer to memory
+    int64_t offset          /// logical offset in units of Element
+  ):
+    ptr_(reinterpret_cast<StoragePointer>(ptr)),
+    offset_(0) {
+
+    int64_t offset_in_vectors = offset / kElementsPerVector;
+    int64_t offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ += offset_in_vectors;
+    offset_ = int(offset_in_elements);
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference(
+    Element *ptr = nullptr
+  ): ConstSubbyteReference(ptr, 0) { }
+
+  /// Gets storage pointer
+  MUTLASS_HOST_DEVICE
+  StoragePointer storage_pointer() const {
+    return ptr_;
+  }
+
+  /// Gets element offset within storage vector
+  MUTLASS_HOST_DEVICE
+  int element_offset() const {
+    return offset_;
+  }
+
+  /// Unpacks an element from memory
+  MUTLASS_HOST_DEVICE
+  Element get() const {
+    Storage item = Storage((*ptr_ >> (offset_ * sizeof_bits<Element>::value)) & kMask);
+    return reinterpret_cast<Element const &>(item);
+  }
+
+  /// Unpacks an element from memory
+  MUTLASS_HOST_DEVICE
+  operator Element() const {
+    return get();
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference &operator+=(int offset) {
+
+    offset += offset_;
+
+    int offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ += offset_in_vectors;
+    offset_ = offset_in_elements;
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference &operator+=(long long offset) {
+
+    offset += offset_;
+
+    long long offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = int(offset % kElementsPerVector);
+
+    ptr_ += offset_in_vectors;
+    offset_ = offset_in_elements;
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference &operator-=(int offset) {
+
+    int offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ -= offset_in_vectors;
+    offset_ -= offset_in_elements;
+
+    if (offset_ < 0) {
+      offset_ += kElementsPerVector;
+      --ptr_;
+    }
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference &operator-=(long long offset) {
+
+    long long offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = int(offset % kElementsPerVector);
+
+    ptr_ -= offset_in_vectors;
+    offset_ -= offset_in_elements;
+
+    if (offset_ < 0) {
+      offset_ += kElementsPerVector;
+      --ptr_;
+    }
+
+    return *this;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference operator+(int offset) const {
+
+    ConstSubbyteReference ref(ptr_, offset_);
+    ref += offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference operator+(long long offset) const {
+
+    ConstSubbyteReference ref(ptr_, offset_);
+    ref += offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference operator-(int offset) const {
+
+    ConstSubbyteReference ref(ptr_, offset_);
+    ref -= offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference operator-=(long long offset) const {
+
+    ConstSubbyteReference ref(ptr_, offset_);
+    ref -= offset;
+
+    return ref;
+  }
+
+  /// Computes the difference in elements between references
+  MUTLASS_HOST_DEVICE
+  ptrdiff_t operator-(ConstSubbyteReference ref) const {
+    return (ptr_ - ref.ptr_) * kElementsPerVector + (offset_ - ref.offset_);
+  }
+
+  /// Explicit cast to int
+  MUTLASS_HOST_DEVICE
+  explicit operator int() const {
+    return int(get());
+  }
+
+  /// Explicit cast to signed 64-bit integer
+  MUTLASS_HOST_DEVICE
+  explicit operator int64_t() const {
+    return int64_t(get());
+  }
+
+  /// Explicit cast to unsigned 64-bit integer
+  MUTLASS_HOST_DEVICE
+  explicit operator uint64_t() const {
+    return uint64_t(get());
+  }
+
+  /// Explicit cast to float
+  MUTLASS_HOST_DEVICE
+  explicit operator float() const {
+    return float(get());
+  }
+
+  /// Explicit cast to double
+  MUTLASS_HOST_DEVICE
+  explicit operator double() const {
+    return double(get());
+  }
+};
+
+template <
+  typename Element_,              /// MUTLASS numeric element type.
+  typename Storage_ =             /// Underlying storage type. Must be able to hold an integer
+                                  ///   number of objects of type Element.
+
+#if defined(__MUSA_ARCH__)        /// Default size depends on width of atomicCas() overloads.
+  #if (__MUSA_ARCH__ >= 700)      ///
+  uint16_t
+  #else
+  uint32_t
+  #endif
+#else
+  uint8_t
+#endif
+  ,
+  class = void
+>
+class SubbyteReference {
+public:
+
+  using Element = Element_;
+  using Storage = Storage_;
+  using StoragePointer = Storage *;
+
+  static_assert(sizeof_bits<Element>::value <= sizeof_bits<Storage>::value,
+    "Size of Element must not be greater than Storage.");
+
+  static_assert(!(sizeof_bits<Storage>::value % sizeof_bits<Element>::value),
+    "Storage must be divisible by Element");
+
+private:
+
+  ///! Number of elements per storage vector
+  int const kElementsPerVector = sizeof_bits<Storage>::value / sizeof_bits<Element>::value;
+
+  ///! Bit mask
+  Storage const kMask =
+    ((sizeof_bits<Element>::value < sizeof_bits<Storage>::value) ?
+      (Storage(1) << sizeof_bits<Element>::value) - Storage(1) :
+      ~Storage(0));
+
+private:
+
+  /// Pointer to array containing element
+  StoragePointer ptr_;
+
+  /// Offset (in units of elements) from pointer.
+  ///
+  /// Invariant: must always be in range [0, kElementsPerVector)
+  int offset_;
+
+public:
+
+  MUTLASS_HOST_DEVICE
+  SubbyteReference(): ptr_(nullptr), offset_(0) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  SubbyteReference(
+    Element *ptr,           /// pointer to memory
+    int64_t offset          /// logical offset in units of Element
+  ):
+    ptr_(reinterpret_cast<StoragePointer>(ptr)),
+    offset_(0) {
+
+    int64_t offset_in_vectors = offset / kElementsPerVector;
+    int64_t offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ += offset_in_vectors;
+    offset_ = int(offset_in_elements);
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  SubbyteReference(
+    Element *ptr = nullptr
+  ): SubbyteReference(ptr, 0) { }
+
+  /// Gets storage pointer
+  MUTLASS_HOST_DEVICE
+  StoragePointer storage_pointer() const {
+    return ptr_;
+  }
+
+  /// Gets storage pointer
+  MUTLASS_HOST_DEVICE
+  Element * operator&() const {
+    return reinterpret_cast<Element *>(ptr_);
+  }
+
+  /// Gets element offset within storage vector
+  MUTLASS_HOST_DEVICE
+  int element_offset() const {
+    return offset_;
+  }
+
+  /// Unpacks an element from memory
+  MUTLASS_HOST_DEVICE
+  Element get() const {
+    uint8_t const* byte_ptr = reinterpret_cast<uint8_t const*>(ptr_);
+    // Convert offset in elements to offset in bytes
+    constexpr int elements_per_byte = mutlass::sizeof_bits<uint8_t>::value / mutlass::sizeof_bits<Element>::value;
+    byte_ptr += offset_ / elements_per_byte;
+    // Offset of element within a byte
+    int byte_offset = offset_ % elements_per_byte;
+    uint8_t item = uint8_t((*byte_ptr >> (byte_offset * mutlass::sizeof_bits<Element>::value)) & kMask);
+    return reinterpret_cast<Element const &>(item);
+  }
+
+  /// Stores an element to memory
+  MUTLASS_HOST_DEVICE
+  SubbyteReference & set(Element const &x) {
+
+    Storage item        = (reinterpret_cast<Storage const &>(x) & kMask);
+    Storage kUpdateMask = Storage(~(kMask << (offset_ * mutlass::sizeof_bits<Element>::value)));
+    Storage new_bits    = Storage(item << (offset_ * mutlass::sizeof_bits<Element>::value));
+
+#if defined(__MUSA_ARCH__)
+
+    //
+    // Homebrew read-modify-write
+    //
+    Storage original;
+    Storage updated;
+
+    do {
+
+      original = (*ptr_);
+
+      updated  = Storage((original & kUpdateMask) | new_bits);
+
+      original = atomicCAS(ptr_, original, updated);
+
+    } while (updated != original);
+
+#else
+
+    Storage original = (*ptr_);
+    Storage updated  = Storage((original & kUpdateMask) | new_bits);
+    *ptr_ = updated;
+
+#endif
+
+    return *this;
+  }
+
+  ////
+
+  /// Unpacks an element from memory
+  MUTLASS_HOST_DEVICE
+  operator Element() const {
+    return get();
+  }
+
+  /// Stores an element to memory
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator=(Element const & x) {
+    return set(x);
+  }
+
+  /// Stores an element to memory
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator=(SubbyteReference const & x) {
+    return set(x.get());
+  }
+
+  /// Stores an element to memory
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator=(
+      ConstSubbyteReference<Element, Storage> const &x) {
+    return set(x.get());
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator+=(int offset) {
+
+    offset += offset_;
+
+    int offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ += offset_in_vectors;
+    offset_ = offset_in_elements;
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator+=(long long offset) {
+
+    offset += offset_;
+
+    long long offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = int(offset % kElementsPerVector);
+
+    ptr_ += offset_in_vectors;
+    offset_ = offset_in_elements;
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator-=(int offset) {
+
+    int offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ -= offset_in_vectors;
+    offset_ -= offset_in_elements;
+
+    if (offset_ < 0) {
+      offset_ += kElementsPerVector;
+      --ptr_;
+    }
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator-=(long long offset) {
+
+    long long offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = int(offset % kElementsPerVector);
+
+    ptr_ -= offset_in_vectors;
+    offset_ -= offset_in_elements;
+
+    if (offset_ < 0) {
+      offset_ += kElementsPerVector;
+      --ptr_;
+    }
+
+    return *this;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference operator+(int offset) const {
+
+    SubbyteReference ref(ptr_, offset_);
+    ref += offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference operator+(long long offset) const {
+
+    SubbyteReference ref(ptr_, offset_);
+    ref += offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference operator-(int offset) const {
+
+    SubbyteReference ref(ptr_, offset_);
+    ref -= offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference operator-=(long long offset) const {
+
+    SubbyteReference ref(ptr_, offset_);
+    ref -= offset;
+
+    return ref;
+  }
+
+  /// Computes the difference in elements between references
+  MUTLASS_HOST_DEVICE
+  ptrdiff_t operator-(SubbyteReference ref) const {
+    return (ptr_ - ref.ptr_) * kElementsPerVector + (offset_ - ref.offset_);
+  }
+
+  /// Explicit cast to int
+  MUTLASS_HOST_DEVICE
+  explicit operator int() const {
+    return int(get());
+  }
+
+  /// Explicit cast to signed 64-bit integer
+  MUTLASS_HOST_DEVICE
+  explicit operator int64_t() const {
+    return int64_t(get());
+  }
+
+  /// Explicit cast to unsigned 64-bit integer
+  MUTLASS_HOST_DEVICE
+  explicit operator uint64_t() const {
+    return uint64_t(get());
+  }
+
+  /// Explicit cast to float
+  MUTLASS_HOST_DEVICE
+  explicit operator float() const {
+    return float(get());
+  }
+
+  /// Explicit cast to double
+  MUTLASS_HOST_DEVICE
+  explicit operator double() const {
+    return double(get());
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T> using _war = T;
+template <
+  typename Element_,              /// MUTLASS numeric element type.
+  typename Storage_               /// Underlying basic storage type.
+>
+class SubbyteReference<Element_, Storage_,
+    typename platform::enable_if<sizeof_bits<Storage_>::value % sizeof_bits<Element_>::value != 0>::type> {
+public:
+
+  using Element = Element_;
+  ///! Note: Storage unit could not be divisibale by Element,
+  ///   Type element may be stored across 2 storage units, so need a storage vector to hold integer
+  ///   number of objects of type Element.
+  using StorageUnit = Storage_;
+  static int const kBitsStoredVec = mutlass::lcm_cxx11(sizeof_bits<Element>::value, sizeof_bits<StorageUnit>::value);
+  static int const kNumStorageUnitPerStoredVec = kBitsStoredVec / sizeof_bits<StorageUnit>::value;
+
+  using StorageVec = StorageUnit[kNumStorageUnitPerStoredVec];
+  using StorageVecPointer = StorageVec *;
+
+  using MusaAtomicType = typename platform::conditional<
+      sizeof_bits<StorageUnit>::value == 16,
+      uint32_t,
+      uint64_t
+    >::type;
+
+  static_assert(sizeof_bits<Element>::value <= sizeof_bits<StorageVec>::value,
+    "Size of Element must not be greater than StorageVec.");
+
+  static_assert(!(sizeof_bits<StorageVec>::value % sizeof_bits<Element>::value),
+    "StorageVec must be divisible by Element");
+
+private:
+
+  ///! Number of elements per storage vector
+  int const kElementsPerVector = sizeof_bits<StorageVec>::value / sizeof_bits<Element>::value;
+
+  ///! Bit mask for storage unit.
+  StorageUnit const kMask = (StorageUnit(1) << sizeof_bits<Element>::value) - StorageUnit(1);
+
+  /// Pointer to array containing element
+  _war<StorageVecPointer> ptr_;
+
+  /// Offset (in units of elements) from pointer.
+  ///
+  /// Invariant: must always be in range [0, kElementsPerVector)
+  int offset_;
+
+  /// Element may be stored across 2 storage unit.
+  ///   Low storage unit index in StorageVec
+  ///   High storage unit index in StorageVec
+  int low_storage_unit_idx_;
+  int high_storage_unit_idx_;
+
+  /// Full Mask to extract the entire element
+  uint64_t full_element_mask_;
+
+  /// Mask to extract the Element from Low storage unit and High storage unit.
+  StorageUnit low_storage_mask_;
+  StorageUnit high_storage_mask_;
+
+  /// Start bit index inside the storage unit.
+  int start_bit_idx_;
+
+private:
+
+  MUTLASS_HOST_DEVICE
+  void update_element_status() {
+    int num_bits = offset_ * sizeof_bits<Element>::value;
+
+    start_bit_idx_ = num_bits % sizeof_bits<StorageUnit>::value;
+
+    low_storage_unit_idx_ = num_bits / sizeof_bits<StorageUnit>::value;
+    high_storage_unit_idx_ = sizeof_bits<StorageUnit>::value - (start_bit_idx_) < sizeof_bits<Element>::value
+                              ? low_storage_unit_idx_ + 1 : low_storage_unit_idx_;
+
+    full_element_mask_ = uint64_t(kMask) << start_bit_idx_;
+    low_storage_mask_ = StorageUnit(full_element_mask_ & ~StorageUnit(0));
+    high_storage_mask_ = StorageUnit((full_element_mask_ >> sizeof_bits<StorageUnit>::value) & ~StorageUnit(0));
+  }
+
+public:
+
+  MUTLASS_HOST_DEVICE
+  SubbyteReference(): ptr_(nullptr), offset_(0) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  SubbyteReference(
+    Element *ptr,           /// pointer to memory
+    int64_t offset          /// logical offset in units of Element
+  ):
+    ptr_(reinterpret_cast<StorageVecPointer>(ptr)),
+    offset_(0) {
+    int64_t offset_in_vectors = offset / kElementsPerVector;
+    int64_t offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ += offset_in_vectors;
+    offset_ = int(offset_in_elements);
+
+    update_element_status();
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  SubbyteReference(
+    Element *ptr = nullptr
+  ): SubbyteReference(ptr, 0) { }
+
+  /// Gets StorageVec pointer
+  MUTLASS_HOST_DEVICE
+  StorageVecPointer storage_pointer() const {
+    return ptr_;
+  }
+
+  /// Gets StorageVec pointer
+  MUTLASS_HOST_DEVICE
+  Element * operator&() const {
+    return reinterpret_cast<Element *>(ptr_);
+  }
+
+  /// Gets element offset within StorageVec vector
+  MUTLASS_HOST_DEVICE
+  int element_offset() const {
+    return offset_;
+  }
+
+  /// Unpacks an element from memory
+  MUTLASS_HOST_DEVICE
+  Element get() const {
+    StorageUnit low_bits = (*ptr_)[low_storage_unit_idx_] & low_storage_mask_;
+    StorageUnit high_bits = low_storage_unit_idx_ != high_storage_unit_idx_ ? (*ptr_)[high_storage_unit_idx_] & high_storage_mask_ : 0;
+
+    uint64_t full_item = ((uint64_t)high_bits << sizeof_bits<StorageUnit>::value) | low_bits;
+    uint8_t result = uint8_t(full_item >> start_bit_idx_);
+
+    return reinterpret_cast<Element const &>(result);
+  }
+
+  /// Stores an element to memory
+  MUTLASS_HOST_DEVICE
+  SubbyteReference & set(Element const &x) {
+
+    uint64_t item = static_cast<uint64_t>((reinterpret_cast<uint8_t const &>(x) & kMask)) << start_bit_idx_;
+
+    StorageUnit low_new_bits  = StorageUnit(item & ~StorageUnit(0));
+    StorageUnit high_new_bits = StorageUnit(item >> sizeof_bits<StorageUnit>::value);
+
+    StorageUnit const kLowUpdateMask  = StorageUnit((~full_element_mask_) & (~StorageUnit(0)));
+    StorageUnit const kHighUpdateMask = StorageUnit(((~full_element_mask_) >> sizeof_bits<StorageUnit>::value) & (~StorageUnit(0)));
+
+#if defined(__MUSA_ARCH__)
+    //
+    // Homebrew read-modify-write
+    //
+    if(high_storage_unit_idx_ != low_storage_unit_idx_){
+      /// Only need update 2 storage unit at once.
+      /// consider misaligned address issue, we need to do atomicCAS twice
+      StorageUnit original_low_bits, original_high_bits, update_low_bits, update_high_bits;
+      do {
+        original_low_bits  = ((*ptr_)[low_storage_unit_idx_]);
+        update_low_bits  = (original_low_bits & kLowUpdateMask) | low_new_bits;
+        original_low_bits = atomicCAS(&((*ptr_)[low_storage_unit_idx_]), original_low_bits, update_low_bits);
+      } while (update_low_bits != original_low_bits);
+      do {
+        original_high_bits = ((*ptr_)[high_storage_unit_idx_]);
+        update_high_bits  = (original_high_bits & kHighUpdateMask) | high_new_bits;
+        original_high_bits = atomicCAS(&((*ptr_)[high_storage_unit_idx_]), original_high_bits, update_high_bits);
+      } while (update_high_bits != original_high_bits);
+    }
+    else {
+      /// Only need update 1 storage unit.
+      StorageUnit original, updated;
+      do {
+        original = ((*ptr_)[low_storage_unit_idx_]);
+
+        updated = (original & kLowUpdateMask) | low_new_bits;
+
+        original = atomicCAS(&((*ptr_)[low_storage_unit_idx_]), original, updated);
+
+      } while (updated != original);
+    }
+#else
+
+
+    StorageUnit update_low_bits  = ((*ptr_)[low_storage_unit_idx_] & kLowUpdateMask) | low_new_bits;
+    StorageUnit update_high_bits = ((*ptr_)[high_storage_unit_idx_] & kHighUpdateMask) | high_new_bits;
+
+    (*ptr_)[low_storage_unit_idx_] = update_low_bits;
+
+    if(low_storage_unit_idx_ != high_storage_unit_idx_)
+      (*ptr_)[high_storage_unit_idx_] = update_high_bits;
+#endif
+
+    return *this;
+  }
+
+  ////
+
+  /// Unpacks an element from memory
+  MUTLASS_HOST_DEVICE
+  operator Element() const {
+    return get();
+  }
+
+  /// Stores an element to memory
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator=(Element const & x) {
+    return set(x);
+  }
+
+  /// Stores an element to memory
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator=(SubbyteReference const & x) {
+    return set(x.get());
+  }
+
+  /// Stores an element to memory
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator=(
+      ConstSubbyteReference<Element, StorageVec> const &x) {
+    return set(x.get());
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator+=(int offset) {
+
+    offset += offset_;
+
+    int offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ += offset_in_vectors;
+    offset_ = offset_in_elements;
+
+    update_element_status();
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator+=(long long offset) {
+
+    offset += offset_;
+
+    long long offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = int(offset % kElementsPerVector);
+
+    ptr_ += offset_in_vectors;
+    offset_ = offset_in_elements;
+
+    update_element_status();
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator-=(int offset) {
+
+    int offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ -= offset_in_vectors;
+    offset_ -= offset_in_elements;
+
+    if (offset_ < 0) {
+      offset_ += kElementsPerVector;
+      --ptr_;
+    }
+
+    update_element_status();
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference &operator-=(long long offset) {
+
+    long long offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = int(offset % kElementsPerVector);
+
+    ptr_ -= offset_in_vectors;
+    offset_ -= offset_in_elements;
+
+    if (offset_ < 0) {
+      offset_ += kElementsPerVector;
+      --ptr_;
+    }
+
+    update_element_status();
+    return *this;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference operator+(int offset) const {
+
+    SubbyteReference ref(ptr_, offset_);
+    ref += offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference operator+(long long offset) const {
+
+    SubbyteReference ref(ptr_, offset_);
+    ref += offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference operator-(int offset) const {
+
+    SubbyteReference ref(ptr_, offset_);
+    ref -= offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  SubbyteReference operator-=(long long offset) const {
+
+    SubbyteReference ref(ptr_, offset_);
+    ref -= offset;
+
+    return ref;
+  }
+
+  /// Computes the difference in elements between references
+  MUTLASS_HOST_DEVICE
+  ptrdiff_t operator-(SubbyteReference ref) const {
+    return (ptr_ - ref.ptr_) * kElementsPerVector + (offset_ - ref.offset_);
+  }
+
+  /// Explicit cast to int
+  MUTLASS_HOST_DEVICE
+  explicit operator int() const {
+    return int(get());
+  }
+
+  /// Explicit cast to signed 64-bit integer
+  MUTLASS_HOST_DEVICE
+  explicit operator int64_t() const {
+    return int64_t(get());
+  }
+
+  /// Explicit cast to unsigned 64-bit integer
+  MUTLASS_HOST_DEVICE
+  explicit operator uint64_t() const {
+    return uint64_t(get());
+  }
+
+  /// Explicit cast to float
+  MUTLASS_HOST_DEVICE
+  explicit operator float() const {
+    return float(get());
+  }
+
+  /// Explicit cast to double
+  MUTLASS_HOST_DEVICE
+  explicit operator double() const {
+    return double(get());
+  }
+};
+
+template<typename T> using _war = T;
+template <
+  typename Element_,              /// MUTLASS numeric element type.
+  typename Storage_               /// Underlying storage type. Must be able to hold an integer
+>
+class ConstSubbyteReference<Element_, Storage_,
+    typename platform::enable_if<sizeof_bits<Storage_>::value % sizeof_bits<Element_>::value != 0>::type> {
+public:
+
+  using Element = Element_;
+  ///! Note: Storage unit could not be divisibale by Element,
+  ///   Type element may be stored across 2 storage units, so need a storage vector to hold integer
+  ///   number of objects of type Element.
+  using StorageUnit = Storage_;
+  static int const kBitsStoredVec = mutlass::lcm_cxx11(sizeof_bits<Element>::value, sizeof_bits<StorageUnit>::value);
+  static int const kNumStorageUnitPerStoredVec = kBitsStoredVec / sizeof_bits<StorageUnit>::value;
+
+  using StorageVec = StorageUnit[kNumStorageUnitPerStoredVec];
+  using StorageVecPointer = StorageVec const *;
+
+  using MusaAtomicType = typename platform::conditional<
+      sizeof_bits<StorageUnit>::value == 16,
+      uint32_t,
+      uint64_t
+    >::type;
+
+  static_assert(sizeof_bits<Element>::value <= sizeof_bits<StorageVec>::value,
+    "Size of Element must not be greater than StorageVec.");
+
+  static_assert(!(sizeof_bits<StorageVec>::value % sizeof_bits<Element>::value),
+    "StorageVec must be divisible by Element");
+
+private:
+
+  ///! Number of elements per storage vector
+  int const kElementsPerVector = sizeof_bits<StorageVec>::value / sizeof_bits<Element>::value;
+
+  ///! Bit mask for storage unit.
+  StorageUnit const kMask = (StorageUnit(1) << sizeof_bits<Element>::value) - StorageUnit(1);
+
+  /// Pointer to array containing element
+  _war<StorageVecPointer> ptr_;
+
+  /// Offset (in units of elements) from pointer.
+  ///
+  /// Invariant: must always be in range [0, kElementsPerVector)
+  int offset_;
+
+  /// Element may be stored across 2 storage unit.
+  ///   Low storage unit index in StorageVec
+  ///   High storage unit index in StorageVec
+  int low_storage_unit_idx_;
+  int high_storage_unit_idx_;
+
+  /// Full Mask to extract the entire element
+  uint64_t full_element_mask_;
+
+  /// Mask to extract the Element from Low storage unit and High storage unit.
+  StorageUnit low_storage_mask_;
+  StorageUnit high_storage_mask_;
+
+  /// Start bit index inside the storage unit.
+  int start_bit_idx_;
+
+private:
+
+  MUTLASS_HOST_DEVICE
+  void update_element_status() {
+    int num_bits = offset_ * sizeof_bits<Element>::value;
+
+    start_bit_idx_ = num_bits % sizeof_bits<StorageUnit>::value;
+
+    low_storage_unit_idx_ = num_bits / sizeof_bits<StorageUnit>::value;
+    high_storage_unit_idx_ = sizeof_bits<StorageUnit>::value - (start_bit_idx_) < sizeof_bits<Element>::value
+                              ? low_storage_unit_idx_ + 1 : low_storage_unit_idx_;
+
+    full_element_mask_ = uint64_t(kMask) << start_bit_idx_;
+    low_storage_mask_ = StorageUnit(full_element_mask_ & ~StorageUnit(0));
+    high_storage_mask_ = StorageUnit((full_element_mask_ >> sizeof_bits<StorageUnit>::value) & ~StorageUnit(0));
+  }
+
+public:
+
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference(): ptr_(nullptr), offset_(0) { }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference(
+    Element const *ptr,           /// pointer to memory
+    int64_t offset          /// logical offset in units of Element
+  ):
+    ptr_(reinterpret_cast<StorageVecPointer>(ptr)),
+    offset_(0) {
+
+    int64_t offset_in_vectors = offset / kElementsPerVector;
+    int64_t offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ += offset_in_vectors;
+    offset_ = int(offset_in_elements);
+
+    update_element_status();
+  }
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference(
+    Element *ptr = nullptr
+  ): ConstSubbyteReference(ptr, 0) { }
+
+  /// Gets storage pointer
+  MUTLASS_HOST_DEVICE
+  StorageVecPointer storage_pointer() const {
+    return ptr_;
+  }
+
+  /// Gets element offset within storage vector
+  MUTLASS_HOST_DEVICE
+  int element_offset() const {
+    return offset_;
+  }
+
+  /// Unpacks an element from memory
+  MUTLASS_HOST_DEVICE
+  Element get() const {
+    StorageUnit low_bits = (*ptr_)[low_storage_unit_idx_] & low_storage_mask_;
+    StorageUnit high_bits = low_storage_unit_idx_ != high_storage_unit_idx_ ? (*ptr_)[high_storage_unit_idx_] & high_storage_mask_ : 0;
+
+    uint64_t full_item = ((uint64_t)high_bits << sizeof_bits<StorageUnit>::value) | low_bits;
+    uint8_t result = uint8_t(full_item >> start_bit_idx_);
+
+    return reinterpret_cast<Element const &>(result);
+  }
+
+  /// Unpacks an element from memory
+  MUTLASS_HOST_DEVICE
+  operator Element() const {
+    return get();
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference &operator+=(int offset) {
+
+    offset += offset_;
+
+    int offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ += offset_in_vectors;
+    offset_ = offset_in_elements;
+
+    update_element_status();
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference &operator+=(long long offset) {
+
+    offset += offset_;
+
+    long long offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = int(offset % kElementsPerVector);
+
+    ptr_ += offset_in_vectors;
+    offset_ = offset_in_elements;
+
+    update_element_status();
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference &operator-=(int offset) {
+
+    int offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = offset % kElementsPerVector;
+
+    ptr_ -= offset_in_vectors;
+    offset_ -= offset_in_elements;
+
+    if (offset_ < 0) {
+      offset_ += kElementsPerVector;
+      --ptr_;
+    }
+
+    update_element_status();
+
+    return *this;
+  }
+
+  /// Adds an offset in units of elements to the reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference &operator-=(long long offset) {
+
+    long long offset_in_vectors = offset / kElementsPerVector;
+    int offset_in_elements = int(offset % kElementsPerVector);
+
+    ptr_ -= offset_in_vectors;
+    offset_ -= offset_in_elements;
+
+    if (offset_ < 0) {
+      offset_ += kElementsPerVector;
+      --ptr_;
+    }
+
+    update_element_status();
+
+    return *this;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference operator+(int offset) const {
+
+    ConstSubbyteReference ref(ptr_, offset_);
+    ref += offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference operator+(long long offset) const {
+
+    ConstSubbyteReference ref(ptr_, offset_);
+    ref += offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference operator-(int offset) const {
+
+    ConstSubbyteReference ref(ptr_, offset_);
+    ref -= offset;
+
+    return ref;
+  }
+
+  /// Returns a reference to an element with a given offset from the current reference
+  MUTLASS_HOST_DEVICE
+  ConstSubbyteReference operator-=(long long offset) const {
+
+    ConstSubbyteReference ref(ptr_, offset_);
+    ref -= offset;
+
+    return ref;
+  }
+
+  /// Computes the difference in elements between references
+  MUTLASS_HOST_DEVICE
+  ptrdiff_t operator-(ConstSubbyteReference ref) const {
+    return (ptr_ - ref.ptr_) * kElementsPerVector + (offset_ - ref.offset_);
+  }
+
+  /// Explicit cast to int
+  MUTLASS_HOST_DEVICE
+  explicit operator int() const {
+    return int(get());
+  }
+
+  /// Explicit cast to signed 64-bit integer
+  MUTLASS_HOST_DEVICE
+  explicit operator int64_t() const {
+    return int64_t(get());
+  }
+
+  /// Explicit cast to unsigned 64-bit integer
+  MUTLASS_HOST_DEVICE
+  explicit operator uint64_t() const {
+    return uint64_t(get());
+  }
+
+  /// Explicit cast to float
+  MUTLASS_HOST_DEVICE
+  explicit operator float() const {
+    return float(get());
+  }
+
+  /// Explicit cast to double
+  MUTLASS_HOST_DEVICE
+  explicit operator double() const {
+    return double(get());
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Element, bool subbyte = (sizeof_bits<Element>::value < 8)>
+struct ReferenceFactory;
+
+template <typename Element>
+struct ReferenceFactory<Element, false> {
+
+  ///! Number of elements per storage vector
+  static int const kElementsPerVector = 1;
+
+  MUTLASS_HOST_DEVICE
+  static Element &get(Element *ptr, int64_t offset) {
+    return ptr[offset];
+  }
+
+  MUTLASS_HOST_DEVICE
+  static Element const &get(Element const *ptr, int64_t offset) {
+    return ptr[offset];
+  }
+
+  MUTLASS_HOST_DEVICE
+  static Element *add_pointer_offset(Element *ptr, int64_t offset) {
+    return ptr + offset;
+  }
+
+  MUTLASS_HOST_DEVICE
+  static Element const *add_pointer_offset(Element const *ptr, int64_t offset) {
+    return ptr + offset;
+  }
+};
+
+template <typename Element>
+struct ReferenceFactory<Element, true> {
+
+  //
+  // Static methods
+  //
+
+  MUTLASS_HOST_DEVICE
+  static SubbyteReference<Element> get(Element *ptr, int64_t offset) {
+    return SubbyteReference<Element>(ptr, offset);
+  }
+
+  MUTLASS_HOST_DEVICE
+  static ConstSubbyteReference<Element> get(Element const *ptr,
+                                             int64_t offset) {
+    return ConstSubbyteReference<Element>(ptr, offset);
+  }
+
+  /// Helper to add an offset in number of elements, assuming this offset is divisible
+  /// by the vector size.
+  MUTLASS_HOST_DEVICE
+  static Element *add_pointer_offset(Element *ptr, int64_t offset_in_elements) {
+
+    return ptr + offset_in_elements * sizeof_bits<Element>::value / sizeof(Element) / 8;
+  }
+
+  /// Helper to add an offset in number of elements, assuming this offset is divisible
+  /// by the vector size.
+  MUTLASS_HOST_DEVICE
+  static Element const *add_pointer_offset(Element const *ptr, int64_t offset_in_elements) {
+
+    return ptr + offset_in_elements * sizeof_bits<Element>::value / sizeof(Element) / 8;
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/tensor_coord.h b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_coord.h
new file mode 100644
index 000000000..7fb59bed7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_coord.h
@@ -0,0 +1,327 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a canonical coordinate for rank=4 tensors offering named indices.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines a canonical 4D coordinate used by tensor operations.
+struct Tensor4DCoord : public Coord<4> {
+
+  /// Base class
+  using Base = Coord<4>;
+
+  /// Index type
+  using Index = typename Base::Index;
+
+  /// LongIndex type
+  using LongIndex = typename Base::LongIndex;
+
+  /// Batch dimension
+  static int const kN = 0;
+
+  /// Height dimension
+  static int const kH = 1;
+
+  /// Width dimension
+  static int const kW = 2;
+
+  /// Channels dimension
+  static int const kC = 3;
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord() { }
+
+  /// Constructs from Coord<4>
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord(Coord<4> const &coord): Base(coord) { }
+
+  /// Helper to construct from N, H, W, and C.
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord(Index n, Index h, Index w, Index c): Base(make_Coord(n, h, w, c)) { }
+
+  /// Helper to construct from N, H, W, and C, which are LongIndex type
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord(LongIndex n, LongIndex h, LongIndex w, LongIndex c)
+    : Base(make_Coord(Index(n), Index(h), Index(w), Index(c))) { }
+
+  /// Returns the batch of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & n() const { return this->at(kN); }
+
+  /// Returns the batch of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & n() { return this->at(kN); }
+
+  /// Returns the row of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & h() const { return this->at(kH); }
+
+  /// Returns the row of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & h() { return this->at(kH); }
+
+  /// Returns the column of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & w() const { return this->at(kW); }
+
+  /// Returns the column of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & w() { return this->at(kW); }
+
+  /// Returns the channel of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & c() const { return this->at(kC); }
+
+  /// Returns the channel of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & c() { return this->at(kC); }
+
+  //
+  // Coord operators
+  //
+
+  /// Element-wise addition
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord operator+(Base const& b) const {
+    return Tensor4DCoord(Base::operator+(b));
+  }
+
+  /// Element-wise subtraction
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord operator-(Base const& b) const {
+    return Tensor4DCoord(Base::operator-(b));
+  }
+
+  /// Element-wise multiplication
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord operator*(Base const& b) const {
+    return Tensor4DCoord(Base::operator*(b));
+  }
+
+  /// Element-wise division
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord operator/(Base const& b) const {
+    return Tensor4DCoord(Base::operator/(b));
+  }
+
+  /// In-place addition
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord& operator+=(Base const& b) {
+    Base::operator+=(b);
+    return *this;
+  }
+
+  /// In-place subtraction
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord& operator-=(Base const& b) {
+    Base::operator-=(b);
+    return *this;
+  }
+
+  /// In-place multiplication
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord& operator*=(Base const& b) {
+    Base::operator*=(b);
+    return *this;
+  }
+
+  /// In-place division
+  MUTLASS_HOST_DEVICE
+  Tensor4DCoord& operator/=(Base const& b) {
+    Base::operator/=(b);
+    return *this;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines a canonical 5D coordinate used by tensor operations.
+struct Tensor5DCoord : public Coord<5> {
+
+  /// Base class
+  using Base = Coord<5>;
+
+  /// Index type
+  using Index = typename Base::Index;
+
+  /// LongIndex type
+  using LongIndex = typename Base::LongIndex;
+
+  /// Batch dimension
+  static int const kN = 0;
+
+  /// Depth dimension
+  static int const kD = 1;
+
+  /// Height dimension
+  static int const kH = 2;
+
+  /// Width dimension
+  static int const kW = 3;
+
+  /// Channels dimension
+  static int const kC = 4;
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord() { }
+
+  /// Constructs from Coord<5>
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord(Coord<5> const &coord): Base(coord) { }
+
+  /// Helper to construct from N, D, H, W, and C.
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord(Index n, Index d, Index h, Index w, Index c): Base(make_Coord(n, d, h, w, c)) { }
+
+  /// Helper to construct from N, D, H, W, and C, which are LongIndex type
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord(LongIndex n, LongIndex d, LongIndex h, LongIndex w, LongIndex c)
+    : Base(make_Coord(Index(n), Index(d), Index(h), Index(w), Index(c))) { }
+
+  /// Returns the batch of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & n() const { return this->at(kN); }
+
+  /// Returns the batch of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & n() { return this->at(kN); }
+
+  /// Returns the batch of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & d() const { return this->at(kD); }
+
+  /// Returns the batch of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & d() { return this->at(kD); }
+
+  /// Returns the row of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & h() const { return this->at(kH); }
+
+  /// Returns the row of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & h() { return this->at(kH); }
+
+  /// Returns the column of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & w() const { return this->at(kW); }
+
+  /// Returns the column of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & w() { return this->at(kW); }
+
+  /// Returns the channel of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index const & c() const { return this->at(kC); }
+
+  /// Returns the channel of the coordinate
+  MUTLASS_HOST_DEVICE
+  Index & c() { return this->at(kC); }
+
+  //
+  // Coord operators
+  //
+
+  /// Element-wise addition
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord operator+(Base const& b) const {
+    return Tensor5DCoord(Base::operator+(b));
+  }
+
+  /// Element-wise subtraction
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord operator-(Base const& b) const {
+    return Tensor5DCoord(Base::operator-(b));
+  }
+
+  /// Element-wise multiplication
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord operator*(Base const& b) const {
+    return Tensor5DCoord(Base::operator*(b));
+  }
+
+  /// Element-wise division
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord operator/(Base const& b) const {
+    return Tensor5DCoord(Base::operator/(b));
+  }
+
+  /// In-place addition
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord& operator+=(Base const& b) {
+    Base::operator+=(b);
+    return *this;
+  }
+
+  /// In-place subtraction
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord& operator-=(Base const& b) {
+    Base::operator-=(b);
+    return *this;
+  }
+
+  /// In-place multiplication
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord& operator*=(Base const& b) {
+    Base::operator*=(b);
+    return *this;
+  }
+
+  /// In-place division
+  MUTLASS_HOST_DEVICE
+  Tensor5DCoord& operator/=(Base const& b) {
+    Base::operator/=(b);
+    return *this;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/tensor_ref.h b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_ref.h
new file mode 100644
index 000000000..83e8328e0
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_ref.h
@@ -0,0 +1,419 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a structure containing strides, bounds, and a pointer to tensor data.
+*/
+#pragma once
+
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+#include "mutlass/platform/platform.h"
+#include "mutlass/subbyte_reference.h"
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Default layout function from coordinates in a tensor's index space into the n-D array held
+/// in memory.
+///
+/// All layout functions must define at least the members shown in IdentityTensorLayout<>.
+template <int Rank>
+class IdentityTensorLayout {
+public:
+  /// Logical rank of tensor
+  static int const kRank = Rank;
+
+  /// Rank of stride vector
+  static int const kStrideRank = Rank;
+
+  /// Index type used for coordinates
+  using Index = int32_t;
+
+  /// Long index type used for offsets
+  using LongIndex = int64_t;
+
+  /// Logical coordinate
+  using TensorCoord = Coord<kRank, Index>;
+
+  /// Stride vector
+  using Stride = Coord<kStrideRank, Index>;
+
+private:
+
+  //
+  // Data members
+  //
+
+  /// Stride data member
+  Stride stride_;
+
+public:
+
+  //
+  // Methods
+  //
+
+  MUTLASS_HOST_DEVICE
+  IdentityTensorLayout(Stride const &stride = Stride()): stride_(stride) { }
+
+  /// Returns the offset of a coordinate in linear memory
+  MUTLASS_HOST_DEVICE
+  LongIndex operator()(Coord<Rank> const &coord) const {
+    return coord.dot(stride_);
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return stride_;
+  }
+
+  /// Returns the stride of the layout
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return stride_;
+  }
+
+  /// Compute the number of contiguous elements needed to store a tensor with the given size
+  MUTLASS_HOST_DEVICE
+  LongIndex capacity(TensorCoord const &size) const {
+    int idx = stride_.max_dim_index();
+    return stride_[idx] * size[idx];
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/* \brief TensorRef is a template for objects pointing to the start of tensors of arbitrary rank
+          and layout within memory. A TensorRef combines a pointer and a Layout concept
+
+  Examples:
+
+  (These examples use helpers for matrix layouts defined in mutlass/layout/matrix.h)
+
+  1. Column-major matrix may be represented as a rank=2 tensor:
+
+    TensorRef<float, layout::ColumnMajor> A(ptr_A, ldm);
+
+  2. Row-major matrix may be represented as a rank=2 tensor:
+
+    TensorRef<float, layout::RowMajor> B(ptr_A, ldm);
+
+  3. An interleaved matrix may be represented as a rank=2 tensor:
+
+    TensorRef<int8_t, layout::ColumnMajorInterleaved<32> > C;
+
+  4. A helper exists to define a TensorRef for a contiguous matrix whose layout
+     is not known at compile time.
+
+    int ldm;                     // leading dimension
+    layout::Matrix kind;         // Could be layout::Matrix::kRowMajor or layout::Matrix::kColumnMajor
+
+
+    TensorRef<int, layout::ContiguousMatrix> E(ptr_E, {ldm, kind});
+
+*/
+template <
+  /// Data type of element stored within tensor (concept: NumericType)
+  typename Element_,
+  /// Defines a mapping from logical coordinate to linear memory (concept: Layout)
+  typename Layout_
+>
+class TensorRef {
+ public:
+  /// Data type of individual access
+  using Element = Element_;
+
+  /// Mapping function from logical coordinate to linear memory
+  using Layout = Layout_;
+
+  /// Reference type to an element
+  using Reference = typename platform::conditional<
+    sizeof_bits<Element>::value >= 8,
+    Element &,
+    SubbyteReference<Element>
+    >::type;
+
+  /// Logical rank of tensor index space
+  static int const kRank = Layout::kRank;
+
+  /// Index type
+  using Index = typename Layout::Index;
+
+  /// Long index used for pointer offsets
+  using LongIndex = typename Layout::LongIndex;
+
+  /// Coordinate in logical tensor space
+  using TensorCoord = typename Layout::TensorCoord;
+
+  /// Layout's stride vector
+  using Stride = typename Layout::Stride;
+
+  /// TensorRef to constant data
+  using ConstTensorRef = TensorRef<
+    typename platform::remove_const<Element>::type const,
+    Layout>;
+
+  /// TensorRef to non-constant data
+  using NonConstTensorRef = TensorRef<
+    typename platform::remove_const<Element>::type,
+    Layout>;
+
+  /// Require at least rank=1. Mathematically, a rank=0 tensor would be considered to be a
+  /// scalar, but degenerate cases such as these are difficult to accommodate without
+  /// extensive C++ metaprogramming or support for zero-length arrays.
+  static_assert(kRank > 0, "Cannot define a zero-rank TensorRef");
+
+ private:
+
+  /// Pointer
+  Element* ptr_;
+
+  /// Layout object maps logical coordinates to linear offsets
+  Layout layout_;
+
+ public:
+
+  //
+  // Methods
+  //
+
+  /// Constructs a TensorRef with a pointer and layout object.
+  MUTLASS_HOST_DEVICE
+  TensorRef(): ptr_(nullptr) {
+
+  }
+
+  /// Constructs a TensorRef with a pointer and layout object.
+  MUTLASS_HOST_DEVICE
+  TensorRef(
+    Element *ptr,                   ///< pointer to start of tensor
+    Layout const &layout            ///< layout object containing stride and mapping function
+  ):
+    ptr_(ptr), layout_(layout) {
+
+  }
+
+  /// Converting constructor from TensorRef to non-constant data.
+  template<typename _Magic = int>
+  MUTLASS_HOST_DEVICE
+  TensorRef(
+    NonConstTensorRef const &ref,              ///< TensorRef to non-const data
+    ///SFINAE trick to avoid creating a copy-constructor when Element_ is already non-const
+    _Magic magic = (typename platform::enable_if< ! platform::is_same<NonConstTensorRef, TensorRef<Element_, Layout_> >::value, _Magic>::type)0
+  ):
+    ptr_(ref.data()), layout_(ref.layout()) { }
+
+  /// Returns a reference to constant-valued tensor.
+  MUTLASS_HOST_DEVICE
+  ConstTensorRef const_ref() const {
+    return ConstTensorRef(ptr_, layout_);
+  }
+
+  MUTLASS_HOST_DEVICE
+  NonConstTensorRef non_const_ref() const {
+    return NonConstTensorRef(const_cast<typename platform::remove_const<Element>::type *>(ptr_), layout_);
+  }
+
+  /// Updates only the pointer
+  MUTLASS_HOST_DEVICE
+  void reset(Element* ptr = nullptr) {
+    ptr_ = ptr;
+  }
+
+  /// Updates the pointer and layout object
+  MUTLASS_HOST_DEVICE
+  void reset(Element* ptr, Layout const &layout) {
+    ptr_ = ptr;
+    layout_ = layout;
+  }
+
+  /// Returns true if the TensorRef is non-null
+  MUTLASS_HOST_DEVICE
+  bool good() const {
+    return ptr_ != nullptr;
+  }
+
+  /// Returns the pointer to referenced data
+  MUTLASS_HOST_DEVICE
+  Element * data() const { return ptr_; }
+
+  /// Returns a reference to the element at a given linear index
+  MUTLASS_HOST_DEVICE
+  Reference data(LongIndex idx) const {
+    return ReferenceFactory<typename platform::remove_const<Element>::type,
+                            (sizeof_bits<Element>::value < 8)>::get(ptr_, idx);
+  }
+
+  /// Returns the layout object
+  MUTLASS_HOST_DEVICE
+  Layout & layout() {
+    return layout_;
+  }
+
+  /// Returns the layout object
+  MUTLASS_HOST_DEVICE
+  Layout layout() const {
+    return layout_;
+  }
+
+  /// Returns the layout object's stride vector
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return layout_.stride();
+  }
+
+  /// Returns the layout object's stride vector
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return layout_.stride();
+  }
+
+  /// Returns the layout object's stride in a given physical dimension
+  MUTLASS_HOST_DEVICE
+  typename Layout::Stride::Index stride(int dim) const {
+    return layout_.stride().at(dim);
+  }
+
+  /// Returns the layout object's stride in a given physical dimension
+  MUTLASS_HOST_DEVICE
+  typename Layout::Stride::Index & stride(int dim) {
+    return layout_.stride().at(dim);
+  }
+
+  /// Computes the offset of an index from the origin of the tensor
+  MUTLASS_HOST_DEVICE
+  LongIndex offset(TensorCoord const& coord) const {
+    return layout_(coord);
+  }
+
+  /// Returns a reference to the element at a given Coord
+  MUTLASS_HOST_DEVICE
+  Reference at(TensorCoord const& coord) const {
+    return data(offset(coord));
+  }
+
+  /// Returns a reference to the element at a given Coord
+  MUTLASS_HOST_DEVICE
+  Reference operator[](TensorCoord const& coord) const {
+    return data(offset(coord));
+  }
+
+  /// Adds an offset to each pointer
+  MUTLASS_HOST_DEVICE
+  TensorRef & add_pointer_offset(LongIndex offset_) {
+    ptr_ += offset_;
+    return *this;
+  }
+
+  /// Adds an offset to each pointer
+  MUTLASS_HOST_DEVICE
+  TensorRef & add_coord_offset(TensorCoord const &coord) {
+    add_pointer_offset(offset(coord));
+    return *this;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorRef operator+(TensorCoord const& b) const {
+    TensorRef result(*this);
+    result.add_coord_offset(b);
+    return result;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorRef & operator+=(TensorCoord const& b) {
+    add_coord_offset(b);
+    return *this;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorRef operator-(TensorCoord const& b) const {
+    TensorRef result(*this);
+    result.add_pointer_offset(-offset(b));
+    return result;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorRef & operator-=(TensorCoord const& b) {
+    add_pointer_offset(-offset(b));
+    return *this;
+  }
+};
+
+/// Constructs a TensorRef, deducing types from arguments.
+template <
+  typename Element,
+  typename Layout
+>
+MUTLASS_HOST_DEVICE
+TensorRef<Element, Layout> make_TensorRef(Element *ptr, Layout const &layout) {
+  return TensorRef<Element, Layout>(ptr, layout);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Partial specializations to handle degenerate and sub-byte cases.
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename Element,
+  typename Layout
+>
+MUTLASS_HOST_DEVICE
+bool TensorRef_aligned(TensorRef<Element, Layout> const &ref, int alignment) {
+
+  int const kStrideRank = Layout::kStrideRank;
+
+  if (reinterpret_cast<uintptr_t>(ref.data()) % alignment) {
+    return false;
+  }
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < kStrideRank; ++i) {
+    if (ref.stride(i) % alignment) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/tensor_ref_planar_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_ref_planar_complex.h
new file mode 100644
index 000000000..4eb5de7f8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_ref_planar_complex.h
@@ -0,0 +1,375 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a structure containing strides, bounds, and a pointer to tensor data.
+*/
+#pragma once
+
+#include <cstdint>
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+#include "mutlass/tensor_ref.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Element_>
+struct PlanarComplexReference {
+
+  //
+  // Type definitions
+  //
+
+  using Element = Element_;
+  using ComplexElement = complex<Element>;
+
+  //
+  // Data members
+  //
+
+  Element *real;
+  Element *imag;
+
+  //
+  // Methods
+  //
+
+  MUTLASS_HOST_DEVICE
+  PlanarComplexReference(
+    Element *real_ = nullptr,
+    Element *imag_ = nullptr
+  ):
+    real(real_), imag(imag_) { }
+
+  /// Loads the complex element
+  MUTLASS_HOST_DEVICE
+  operator complex<Element>() const {
+    return complex<Element>{*real, *imag};
+  }
+
+  /// Stores a complex element to the location pointed to by the reference
+  MUTLASS_HOST_DEVICE
+  PlanarComplexReference &operator=(complex<Element> const &rhs) {
+    *real = rhs.real();
+    *imag = rhs.imag();
+    return *this;
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/* \brief TensorRef is a template for objects pointing to the start of tensors of arbitrary rank
+          and layout within memory. A TensorRef combines a pointer and a Layout concept
+
+*/
+template <
+  /// Data type of element stored within tensor (concept: NumericType)
+  typename Element_,
+  /// Defines a mapping from logical coordinate to linear memory (concept: Layout)
+  typename Layout_
+>
+class TensorRefPlanarComplex {
+ public:
+  /// Data type of individual access
+  using Element = Element_;
+
+  /// Complex element type
+  using ComplexElement = complex<Element>;
+
+  /// Mapping function from logical coordinate to linear memory
+  using Layout = Layout_;
+
+  static_assert(sizeof_bits<Element>::value >= 8,
+    "Planar complex not suitable for subbyte elements at this time");
+
+  /// Reference type to an element
+  using Reference = PlanarComplexReference<Element>;
+
+  /// Logical rank of tensor index space
+  static int const kRank = Layout::kRank;
+
+  /// Index type
+  using Index = typename Layout::Index;
+
+  /// Long index used for pointer offsets
+  using LongIndex = typename Layout::LongIndex;
+
+  /// Coordinate in logical tensor space
+  using TensorCoord = typename Layout::TensorCoord;
+
+  /// Layout's stride vector
+  using Stride = typename Layout::Stride;
+
+  /// TensorRef to constant data
+  using ConstTensorRef = TensorRefPlanarComplex<
+    typename platform::remove_const<Element>::type const,
+    Layout>;
+
+  /// TensorRef to non-constant data
+  using NonConstTensorRef = TensorRefPlanarComplex<
+    typename platform::remove_const<Element>::type,
+    Layout>;
+
+  /// Require at least rank=1. Mathematically, a rank=0 tensor would be considered to be a
+  /// scalar, but degenerate cases such as these are difficult to accommodate without
+  /// extensive C++ metaprogramming or support for zero-length arrays.
+  static_assert(kRank > 0, "Cannot define a zero-rank TensorRef");
+
+ private:
+
+  /// Pointer
+  Element* ptr_;
+
+  /// Layout object maps logical coordinates to linear offsets
+  Layout layout_;
+
+  /// Offset to imaginary part
+  LongIndex imaginary_stride_;
+
+ public:
+
+  //
+  // Methods
+  //
+
+  /// Constructs a TensorRef with a pointer and layout object.
+  MUTLASS_HOST_DEVICE
+  TensorRefPlanarComplex(
+    Element *ptr = nullptr,                   ///< pointer to start of tensor
+    Layout const &layout = Layout(),          ///< layout object containing stride and mapping function
+    LongIndex imaginary_stride = 0
+  ):
+    ptr_(ptr), layout_(layout), imaginary_stride_(imaginary_stride) {
+
+  }
+
+  /// Converting constructor from TensorRef to non-constant data.
+  MUTLASS_HOST_DEVICE
+  TensorRefPlanarComplex(
+    NonConstTensorRef const &ref              ///< TensorRef to non-const data
+  ):
+    ptr_(ref.data()), layout_(ref.layout()), imaginary_stride_(ref.imaginary_stride_) { }
+
+  /// Returns a reference to constant-valued tensor.
+  MUTLASS_HOST_DEVICE
+  ConstTensorRef const_ref() const {
+    return ConstTensorRef(ptr_, layout_, imaginary_stride_);
+  }
+
+  MUTLASS_HOST_DEVICE
+  NonConstTensorRef non_const_ref() const {
+    return NonConstTensorRef(
+      const_cast<typename platform::remove_const<Element>::type *>(ptr_),
+      layout_,
+      imaginary_stride_);
+  }
+
+  /// Updates only the pointer
+  MUTLASS_HOST_DEVICE
+  void reset(Element* ptr = nullptr, LongIndex imaginary_stride = 0) {
+    ptr_ = ptr;
+    imaginary_stride_ = imaginary_stride;
+  }
+
+  /// Updates the pointer and layout object
+  MUTLASS_HOST_DEVICE
+  void reset(Element* ptr, Layout const &layout, LongIndex imaginary_stride) {
+    ptr_ = ptr;
+    layout_ = layout;
+    imaginary_stride_ = imaginary_stride;
+  }
+
+  /// Returns true if the TensorRef is non-null
+  MUTLASS_HOST_DEVICE
+  bool good() const {
+    return ptr_ != nullptr;
+  }
+
+  /// Returns the pointer to referenced data
+  MUTLASS_HOST_DEVICE
+  Element * data() const { return ptr_; }
+
+  /// Returns the pointer to referenced data
+  MUTLASS_HOST_DEVICE
+  Element * imaginary_data() const { return ptr_ + imaginary_stride_; }
+
+  /// Returns a reference to the element at a given linear index
+  MUTLASS_HOST_DEVICE
+  Reference data(LongIndex idx) const {
+    return Reference(ptr_ + idx, ptr_ + idx + imaginary_stride_);
+  }
+
+  /// Returns the layout object
+  MUTLASS_HOST_DEVICE
+  Layout & layout() {
+    return layout_;
+  }
+
+  /// Returns the layout object
+  MUTLASS_HOST_DEVICE
+  Layout layout() const {
+    return layout_;
+  }
+
+  /// Gets the stride to an imaginary element
+  LongIndex imaginary_stride() const {
+    return imaginary_stride_;
+  }
+
+  /// Gets the stride to an imaginary element
+  LongIndex &imaginary_stride() {
+    return imaginary_stride_;
+  }
+
+  /// Returns the layout object's stride vector
+  MUTLASS_HOST_DEVICE
+  Stride stride() const {
+    return layout_.stride();
+  }
+
+  /// Returns the layout object's stride vector
+  MUTLASS_HOST_DEVICE
+  Stride & stride() {
+    return layout_.stride();
+  }
+
+  /// Returns the layout object's stride in a given physical dimension
+  MUTLASS_HOST_DEVICE
+  Index stride(int dim) const {
+    return layout_.stride().at(dim);
+  }
+
+  /// Returns the layout object's stride in a given physical dimension
+  MUTLASS_HOST_DEVICE
+  Index & stride(int dim) {
+    return layout_.stride().at(dim);
+  }
+
+  /// Computes the offset of an index from the origin of the tensor
+  MUTLASS_HOST_DEVICE
+  LongIndex offset(TensorCoord const& coord) const {
+    return layout_(coord);
+  }
+
+  /// Returns a reference to the element at a given Coord
+  MUTLASS_HOST_DEVICE
+  Reference at(TensorCoord const& coord) const {
+    return data(offset(coord));
+  }
+
+  /// Returns a reference to the element at a given Coord
+  MUTLASS_HOST_DEVICE
+  Reference operator[](TensorCoord const& coord) const {
+    return data(offset(coord));
+  }
+
+  /// Adds an offset to each pointer
+  MUTLASS_HOST_DEVICE
+  TensorRefPlanarComplex & add_pointer_offset(LongIndex offset_) {
+    ptr_ += offset_;
+    return *this;
+  }
+
+  /// Adds an offset to each pointer
+  MUTLASS_HOST_DEVICE
+  TensorRefPlanarComplex & add_coord_offset(TensorCoord const &coord) {
+    add_pointer_offset(offset(coord));
+    return *this;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorRefPlanarComplex operator+(TensorCoord const& b) const {
+    TensorRefPlanarComplex result(*this);
+    result.add_coord_offset(b);
+    return result;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorRefPlanarComplex & operator+=(TensorCoord const& b) {
+    add_coord_offset(b);
+    return *this;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorRefPlanarComplex operator-(TensorCoord const& b) const {
+    TensorRefPlanarComplex result(*this);
+    result.add_pointer_offset(-offset(b));
+    return result;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorRefPlanarComplex & operator-=(TensorCoord const& b) {
+    add_pointer_offset(-offset(b));
+    return *this;
+  }
+
+  /// TensorRef to real-valued tensor
+  MUTLASS_HOST_DEVICE
+  mutlass::TensorRef<Element, Layout> ref_real() const {
+    return mutlass::TensorRef<Element, Layout>(data(), layout());
+  }
+
+  /// TensorRef to real-valued tensor
+  MUTLASS_HOST_DEVICE
+  mutlass::TensorRef<Element, Layout> ref_imag() const {
+    return mutlass::TensorRef<Element, Layout>(imaginary_data(), layout());
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Constructs a TensorRef, deducing types from arguments.
+template <
+  typename Element,
+  typename Layout
+>
+MUTLASS_HOST_DEVICE
+TensorRefPlanarComplex<Element, Layout> make_TensorRefPlanarComplex(
+  Element *ptr,
+  Layout const &layout,
+  int64_t imaginary_stride) {
+
+  return TensorRefPlanarComplex<Element, Layout>(ptr, layout, imaginary_stride);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/tensor_view.h b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_view.h
new file mode 100644
index 000000000..cbb260e09
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_view.h
@@ -0,0 +1,298 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a structure containing strides and a pointer to tensor data.
+
+    TensorView is derived from TensorRef and contributes bounds to the tensor's index space. Thus,
+    it is a complete mathematical object and may be used in tensor algorithms. It is decoupled from
+    data storage and is therefore lightweight and may be embedded in larger tensor objects or
+    memory structures.
+
+    See mutlass/tensor_ref.h for more details about the mapping of the logical tensor index space to
+    linear memory.
+*/
+
+#pragma once
+
+#if !defined(__MUSACC_RTC__)
+#include <cmath>
+#endif
+
+#include "mutlass/mutlass.h"
+#include "mutlass/tensor_ref.h"
+
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  /// Data type of element stored within tensor
+  typename Element_,
+  /// Maps a Coord<Rank_> in the logical tensor index space to the internal n-D array
+  typename Layout_
+>
+class TensorView : public TensorRef<Element_, Layout_> {
+ public:
+
+  /// Base tensor reference
+  using Base = mutlass::TensorRef<Element_, Layout_>;
+
+  /// Mapping function from logical coordinate to internal n-D array
+  using Layout = Layout_;
+
+  /// TensorRef pointing to constant memory
+  using ConstTensorRef = typename Base::ConstTensorRef;
+
+  /// Underlying TensorRef type
+  using TensorRef = Base;
+
+  /// Data type of individual access
+  using Element = Element_;
+
+  /// Reference type to an element
+  using Reference = Element &;
+
+  /// Logical rank of tensor index space
+  static int const kRank = Layout::kRank;
+
+  /// Index type
+  using Index = typename Layout::Index;
+
+  /// Long index used for pointer offsets
+  using LongIndex = typename Layout::LongIndex;
+
+  /// Coordinate in logical tensor space
+  using TensorCoord = typename Layout::TensorCoord;
+
+  /// Coordinate in storage n-D array
+  using Stride = typename Layout::Stride;
+
+  /// TensorView pointing to constant memory
+  using ConstTensorView = TensorView<
+    typename platform::remove_const<Element>::type const,
+    Layout>;
+
+  /// TensorView pointing to non-constant memory
+  using NonConstTensorView = TensorView<
+    typename platform::remove_const<Element>::type,
+    Layout>;
+
+  /// Require at least rank=1. Mathematically, a rank=0 tensor would be considered to be a
+  /// scalar, but degenerate cases such as these are difficult to accommodate without
+  /// extensive C++ metaprogramming or support for zero-length arrays.
+  static_assert(kRank > 0, "Cannot define a zero-rank TensorRef");
+
+ private:
+
+  /// View extent
+  TensorCoord extent_;
+
+ public:
+
+  //
+  // Methods
+  //
+
+  /// Constructs a TensorView object
+  MUTLASS_HOST_DEVICE
+  TensorView() { }
+
+  /// Constructs a TensorView object
+  MUTLASS_HOST_DEVICE
+  TensorView(
+    Element *ptr,                         ///< pointer to start of tensor
+    Layout const &layout,                 ///< layout object containing stride and mapping function
+    TensorCoord const &extent             ///< size of the view in logical coordinates
+  ):
+    Base(ptr, layout), extent_(extent) {
+
+  }
+
+  /// Constructs a TensorView object
+  MUTLASS_HOST_DEVICE
+  TensorView(
+    TensorRef const &ref,                 ///< pointer and layout object referencing a tensor
+    TensorCoord const &extent             ///< logical size of tensor
+  ):
+    Base(ref), extent_(extent) {
+
+  }
+
+  /// Converting constructor from TensorRef to non-constant data.
+  MUTLASS_HOST_DEVICE
+  TensorView(
+    NonConstTensorView const &view        ///< TensorView to non-const data
+  ):
+    Base(view), extent_(view.extent_) { }
+
+  /// Updates the pointer and layout object
+  MUTLASS_HOST_DEVICE
+  void reset(Element* ptr, Layout const &layout, TensorCoord const &extent) {
+    Base::reset(ptr, layout);
+    this->resize(extent);
+  }
+
+  /// Updates the pointer
+  MUTLASS_HOST_DEVICE
+  void reset(Element* ptr) {
+    Base::reset(ptr);
+  }
+
+  /// Changes the size of the view without affecting pointer or layout
+  MUTLASS_HOST_DEVICE
+  void resize(TensorCoord const &extent) {
+    this->extent_ = extent;
+  }
+
+  /// Returns the extent of the view (the size along each logical dimension).
+  MUTLASS_HOST_DEVICE
+  TensorCoord const& extent() const { return extent_; }
+
+  /// Returns the extent along a particular logical dimension.
+  MUTLASS_HOST_DEVICE
+  Index extent(int dim) const { return extent_.at(dim); }
+
+  /// Returns the number of logical elements
+  MUTLASS_HOST_DEVICE
+  LongIndex size() const {
+    return extent_.product();
+  }
+
+  /// Determines whether a location is within a tensor
+  MUTLASS_HOST_DEVICE
+  bool contains(TensorCoord const& coord) const {
+    MUTLASS_PRAGMA_UNROLL
+    for (int dim = 0; dim < kRank; ++dim) {
+      if (!(coord[dim] >= 0 && coord[dim] < extent(dim))) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /// Returns a TensorRef pointing to the first element of the tensor.
+  MUTLASS_HOST_DEVICE
+  TensorRef ref() const {
+    return TensorRef(this->data(), this->layout());
+  }
+
+  /// Returns a TensorRef pointing to the first element of the tensor.
+  MUTLASS_HOST_DEVICE
+  ConstTensorRef const_ref() const {
+    return ConstTensorRef(this->data(), this->layout());
+  }
+
+  /// Returns a TensorView to const data
+  MUTLASS_HOST_DEVICE
+  ConstTensorView const_view() const {
+    return ConstTensorView(const_ref(), extent_);
+  }
+
+  /// Returns a Tensor_view given location and size quantities
+  MUTLASS_HOST_DEVICE
+  TensorView subview(
+    TensorCoord extent,                               ///< extent of the resulting view
+    TensorCoord const& location = TensorCoord()       ///< resulting view's origin within the old view
+  ) const {
+
+    TensorView result(this->ref(), extent.clamp(extent_ - location));
+    result.add_coord_offset(location);
+    return result;
+  }
+
+  /// Returns the number of scalar elements needed to store tensor.
+  MUTLASS_HOST_DEVICE
+  size_t capacity() const {
+    return Base::layout().capacity(extent_);
+  }
+
+  /// Returns a TensorView offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorView operator+(
+    TensorCoord const& b            ///< offset in the logical coordinate space of the tensor
+  ) const {
+
+    TensorView result(*this);
+    result.add_pointer_offset(this->offset(b));
+    return result;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorView& operator+=(
+    TensorCoord const& b            ///< offset in the logical coordinate space of the tensor
+  ) {
+
+    this->add_pointer_offset(this->offset(b));
+    return *this;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorView operator-(
+    TensorCoord const& b            ///< offset in the logical coordinate space of the tensor
+  ) const {
+
+    TensorRef result(*this);
+    result.add_pointer_offset(-this->offset(b));
+    return result;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorView& operator-=(
+    TensorCoord const& b            ///< offset in the logical coordinate space of the tensor
+  ) {
+
+    this->add_pointer_offset(-this->offset(b));
+    return *this;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Constructs a TensorRef, deducing types from arguments.
+template <
+  typename Element,
+  typename Layout
+>
+MUTLASS_HOST_DEVICE TensorView<Element, Layout> make_TensorView(
+  Element *ptr,
+  Layout const &layout,
+  typename Layout::TensorCoord const &extent) {
+
+  return TensorView<Element, Layout>(ptr, layout, extent);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/tensor_view_planar_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_view_planar_complex.h
new file mode 100644
index 000000000..325ff5e36
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/tensor_view_planar_complex.h
@@ -0,0 +1,302 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a structure containing strides and a pointer to tensor data.
+
+    TensorView is derived from TensorRef and contributes bounds to the tensor's index space. Thus,
+    it is a complete mathematical object and may be used in tensor algorithms. It is decoupled from
+    data storage and is therefore lightweight and may be embedded in larger tensor objects or
+    memory structures.
+
+    See mutlass/tensor_ref.h for more details about the mapping of the logical tensor index space to
+    linear memory.
+*/
+
+#pragma once
+
+#if !defined(__MUSACC_RTC__)
+#include <cmath>
+#endif
+
+#include "mutlass/mutlass.h"
+#include "mutlass/tensor_ref_planar_complex.h"
+
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  /// Data type of element stored within tensor
+  typename Element_,
+  /// Maps a Coord<Rank_> in the logical tensor index space to the internal n-D array
+  typename Layout_
+>
+class TensorViewPlanarComplex : public TensorRefPlanarComplex<Element_, Layout_> {
+ public:
+
+  /// Base tensor reference
+  using Base = mutlass::TensorRefPlanarComplex<Element_, Layout_>;
+
+  /// Mapping function from logical coordinate to internal n-D array
+  using Layout = Layout_;
+
+  /// TensorRef pointing to constant memory
+  using ConstTensorRef = typename Base::ConstTensorRef;
+
+  /// Underlying TensorRef type
+  using TensorRef = Base;
+
+  /// Data type of individual access
+  using Element = Element_;
+
+  /// Reference type to an element
+  using Reference = Element &;
+
+  /// Logical rank of tensor index space
+  static int const kRank = Layout::kRank;
+
+  /// Index type
+  using Index = typename Layout::Index;
+
+  /// Long index used for pointer offsets
+  using LongIndex = typename Layout::LongIndex;
+
+  /// Coordinate in logical tensor space
+  using TensorCoord = typename Layout::TensorCoord;
+
+  /// Coordinate in storage n-D array
+  using Stride = typename Layout::Stride;
+
+  /// TensorView pointing to constant memory
+  using ConstTensorView = TensorViewPlanarComplex<
+    typename platform::remove_const<Element>::type const,
+    Layout>;
+
+  /// TensorView pointing to non-constant memory
+  using NonConstTensorView = TensorViewPlanarComplex<
+    typename platform::remove_const<Element>::type,
+    Layout>;
+
+  /// Require at least rank=1. Mathematically, a rank=0 tensor would be considered to be a
+  /// scalar, but degenerate cases such as these are difficult to accommodate without
+  /// extensive C++ metaprogramming or support for zero-length arrays.
+  static_assert(kRank > 0, "Cannot define a zero-rank TensorRef");
+
+ private:
+
+  /// View extent
+  TensorCoord extent_;
+
+ public:
+
+  //
+  // Methods
+  //
+
+  /// Constructs a TensorView object
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex(TensorCoord const &extent = TensorCoord()): extent_(extent) {
+
+  }
+
+  /// Constructs a TensorView object
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex(
+    Element *ptr,                         ///< pointer to start of tensor
+    Layout const &layout,                 ///< layout object containing stride and mapping function
+    LongIndex imaginary_stride,           ///< stride between real and imaginary part
+    TensorCoord const &extent             ///< size of the view in logical coordinates
+  ):
+    Base(ptr, layout, imaginary_stride), extent_(extent) {
+
+  }
+
+  /// Constructs a TensorView object
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex(
+    TensorRef const &ref,                 ///< pointer and layout object referencing a tensor
+    TensorCoord const &extent             ///< logical size of tensor
+  ):
+    Base(ref), extent_(extent) {
+
+  }
+
+  /// Converting constructor from TensorRef to non-constant data.
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex(
+    NonConstTensorView const &view        ///< TensorView to non-const data
+  ):
+    Base(view), extent_(view.extent_) { }
+
+  /// Updates the pointer and layout object
+  MUTLASS_HOST_DEVICE
+  void reset(Element* ptr, Layout const &layout, LongIndex imaginary_stride, TensorCoord size) {
+    Base::reset(ptr, layout, imaginary_stride);
+    this->resize(extent_);
+  }
+
+  /// Changes the size of the view without affecting pointer or layout
+  MUTLASS_HOST_DEVICE
+  void resize(TensorCoord extent) {
+    this->extent_ = extent;
+  }
+
+  /// Returns the extent of the view (the size along each logical dimension).
+  MUTLASS_HOST_DEVICE
+  TensorCoord const& extent() const { return extent_; }
+
+  /// Returns the extent along a particular logical dimension.
+  MUTLASS_HOST_DEVICE
+  Index extent(int dim) const { return extent_.at(dim); }
+
+  /// Determines whether a location is within a tensor
+  MUTLASS_HOST_DEVICE
+  bool contains(TensorCoord const& coord) const {
+    MUTLASS_PRAGMA_UNROLL
+    for (int dim = 0; dim < kRank; ++dim) {
+      if (!(coord[dim] >= 0 && coord[dim] < extent(dim))) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /// Returns a TensorRef pointing to the first element of the tensor.
+  MUTLASS_HOST_DEVICE
+  Base ref() const {
+    return Base(this->data(), this->layout(), this->imaginary_stride());
+  }
+
+  /// Returns a TensorRef pointing to the first element of the tensor.
+  MUTLASS_HOST_DEVICE
+  ConstTensorRef const_ref() const {
+    return ConstTensorRef(this->data(), this->layout());
+  }
+
+  /// Returns a TensorView to const data
+  MUTLASS_HOST_DEVICE
+  ConstTensorView const_view() const {
+    return ConstTensorView(const_ref(), extent_);
+  }
+
+  /// Returns a Tensor_view given location and size quantities
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex subview(
+    TensorCoord extent,                               ///< extent of the resulting view
+    TensorCoord const& location = TensorCoord()       ///< resulting view's origin within the old view
+  ) const {
+
+    TensorViewPlanarComplex result(this->ref(), extent.clamp(extent_ - location));
+    result.add_coord_offset(location);
+    return result;
+  }
+
+  /// Returns the number of scalar elements needed to store tensor.
+  MUTLASS_HOST_DEVICE
+  size_t capacity() const {
+    return Base::layout().capacity(extent_);
+  }
+
+  /// Returns a TensorView offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex operator+(
+    TensorCoord const& b            ///< offset in the logical coordinate space of the tensor
+  ) const {
+
+    TensorViewPlanarComplex result(*this);
+    result.add_pointer_offset(this->offset(b));
+    return result;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex& operator+=(
+    TensorCoord const& b            ///< offset in the logical coordinate space of the tensor
+  ) {
+
+    this->add_pointer_offset(this->offset(b));
+    return *this;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex operator-(
+    TensorCoord const& b            ///< offset in the logical coordinate space of the tensor
+  ) const {
+
+    TensorRef result(*this);
+    result.add_pointer_offset(-this->offset(b));
+    return result;
+  }
+
+  /// Returns a TensorRef offset by a given amount
+  MUTLASS_HOST_DEVICE
+  TensorViewPlanarComplex& operator-=(
+    TensorCoord const& b            ///< offset in the logical coordinate space of the tensor
+  ) {
+
+    this->add_pointer_offset(-this->offset(b));
+    return *this;
+  }
+
+  /// TensorRef to real-valued tensor
+  MUTLASS_HOST_DEVICE
+  mutlass::TensorView<Element, Layout> view_real() const {
+    return mutlass::TensorView<Element, Layout>(this->data(), this->layout(), extent_);
+  }
+
+  /// TensorRef to real-valued tensor
+  MUTLASS_HOST_DEVICE
+  mutlass::TensorView<Element, Layout> view_imag() const {
+    return mutlass::TensorView<Element, Layout>(this->imaginary_data(), this->layout(), extent_);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Constructs a TensorRef, deducing types from arguments.
+template <
+  typename Element,
+  typename Layout
+>
+MUTLASS_HOST_DEVICE TensorViewPlanarComplex<Element, Layout> make_TensorViewPlanarComplex(
+  Element *ptr,
+  Layout const &layout,
+  typename Layout::LongIndex imaginary_stride,
+  typename Layout::TensorCoord const &extent) {
+
+  return TensorViewPlanarComplex<Element, Layout>(ptr, layout, imaginary_stride, extent);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/tfloat32.h b/vllm_musa/csrc/flash_attn/include/mutlass/tfloat32.h
new file mode 100644
index 000000000..9fb527367
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/tfloat32.h
@@ -0,0 +1,477 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 a proxy class for storing Tensor Float 32 data type.
+*/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include "mutlass/floating_point_mtrtc.h"
+#else
+#include <cmath>
+#include <limits>
+#include <cstdint>
+#endif
+
+#include "mutlass/mutlass.h"
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Tensor Float 32 data type
+struct alignas(4) tfloat32_t {
+
+  //
+  // Data members
+  //
+
+  /// Storage type
+  uint32_t storage;
+
+  //
+  // Methods
+  //
+
+  /// Constructs from an unsigned int
+  MUTLASS_HOST_DEVICE
+  static tfloat32_t bitcast(uint32_t x) {
+    tfloat32_t h;
+    h.storage = x;
+    return h;
+  }
+
+  /// Emulated rounding is fast in device code
+  MUTLASS_HOST_DEVICE
+  static tfloat32_t round_half_ulp_truncate(float const &s) {
+    uint32_t x = reinterpret_cast<uint32_t const &>(s);
+
+    #if defined(__MUSA_ARCH__)
+    if (::isfinite(s)) {
+      x += 0x1000u;
+    }
+    #else
+    if (std::isfinite(s)) {
+      x += 0x1000u;
+    }
+    #endif
+
+    return tfloat32_t::bitcast(x);
+  }
+
+  /// Default constructor
+  tfloat32_t() = default;
+
+  /// Floating-point conversion - round toward nearest even
+  MUTLASS_HOST_DEVICE
+//  explicit tfloat32_t(float x): storage(round_half_ulp_truncate(x).storage) { }
+  tfloat32_t(float x): storage(round_half_ulp_truncate(x).storage) { }
+
+  /// Floating-point conversion - round toward nearest even
+  MUTLASS_HOST_DEVICE
+//  explicit tfloat32_t(double x): tfloat32_t(float(x)) {
+  tfloat32_t(double x): tfloat32_t(float(x)) {
+  }
+
+  /// Integer conversion - round toward zero
+  MUTLASS_HOST_DEVICE
+//  explicit tfloat32_t(int x) {
+  tfloat32_t(int x) {
+    float flt = static_cast<float>(x);
+    #if defined(__MUSA_ARCH__)
+    storage = reinterpret_cast<uint32_t const &>(flt);
+    #else
+    std::memcpy(&storage, &flt, sizeof(storage));
+    #endif
+  }
+
+  /// Converts to float
+  MUTLASS_HOST_DEVICE
+  operator float() const {
+
+    // Conversions to IEEE single-precision requires clearing dont-care bits
+    // of the mantissa.
+    unsigned bits = (storage & ~0x1fffu);
+
+    #if defined(__MUSA_ARCH__)
+    return reinterpret_cast<float const &>(bits);
+    #else
+    float flt;
+    std::memcpy(&flt, &bits, sizeof(flt));
+    return flt;
+    #endif
+  }
+
+  /// Converts to float
+  MUTLASS_HOST_DEVICE
+  explicit operator double() const {
+    return double(float(*this));
+  }
+
+  /// Converts to int
+  MUTLASS_HOST_DEVICE
+  explicit operator int() const {
+    return int(float(*this));
+  }
+
+  /// Casts to bool
+  MUTLASS_HOST_DEVICE
+  explicit operator bool() const {
+    return (float(*this) != 0.0f);
+  }
+
+  /// Obtains raw bits
+  MUTLASS_HOST_DEVICE
+  uint32_t raw() const {
+    return storage;
+  }
+
+  /// Returns the sign bit
+  MUTLASS_HOST_DEVICE
+  bool signbit() const {
+    return ((raw() & 0x80000000) != 0);
+  }
+
+  /// Returns the biased exponent
+  MUTLASS_HOST_DEVICE
+  int exponent_biased() const {
+    return int((raw() >> 23) & 0x0ff);
+  }
+
+  /// Returns the unbiased exponent
+  MUTLASS_HOST_DEVICE
+  int exponent() const {
+    return exponent_biased() - 127;
+  }
+
+  /// Returns the mantissa
+  MUTLASS_HOST_DEVICE
+  int mantissa() const {
+    return int(raw() & 0x7fffff);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+MUTLASS_HOST_DEVICE
+bool signbit(mutlass::tfloat32_t const& h) {
+  return h.signbit();
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::tfloat32_t abs(mutlass::tfloat32_t const& h) {
+  return mutlass::tfloat32_t::bitcast(h.raw() & 0x7fffffff);
+}
+
+MUTLASS_HOST_DEVICE
+bool isnan(mutlass::tfloat32_t const& h) {
+  return (h.exponent_biased() == 0x0ff) && h.mantissa();
+}
+
+MUTLASS_HOST_DEVICE
+bool isfinite(mutlass::tfloat32_t const& h) {
+  return (h.exponent_biased() != 0x0ff);
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::tfloat32_t nan_tf32(const char*) {
+  return mutlass::tfloat32_t::bitcast(0x7fffffff);
+}
+
+MUTLASS_HOST_DEVICE
+bool isinf(mutlass::tfloat32_t const& h) {
+  return (h.exponent_biased() == 0x0ff) && !h.mantissa();
+}
+
+MUTLASS_HOST_DEVICE
+bool isnormal(mutlass::tfloat32_t const& h) {
+  return h.exponent_biased() && h.exponent_biased() != 0x0ff;
+}
+
+MUTLASS_HOST_DEVICE
+int fpclassify(mutlass::tfloat32_t const& h) {
+  int exp = h.exponent_biased();
+  int mantissa = h.mantissa();
+  if (exp == 0x0ff) {
+    if (mantissa) {
+      return FP_NAN;
+    }
+    else {
+      return FP_INFINITE;
+    }
+  }
+  else if (!exp) {
+    if (mantissa) {
+      return FP_SUBNORMAL;
+    }
+    else {
+      return FP_ZERO;
+    }
+  }
+  return FP_NORMAL;
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::tfloat32_t sqrt(mutlass::tfloat32_t const& h) {
+#if defined(__MUSACC_RTC__)
+  return mutlass::tfloat32_t(sqrtf(float(h)));
+#else
+  return mutlass::tfloat32_t(std::sqrt(float(h)));
+#endif
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t copysign(tfloat32_t const& a, tfloat32_t const& b) {
+
+  uint32_t a_mag = (reinterpret_cast<uint32_t const &>(a) & 0x7fffffff);
+  uint32_t b_sign = (reinterpret_cast<uint32_t const &>(b) & 0x80000000);
+  uint32_t result = (a_mag | b_sign);
+
+  return reinterpret_cast<tfloat32_t const &>(result);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Standard Library operations and definitions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace std {
+
+#if !defined(__MUSACC_RTC__)
+/// Numeric limits
+template <>
+struct numeric_limits<mutlass::tfloat32_t> {
+  static bool const is_specialized = true;
+  static bool const is_signed = true;
+  static bool const is_integer = false;
+  static bool const is_exact = false;
+  static bool const has_infinity = true;
+  static bool const has_quiet_NaN = true;
+  static bool const has_signaling_NaN = false;
+  static std::float_denorm_style const has_denorm = std::denorm_present;
+  static bool const has_denorm_loss = true;
+  static std::float_round_style const round_style = std::round_to_nearest;
+  static bool const is_iec559 = false;
+  static bool const is_bounded = true;
+  static bool const is_modulo = false;
+  static int const digits = 19;
+
+  /// Least positive value
+  static mutlass::tfloat32_t min() { return mutlass::tfloat32_t::bitcast(0x01); }
+
+  /// Minimum finite value
+  static mutlass::tfloat32_t lowest() { return mutlass::tfloat32_t::bitcast(0xff7fffff); }
+
+  /// Maximum finite value
+  static mutlass::tfloat32_t max() { return mutlass::tfloat32_t::bitcast(0x7f7fffff); }
+
+  /// Returns smallest finite value
+  static mutlass::tfloat32_t epsilon() { return mutlass::tfloat32_t::bitcast(0x1000); }
+
+  /// Returns smallest finite value
+  static mutlass::tfloat32_t round_error() { return mutlass::tfloat32_t(0.5f); }
+
+  /// Returns smallest finite value
+  static mutlass::tfloat32_t infinity() { return mutlass::tfloat32_t::bitcast(0x7f800000); }
+
+  /// Returns smallest finite value
+  static mutlass::tfloat32_t quiet_NaN() { return mutlass::tfloat32_t::bitcast(0x7fffffff); }
+
+  /// Returns smallest finite value
+  static mutlass::tfloat32_t signaling_NaN() { return mutlass::tfloat32_t::bitcast(0x7fffffff); }
+
+  /// Returns smallest finite value
+  static mutlass::tfloat32_t denorm_min() { return mutlass::tfloat32_t::bitcast(0x1); }
+};
+#endif
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace std
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Arithmetic operators
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+MUTLASS_HOST_DEVICE
+bool operator==(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return float(lhs) == float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator!=(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return float(lhs) != float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return float(lhs) < float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator<=(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return float(lhs) <= float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return float(lhs) > float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+bool operator>=(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return float(lhs) >= float(rhs);
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t operator+(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return tfloat32_t(float(lhs) + float(rhs));
+}
+
+
+MUTLASS_HOST_DEVICE
+tfloat32_t operator-(tfloat32_t const& lhs) {
+  union u_tff32 {
+    float val_f32;
+    tfloat32_t val_tf;
+    MUTLASS_HOST_DEVICE u_tff32() : val_f32(0) { }
+  };
+  union u_tff32 x; x.val_f32 = -reinterpret_cast<float const &>(lhs);
+  return x.val_tf;
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t operator-(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return tfloat32_t(float(lhs) - float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t operator*(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return tfloat32_t(float(lhs) * float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t operator/(tfloat32_t const& lhs, tfloat32_t const& rhs) {
+  return tfloat32_t(float(lhs) / float(rhs));
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t& operator+=(tfloat32_t & lhs, tfloat32_t const& rhs) {
+  lhs = tfloat32_t(float(lhs) + float(rhs));
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t& operator-=(tfloat32_t & lhs, tfloat32_t const& rhs) {
+  lhs = tfloat32_t(float(lhs) - float(rhs));
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t& operator*=(tfloat32_t & lhs, tfloat32_t const& rhs) {
+  lhs = tfloat32_t(float(lhs) * float(rhs));
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t& operator/=(tfloat32_t & lhs, tfloat32_t const& rhs) {
+  lhs = tfloat32_t(float(lhs) / float(rhs));
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t& operator++(tfloat32_t & lhs) {
+  float tmp(lhs);
+  ++tmp;
+  lhs = tfloat32_t(tmp);
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t& operator--(tfloat32_t & lhs) {
+  float tmp(lhs);
+  --tmp;
+  lhs = tfloat32_t(tmp);
+  return lhs;
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t operator++(tfloat32_t & lhs, int) {
+  tfloat32_t ret(lhs);
+  float tmp(lhs);
+  tmp++;
+  lhs = tfloat32_t(tmp);
+  return ret;
+}
+
+MUTLASS_HOST_DEVICE
+tfloat32_t operator--(tfloat32_t & lhs, int) {
+  tfloat32_t ret(lhs);
+  float tmp(lhs);
+  tmp--;
+  lhs = tfloat32_t(tmp);
+  return ret;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+//
+// User-defined literals
+//
+
+MUTLASS_HOST_DEVICE
+mutlass::tfloat32_t operator "" _tf32(long double x) {
+  return mutlass::tfloat32_t(float(x));
+}
+
+MUTLASS_HOST_DEVICE
+mutlass::tfloat32_t operator "" _tf32(unsigned long long int x) {
+  return mutlass::tfloat32_t(int(x));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/trace.h b/vllm_musa/csrc/flash_attn/include/mutlass/trace.h
new file mode 100644
index 000000000..0e72720d5
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/trace.h
@@ -0,0 +1,59 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Helpers for optionally tracing through code when debugging.
+
+    This file is to be included after all other headers.
+*/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Tracing options
+#ifndef MUTLASS_DEBUG_TRACE_LEVEL
+#define MUTLASS_DEBUG_TRACE_LEVEL 0
+#endif
+
+#if MUTLASS_DEBUG_TRACE_LEVEL
+#include <iostream>
+#include "mutlass/core_io.h"
+#if defined(__MUSA_ARCH__)
+#define MUTLASS_TRACE_HOST(x)
+#else
+#define MUTLASS_TRACE_HOST(x) { std::cout << __FILE__ << ":" << __LINE__ << "  " << x << std::endl; }
+#endif
+#else
+#define MUTLASS_TRACE_HOST(x)
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/transform/thread/transpose.h b/vllm_musa/csrc/flash_attn/include/mutlass/transform/thread/transpose.h
new file mode 100644
index 000000000..ee54bc503
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/transform/thread/transpose.h
@@ -0,0 +1,108 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Basic copy routines for tensor views
+*/
+
+#pragma once
+
+namespace mutlass {
+namespace transform {
+namespace thread {
+
+/// Transforms a fragment by doing a transpose
+template <
+  int ElementCount,
+  typename TransposeShape,
+  typename Element
+> struct Transpose;
+
+/// Specialization for int8_t 4x4 transpose
+template <int ElementCount_>
+struct Transpose<ElementCount_, layout::PitchLinearShape<4,4> , int8_t> {
+
+    static const int kElementCount = ElementCount_;
+    using TransposeShape = layout::PitchLinearShape<4,4>;
+    using Element = int8_t;
+    using Fragment = mutlass::Array<Element, kElementCount>;
+
+    static_assert(!(kElementCount % TransposeShape::kCount), "Shape needs to be multiple of 16 elements to do a 4x4 transpose");
+
+    MUTLASS_DEVICE
+    void transform(Fragment& dst, Fragment& src) {
+
+    // Expose src/dst as int arrays.
+    int* src_int = reinterpret_cast<int*>(&src);
+    int* dst_int = reinterpret_cast<int*>(&dst);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kElementCount / TransposeShape::kCount; i++){
+
+      int const i0 = 4 * i + 0;
+      int const i1 = 4 * i + 1;
+      int const i2 = 4 * i + 2;
+      int const i3 = 4 * i + 3;
+
+      int a0 = src_int[i0];
+      int a1 = src_int[i1];
+      int a2 = src_int[i2];
+      int a3 = src_int[i3];
+
+      int b0, b1, b2, b3, c0;
+      b0 = __byte_perm(a0, a1, 0x0040);
+      c0 = __byte_perm(a2, a3, 0x0040);
+      b0 = __byte_perm(b0, c0, 0x5410);
+
+      b1 = __byte_perm(a0, a1, 0x0051);
+      c0 = __byte_perm(a2, a3, 0x0051);
+      b1 = __byte_perm(b1, c0, 0x5410);
+
+      b2 = __byte_perm(a0, a1, 0x0062);
+      c0 = __byte_perm(a2, a3, 0x0062);
+      b2 = __byte_perm(b2, c0, 0x5410);
+
+      b3 = __byte_perm(a0, a1, 0x0073);
+      c0 = __byte_perm(a2, a3, 0x0073);
+      b3 = __byte_perm(b3, c0, 0x5410);
+
+      dst_int[i0] = b0;
+      dst_int[i1] = b1;
+      dst_int[i2] = b2;
+      dst_int[i3] = b3;
+    }
+  }
+};
+
+}  // namespace thread
+}  // namespace layout
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/transform/thread/unary_op.h b/vllm_musa/csrc/flash_attn/include/mutlass/transform/thread/unary_op.h
new file mode 100644
index 000000000..10cb2a575
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/transform/thread/unary_op.h
@@ -0,0 +1,106 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+
+namespace mutlass {
+namespace transform {
+namespace thread {
+
+namespace UnaryTransform {
+    struct Identity;    ///< None (i.e., identity)
+    struct Conjugate;   ///< Complex conjugate
+}
+
+/// Element-wise unary operator that transforms one element of a fragment at a time
+template<
+    typename FragmentIn, ///< Input Fragment
+    typename FragmentOut,///< Output Fragment
+    typename Transform>  ///< Unary transform operator
+class UnaryOp
+{
+    public:
+        MUTLASS_DEVICE
+        static FragmentOut execute(FragmentIn &in)
+        {
+            static_assert(FragmentIn::kElements == FragmentOut::kElements, "Number of elements must match.");
+            static_assert(platform::is_same<Transform, UnaryTransform::Identity>::value ||
+                          platform::is_same<Transform, UnaryTransform::Conjugate>::value,
+                          "Unary Operator not supported.");
+
+            FragmentOut out;
+            if (platform::is_same<Transform, UnaryTransform::Identity>::value )
+            {
+                MUTLASS_PRAGMA_UNROLL
+                for (int i=0; i < FragmentIn::kElements; ++i){
+                   out[i] = static_cast<typename FragmentOut::Element>(in[i]);
+                }
+            }
+            else if (platform::is_same<Transform, UnaryTransform::Conjugate>::value )
+            {
+                for (int i=0; i < FragmentIn::kElements; ++i){
+                   out[i] = conj(static_cast<typename FragmentOut::Element>(in[i]));
+                }
+            }
+            return out;
+        }
+};
+
+template<typename FragmentIn, typename Transform>
+class UnaryOp<FragmentIn, FragmentIn, Transform>
+{
+    public:
+        MUTLASS_DEVICE
+        static FragmentIn execute(FragmentIn &in)
+        {
+            static_assert(platform::is_same<Transform, UnaryTransform::Identity>::value ||
+                          platform::is_same<Transform, UnaryTransform::Conjugate>::value,
+                          "Unary Operator not supported.");
+
+            if (platform::is_same<Transform, UnaryTransform::Identity>::value )
+            {
+                return in;
+            }
+            else if (platform::is_same<Transform, UnaryTransform::Conjugate>::value )
+            {
+                for(int i=0; i < FragmentIn::kElements; ++i){
+                   in[i] = conj(in[i]);
+                }
+            }
+            return in;
+        }
+      };
+    }
+  }
+}
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/uint128.h b/vllm_musa/csrc/flash_attn/include/mutlass/uint128.h
new file mode 100644
index 000000000..5fa669a44
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/uint128.h
@@ -0,0 +1,268 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 an unsigned 128b integer with several operators to support 64-bit integer division.
+*/
+#pragma once
+
+#if defined(__MUSACC_RTC__)
+#include <musa/std/cstdint>
+#else
+#include <cstdint>
+#include <cstdlib>
+#include <cmath>
+#include <type_traits>
+#include <stdexcept>
+#endif
+
+#include "mutlass/mutlass.h"
+
+/// Optionally enable GCC's built-in type
+#if (defined(__x86_64) || defined (__aarch64__)) && defined(__GNUC__)
+#define MUTLASS_UINT128_NATIVE
+#elif defined(_MSC_VER) && defined(_M_AMD64)
+#define MUTLASS_INT128_ARITHMETIC
+#include <intrin.h>
+#if _MSC_VER >= 1920 && !defined(__MUSA_ARCH__)
+#define MUTLASS_INT128_ARITHMETIC_DIV
+#include <immintrin.h>
+#endif
+#endif
+
+namespace mutlass {
+
+///! Unsigned 128b integer type
+struct alignas(16) uint128_t
+{
+  /// Size of one part of the uint's storage in bits
+  static constexpr int storage_bits_ = 64;
+
+  struct hilo
+  {
+    uint64_t lo;
+    uint64_t hi;
+  };
+
+  // Use a union to store either low and high parts or, if present, a built-in 128b integer type.
+  union {
+    struct hilo hilo_;
+
+#if defined(MUTLASS_UINT128_NATIVE)
+    unsigned __int128 native;
+#endif // defined(MUTLASS_UINT128_NATIVE)
+  };
+
+  //
+  // Methods
+  //
+
+  /// Default ctor
+  MUTLASS_HOST_DEVICE
+  uint128_t() : hilo_{0, 0} {}
+
+  /// Constructor from uint64
+  MUTLASS_HOST_DEVICE
+  uint128_t(uint64_t lo_) : hilo_{lo_, 0} {}
+
+  /// Constructor from two 64b unsigned integers
+  MUTLASS_HOST_DEVICE
+  uint128_t(uint64_t lo_, uint64_t hi_) : hilo_{lo_, hi_} {}
+
+  /// Optional constructor from native value
+#if defined(MUTLASS_UINT128_NATIVE)
+  uint128_t(unsigned __int128 value) : native(value) { }
+#endif
+
+  /// Lossily cast to uint64
+  MUTLASS_HOST_DEVICE
+  explicit operator uint64_t() const
+  {
+    return hilo_.lo;
+  }
+
+  MUTLASS_HOST_DEVICE
+  static void exception()
+  {
+#if defined(__MUSA_ARCH__)
+    __musa_exit();
+#else
+    // throw std::runtime_error("Not yet implemented.");
+    abort();
+#endif
+  }
+
+  /// Add
+  MUTLASS_HOST_DEVICE
+  uint128_t operator+(uint128_t const& rhs) const
+  {
+    uint128_t y{};
+#if defined(MUTLASS_UINT128_NATIVE)
+    y.native = native + rhs.native;
+#else
+    y.hilo_.lo = hilo_.lo + rhs.hilo_.lo;
+    y.hilo_.hi = hilo_.hi + rhs.hilo_.hi + (y.hilo_.lo < hilo_.lo);
+#endif
+    return y;
+  }
+
+  /// Subtract
+  MUTLASS_HOST_DEVICE
+  uint128_t operator-(uint128_t const& rhs) const
+  {
+    uint128_t y{};
+#if defined(MUTLASS_UINT128_NATIVE)
+    y.native = native - rhs.native;
+#else
+    y.hilo_.lo = hilo_.lo - rhs.hilo_.lo;
+    y.hilo_.hi = hilo_.hi - rhs.hilo_.hi - (rhs.hilo_.lo && y.hilo_.lo > hilo_.lo);
+#endif
+    return y;
+  }
+
+  /// Multiply by unsigned 64b integer yielding 128b integer
+  MUTLASS_HOST_DEVICE
+  uint128_t operator*(uint64_t const& rhs) const
+  {
+    uint128_t y{};
+#if defined(MUTLASS_UINT128_NATIVE)
+    y.native = native * rhs;
+#elif defined(MUTLASS_INT128_ARITHMETIC)
+    // Multiply by the low part
+    y.hilo_.lo = _umul128(hilo_.lo, rhs, &y.hilo_.hi);
+
+    // Add the high part and ignore the overflow
+    uint64_t overflow{0};
+    y.hilo_.hi += _umul128(hilo_.hi, rhs, &overflow);
+#else
+    MUTLASS_UNUSED(rhs);
+    exception();
+#endif
+    return y;
+  }
+
+  /// Divide 128b operation by 64b operation yielding a 64b quotient
+  MUTLASS_HOST_DEVICE
+  uint64_t operator/(uint64_t const& divisor) const
+  {
+    uint64_t quotient{0};
+#if defined(MUTLASS_UINT128_NATIVE)
+    quotient = uint64_t(native / divisor);
+#elif defined(MUTLASS_INT128_ARITHMETIC_DIV)
+    // implemented using MSVC's arithmetic intrinsics
+    uint64_t remainder{0};
+    quotient = _udiv128(hilo_.hi, hilo_.lo, divisor, &remainder);
+#else
+    MUTLASS_UNUSED(divisor);
+    exception();
+#endif
+    return quotient;
+  }
+
+  /// Divide 128b operation by 64b operation yielding a 64b quotient
+  MUTLASS_HOST_DEVICE
+  uint64_t operator%(uint64_t const& divisor) const
+  {
+    uint64_t remainder{0};
+#if defined(MUTLASS_UINT128_NATIVE)
+    remainder = uint64_t(native % divisor);
+#elif defined(MUTLASS_INT128_ARITHMETIC_DIV)
+    // implemented using MSVC's arithmetic intrinsics
+    (void)_udiv128(hilo_.hi, hilo_.lo, divisor, &remainder);
+#else
+    MUTLASS_UNUSED(divisor);
+    exception();
+#endif
+    return remainder;
+  }
+
+  /// Computes the quotient and remainder in a single method.
+  MUTLASS_HOST_DEVICE
+  uint64_t divmod(uint64_t &remainder, uint64_t divisor) const
+  {
+    uint64_t quotient{0};
+#if defined(MUTLASS_UINT128_NATIVE)
+    quotient = uint64_t(native / divisor);
+    remainder = uint64_t(native % divisor);
+#elif defined(MUTLASS_INT128_ARITHMETIC_DIV)
+    // implemented using MSVC's arithmetic intrinsics
+    quotient = _udiv128(hilo_.hi, hilo_.lo, divisor, &remainder);
+#else
+    MUTLASS_UNUSED(remainder);
+    MUTLASS_UNUSED(divisor);
+    exception();
+#endif
+    return quotient;
+  }
+
+  /// Left-shifts a 128b unsigned integer
+  MUTLASS_HOST_DEVICE
+  uint128_t operator<<(int sh) const
+  {
+    if (sh == 0) {
+      return *this;
+    }
+    else if (sh >= storage_bits_) {
+      return uint128_t(0, hilo_.lo << (sh - storage_bits_));
+    }
+    else {
+      return uint128_t(
+        (hilo_.lo << sh),
+        (hilo_.hi << sh) | uint64_t(hilo_.lo >> (storage_bits_ - sh))
+      );
+    }
+  }
+
+  /// Right-shifts a 128b unsigned integer
+  MUTLASS_HOST_DEVICE
+  uint128_t operator>>(int sh) const
+  {
+    if (sh == 0) {
+      return *this;
+    }
+    else if (sh >= storage_bits_) {
+      return uint128_t((hilo_.hi >> (sh - storage_bits_)), 0);
+    }
+    else {
+      return uint128_t(
+        (hilo_.lo >> sh) | (hilo_.hi << (storage_bits_ - sh)),
+        (hilo_.hi >> sh)
+      );
+    }
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/GPU_Clock.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/GPU_Clock.hpp
new file mode 100644
index 000000000..4d52b9520
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/GPU_Clock.hpp
@@ -0,0 +1,68 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <musa_runtime.h>
+
+struct GPU_Clock
+{
+  GPU_Clock() {
+    musaEventCreate(&start_);
+    musaEventCreate(&stop_);
+    musaEventRecord(start_);
+  }
+
+  ~GPU_Clock() {
+    musaEventDestroy(start_);
+    musaEventDestroy(stop_);
+  }
+
+  void start() {
+    musaEventRecord(start_);
+  }
+
+  float milliseconds() {
+    musaEventRecord(stop_);
+    musaEventSynchronize(stop_);
+    float time;
+    musaEventElapsedTime(&time, start_, stop_);
+    return time;
+  }
+
+  float seconds() {
+    return milliseconds() * float(1e-3);
+  }
+
+ private:
+  musaEvent_t start_, stop_;
+};
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/command_line.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/command_line.h
new file mode 100644
index 000000000..1b6813f67
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/command_line.h
@@ -0,0 +1,314 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * Utility for parsing command line arguments
+ */
+
+#include <iostream>
+#include <limits>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include <musa_runtime.h>
+
+#include "mutlass/mutlass.h"
+
+namespace mutlass {
+
+/******************************************************************************
+ * command_line
+ ******************************************************************************/
+
+/**
+ * Utility for parsing command line arguments
+ */
+struct CommandLine {
+  std::vector<std::string> keys;
+  std::vector<std::string> values;
+  std::vector<std::string> args;
+
+  /**
+   * Constructor
+   */
+  CommandLine(int argc, const char** argv) {
+    using namespace std;
+
+    for (int i = 1; i < argc; i++) {
+      string arg = argv[i];
+
+      if ((arg[0] != '-') || (arg[1] != '-')) {
+        args.push_back(arg);
+        continue;
+      }
+
+      string::size_type pos;
+      string key, val;
+      if ((pos = arg.find('=')) == string::npos) {
+        key = string(arg, 2, arg.length() - 2);
+        val = "";
+      } else {
+        key = string(arg, 2, pos - 2);
+        val = string(arg, pos + 1, arg.length() - 1);
+      }
+
+      keys.push_back(key);
+      values.push_back(val);
+    }
+  }
+
+  /**
+   * Checks whether a flag "--<flag>" is present in the commandline
+   */
+  bool check_cmd_line_flag(const char* arg_name) const {
+    using namespace std;
+
+    for (int i = 0; i < int(keys.size()); ++i) {
+      if (keys[i] == string(arg_name)) return true;
+    }
+    return false;
+  }
+
+  /**
+   * Returns number of naked (non-flag and non-key-value) commandline parameters
+   */
+  size_t num_naked_args() const {
+    return args.size();
+  }
+
+  /**
+   * Print naked (non-flag and non-key-value) commandline parameters
+   */
+  void print_naked_args(std::ostream &out) const {
+    for (auto arg : args) {
+      out << "   " << arg <<"\n";
+    }
+  }
+
+  /**
+   * Returns the commandline parameter for a given index (not including flags)
+   */
+  template <typename value_t>
+  void get_cmd_line_argument(int index, value_t& val) const {
+    using namespace std;
+    if (index < args.size()) {
+      istringstream str_stream(args[index]);
+      str_stream >> val;
+    }
+  }
+
+  /**
+   * Obtains the boolean value specified for a given commandline parameter --<flag>=<bool>
+   */
+  void get_cmd_line_argument(const char* arg_name, bool& val, bool _default) const {
+    val = _default;
+    if (check_cmd_line_flag(arg_name)) {
+      std::string value;
+      get_cmd_line_argument(arg_name, value);
+
+      val = !(value == "0" || value == "false");
+    }
+  }
+
+  /**
+   * Obtains the value specified for a given commandline parameter --<flag>=<value>
+   */
+  template <typename value_t>
+  void get_cmd_line_argument(const char* arg_name,
+                             value_t& val) const {
+
+    get_cmd_line_argument(arg_name, val, val);
+  }
+
+  /**
+   * Obtains the value specified for a given commandline parameter --<flag>=<value>
+   */
+  template <typename value_t>
+  void get_cmd_line_argument(const char* arg_name,
+                             value_t& val,
+                             value_t const& _default) const {
+    using namespace std;
+
+    val = _default;
+
+    for (int i = 0; i < int(keys.size()); ++i) {
+      if (keys[i] == string(arg_name)) {
+        istringstream str_stream(values[i]);
+        str_stream >> val;
+      }
+    }
+  }
+
+  /**
+   * Returns the values specified for a given commandline parameter --<flag>=<value>,<value>*
+   */
+  template <typename value_t>
+  void get_cmd_line_arguments(const char* arg_name,
+                              std::vector<value_t>& vals,
+                              char sep = ',') const {
+    using namespace std;
+
+    if (check_cmd_line_flag(arg_name)) {
+      // Clear any default values
+      vals.clear();
+
+      // Recover from multi-value string
+      for (int i = 0; i < keys.size(); ++i) {
+        if (keys[i] == string(arg_name)) {
+          string val_string(values[i]);
+          separate_string(val_string, vals, sep);
+        }
+      }
+    }
+  }
+
+  /**
+   * Returns the values specified for a given commandline parameter
+   * --<flag>=<value>,<value_start:value_end>*
+   */
+  void get_cmd_line_argument_pairs(const char* arg_name,
+                                   std::vector<std::pair<std::string, std::string> >& tokens,
+                                   char delim = ',',
+                                   char sep = ':') const {
+    if (check_cmd_line_flag(arg_name)) {
+      std::string value;
+      get_cmd_line_argument(arg_name, value);
+
+      tokenize(tokens, value, delim, sep);
+    }
+  }
+
+  /**
+   * Returns a list of ranges specified for a given commandline parameter
+   * --<flag>=<key:value>,<key:value>*
+   */
+  void get_cmd_line_argument_ranges(const char* arg_name,
+                                    std::vector<std::vector<std::string> >& vals,
+                                    char delim = ',',
+                                    char sep = ':') const {
+    std::vector<std::string> ranges;
+    get_cmd_line_arguments(arg_name, ranges, delim);
+
+    for (std::vector<std::string>::const_iterator range = ranges.begin();
+      range != ranges.end(); ++range) {
+
+      std::vector<std::string> range_vals;
+      separate_string(*range, range_vals, sep);
+      vals.push_back(range_vals);
+    }
+  }
+
+  /**
+   * The number of pairs parsed
+   */
+  int parsed_argc() const { return (int)keys.size(); }
+
+  //-------------------------------------------------------------------------
+  // Utility functions
+  //-------------------------------------------------------------------------
+
+  /// Tokenizes a comma-delimited list of string pairs delimited by ':'
+  static void tokenize(std::vector<std::pair<std::string, std::string> >& tokens,
+                       std::string const& str,
+                       char delim = ',',
+                       char sep = ':') {
+    // Home-built to avoid Boost dependency
+    size_t s_idx = 0;
+    size_t d_idx = std::string::npos;
+    while (s_idx < str.size()) {
+      d_idx = str.find_first_of(delim, s_idx);
+
+      size_t end_idx = (d_idx != std::string::npos ? d_idx : str.size());
+      size_t sep_idx = str.find_first_of(sep, s_idx);
+      size_t offset = 1;
+      if (sep_idx == std::string::npos || sep_idx >= end_idx) {
+        sep_idx = end_idx;
+        offset = 0;
+      }
+
+      std::pair<std::string, std::string> item(
+          str.substr(s_idx, sep_idx - s_idx),
+          str.substr(sep_idx + offset, end_idx - sep_idx - offset));
+
+      tokens.push_back(item);
+      s_idx = end_idx + 1;
+    }
+  }
+
+  /// Tokenizes a comma-delimited list of string pairs delimited by ':'
+  static void tokenize(std::vector<std::string>& tokens,
+                       std::string const& str,
+                       char delim = ',',
+                       char sep = ':') {
+    typedef std::vector<std::pair<std::string, std::string> > TokenVector;
+    typedef TokenVector::const_iterator token_iterator;
+
+    std::vector<std::pair<std::string, std::string> > token_pairs;
+    tokenize(token_pairs, str, delim, sep);
+    for (token_iterator tok = token_pairs.begin(); tok != token_pairs.end(); ++tok) {
+      tokens.push_back(tok->first);
+    }
+  }
+
+  template <typename value_t>
+  static void separate_string(std::string const& str,
+                              std::vector<value_t>& vals,
+                              char sep = ',') {
+    std::istringstream str_stream(str);
+    std::string::size_type old_pos = 0;
+    std::string::size_type new_pos = 0;
+
+    // Iterate <sep>-delimited values
+    value_t val;
+    while ((new_pos = str.find(sep, old_pos)) != std::string::npos) {
+      if (new_pos != old_pos) {
+        str_stream.width(new_pos - old_pos);
+        str_stream >> val;
+        vals.push_back(val);
+      }
+
+      // skip over delimiter
+      str_stream.ignore(1);
+      old_pos = new_pos + 1;
+    }
+
+    // Read last value
+    str_stream >> val;
+    vals.push_back(val);
+  }
+};
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/debug.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/debug.h
new file mode 100644
index 000000000..cc9325711
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/debug.h
@@ -0,0 +1,144 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Contains code for debugging mutlass code
+*/
+
+#pragma once
+
+#include "device_dump.h"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/******************************************************************************
+ * Debug and logging macros
+ ******************************************************************************/
+
+/**
+ * Formats and prints the given message to stdout
+ */
+#if !defined(MUSA_LOG)
+#if !defined(__MUSA_ARCH__)
+#define MUSA_LOG(format, ...) printf(format, __VA_ARGS__)
+#else
+#define MUSA_LOG(format, ...)                              \
+  printf("[block (%d,%d,%d), thread (%d,%d,%d)]: " format, \
+         blockIdx.x,                                       \
+         blockIdx.y,                                       \
+         blockIdx.z,                                       \
+         threadIdx.x,                                      \
+         threadIdx.y,                                      \
+         threadIdx.z,                                      \
+         __VA_ARGS__);
+#endif
+#endif
+
+/**
+ * Formats and prints the given message to stdout only if DEBUG is defined
+ */
+#if !defined(MUSA_LOG_DEBUG)
+#ifdef DEBUG
+#define MUSA_LOG_DEBUG(format, ...) MUSA_LOG(format, __VA_ARGS__)
+#else
+#define MUSA_LOG_DEBUG(format, ...)
+#endif
+#endif
+
+/**
+ * \brief The corresponding error message is printed to \p stderr (or \p stdout in device code)
+ * along with the supplied source context.
+ *
+ * \return The MUSA error.
+ */
+__host__ MUTLASS_DEVICE musaError_t musa_perror_impl(musaError_t error,
+                                                     const char* expression,
+                                                     const char* filename,
+                                                     int line) {
+  (void)filename;
+  (void)line;
+  if (error) {
+#if !defined(__MUSA_ARCH__)
+    fprintf(
+        stderr, "MUSA error %d [%s, %d] in expression '%s': %s\n", error, filename, line, expression, musaGetErrorString(error));
+    fflush(stderr);
+#else
+    printf("MUSA error %d [%s, %d] in expression '%s'\n", error, filename, line, expression);
+#endif
+  }
+  return error;
+}
+
+/**
+ * \brief Perror macro
+ */
+#ifndef MUSA_PERROR
+#define MUSA_PERROR(e) musa_perror_impl((musaError_t)(e), #e, __FILE__, __LINE__)
+#endif
+
+/**
+ * \brief Perror macro with exit
+ */
+#ifndef MUSA_PERROR_EXIT
+#define MUSA_PERROR_EXIT(e)                                     \
+  do { if (musa_perror_impl((musaError_t)(e), #e, __FILE__, __LINE__)) { \
+    exit(1);                                                    \
+  } } while (0)
+#endif
+
+/**
+ * \brief Perror macro only if DEBUG is defined
+ */
+#ifndef MUSA_PERROR_DEBUG
+#ifdef DEBUG
+#define MUSA_PERROR_DEBUG(e) MUSA_PERROR(e)
+#else
+#define MUSA_PERROR_DEBUG(e) (e)
+#endif
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// A small helper class to dump a type at compile time
+// Usage:: DumpType<Class>::Class
+template <typename T>
+struct DebugType {};
+
+template <typename T>
+void DebugTypeFunc(T const& t) {
+  T::t;
+}
+
+// A small helper class to dump a compile time constant at compile time
+// Usage: DumpValue<Class::kConstant>::kConstant
+template <int Value>
+struct DebugValue {};
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_dump.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_dump.h
new file mode 100644
index 000000000..60c6bd928
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_dump.h
@@ -0,0 +1,188 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <stdio.h>
+#include "mutlass/mutlass.h"
+
+/**
+ * \file
+ * \brief C++ interface to dump fragments and shared memory contents for
+ * debugging.
+ */
+
+namespace mutlass {
+namespace debug {
+
+/******************************************************************************
+ * Dump the fragments
+ ******************************************************************************/
+
+/// The first N threads dump the first M elements from their fragments with a
+/// stride of S elements.  If N is not specified, dump the data of all the
+/// threads.  If M is not specified, dump all the elements of the fragment.
+template <typename Fragment>
+MUTLASS_DEVICE void dump_fragment(Fragment const& frag, int N = 0, int M = 0,
+                                  int S = 1) {
+  int total_threads = blockDim.x * blockDim.y * blockDim.z;
+  int block_id =
+      blockIdx.x + blockIdx.y * gridDim.x + gridDim.x * gridDim.y * blockIdx.z;
+  int thread_id = (threadIdx.z * (blockDim.x * blockDim.y)) +
+                  (threadIdx.y * blockDim.x) + threadIdx.x;
+
+  if (N < 0 || N > total_threads) {
+    if (thread_id == 0 && block_id == 0)
+      printf("Thread number N = %d should between [1, %d].\n", N,
+             total_threads);
+
+    __syncthreads();
+
+    return;
+  }
+
+  int total_elements = frag.size();
+
+  if (M < 0 || M > total_elements) {
+    if (thread_id == 0 && block_id == 0)
+      printf("Element number M = %d should between [1, %d].\n", M,
+             total_elements);
+
+    __syncthreads();
+
+    return;
+  }
+
+  if (N == 0) N = total_threads;
+
+  if (M == 0) M = total_elements;
+
+  if (S < 1 || S > M) {
+    if (thread_id == 0 && block_id == 0)
+      printf("Stride S = %d should between [1, %d].\n", S, M);
+
+    __syncthreads();
+
+    return;
+  }
+
+  if (thread_id == 0 && block_id == 0)
+    printf("\n*******************Dumping the fragments*******************\n\n");
+
+  MUTLASS_PRAGMA_NO_UNROLL
+  for (int tid = 0; tid < N; ++tid) {
+    if (tid == thread_id) {
+      printf("TB%d W%d T%d: ", block_id, tid / 32, tid & 31);
+      MUTLASS_PRAGMA_NO_UNROLL
+      for (int i = 0; i < M; i += S) {
+        printf("%.0f ", float(typename Fragment::value_type(frag[i])));
+      }
+      printf("\n");
+    }
+
+    __syncthreads();
+  }
+
+  if (thread_id == 0 && block_id == 0)
+    printf("\n***********************************************************\n\n");
+
+  __syncthreads();
+
+  return;
+}
+
+/******************************************************************************
+ * Dump the shared memory
+ ******************************************************************************/
+
+#define SHMEM_ROW_SIZE 128
+
+/// Dump the shared memory contents.  ptr is the begin address, size specifies
+/// the number of elements that need to be dumped, and S specifies the stride.
+template <typename Element>
+MUTLASS_DEVICE void dump_shmem(Element const* ptr, size_t size, int S = 1) {
+  int block_id =
+      blockIdx.x + blockIdx.y * gridDim.x + gridDim.x * gridDim.y * blockIdx.z;
+  int thread_id = (threadIdx.z * (blockDim.x * blockDim.y)) +
+                  (threadIdx.y * blockDim.x) + threadIdx.x;
+
+  if (ptr == nullptr) {
+    if (thread_id == 0 && block_id == 0) printf("ptr is null.\n");
+
+    __syncthreads();
+    return;
+  }
+
+  if (size < 1) {
+    if (thread_id == 0 && block_id == 0)
+      printf("Element size is less than 1\n");
+
+    __syncthreads();
+
+    return;
+  }
+
+  int row_elements = SHMEM_ROW_SIZE / sizeof(Element);
+
+  if (S < 1 || S > row_elements) {
+    if (thread_id == 0 && block_id == 0)
+      printf("Stride S = %d should between [1, %d].\n", S, row_elements);
+
+    __syncthreads();
+
+    return;
+  }
+
+  __syncthreads();
+
+  if (thread_id == 0)
+    printf("\n********Dumping the shared memory of TB %d*******\n\n", block_id);
+
+  if (thread_id == 0) {
+    for (int i = 0; i < size; i += row_elements) {
+      for (int j = 0; j < row_elements; j += S) {
+        printf("%.0f ", float(ptr[i + j]));
+      }
+
+      printf("\n");
+    }
+  }
+
+  if (thread_id == 0)
+    printf("\n***********************************************************\n\n");
+
+  __syncthreads();
+
+  return;
+}
+}  // namespace debug
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_groupnorm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_groupnorm.h
new file mode 100644
index 000000000..b1a1b81c4
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_groupnorm.h
@@ -0,0 +1,403 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief musa kernels to do group norm on a device memory tensor with NHWC layout. The tensor will be divided into [N, H, W, G, C'] and then we do normalization on [H, W, C'].
+ */
+
+#include "mutlass/mutlass.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/tensor_ref.h"
+#include "device_utils.h"
+#include <float.h>
+
+namespace mutlass {
+
+/** \brief interface to do group norm on a device memory tensor with NHWC layout.
+ * \tparam T: data type
+ */
+template <typename T>
+void groupnorm(mutlass::Tensor4DCoord input_size,
+               const int num_groups,
+               const float eps,
+               TensorRef<T, layout::TensorNHWC> ref_output,
+               TensorRef<T, layout::TensorNHWC> ref_input,
+               TensorRef<T, layout::TensorNHWC> ref_gamma,
+               TensorRef<T, layout::TensorNHWC> ref_beta,
+               musaStream_t stream);
+
+extern __shared__ char groupnorm_shm[];
+
+// For small prod_dim1_to_last_dim/num_groups, to avoid multiple loads from global memory,
+// we store the input in the shared memory.
+// grid(num_groups, dim0)
+// block(BLOCKSIZE)
+// BLOCKSIZE * TVecs_PER_THREAD <= prod_dim1_to_last_dim/num_group
+template<typename TVec, typename T, int T_PER_TVec>
+__global__ void groupnorm_twopass_store_locally(T*          output,
+                                                const T*    input,
+                                                const T*    gamma,
+                                                const T*    beta,
+                                                int         num_groups,
+                                                int         prod_dim1_to_last_dim,
+                                                int         last_dim,
+                                                const float eps,
+                                                const int   TVecs_PER_THREAD)
+{
+    const int   bid               = blockIdx.y;   // index of batch
+    const int   gid               = blockIdx.x;   // index of group
+    const int   tid               = threadIdx.x;  // index of thread
+    const int   bdimx             = blockDim.x;
+    const int   s_reduce_elements = prod_dim1_to_last_dim / num_groups;
+    const int   v_reduce_elements = s_reduce_elements / T_PER_TVec;
+    const int   s_group_stride    = last_dim / num_groups;
+    const int   v_group_stride    = s_group_stride / T_PER_TVec;
+    const int   offset_of_group   = (bid * prod_dim1_to_last_dim + gid * s_group_stride) / T_PER_TVec;
+    const TVec* input_TVec_ptr    = (const TVec*)(input) + offset_of_group;
+    TVec*       output_TVec_ptr   = (TVec*)(output) + offset_of_group;
+    T*       local_val         = ((T*)groupnorm_shm) + TVecs_PER_THREAD * T_PER_TVec * tid;
+    float       local_sum[1]      = {0.0f};
+
+// load from global memory into shared memory
+#pragma unroll
+    for (int i = 0; i < TVecs_PER_THREAD; i += 1) {
+        const int current_load_start_idx = (i * bdimx + tid) * T_PER_TVec;
+        const int offset_in_group =
+            ((current_load_start_idx / s_group_stride) * last_dim + (current_load_start_idx % s_group_stride))
+            / T_PER_TVec;
+        if (current_load_start_idx < s_reduce_elements) {
+            TVec      tmp_vec          = input_TVec_ptr[offset_in_group];
+            T*        tmp_vec_ptr      = (T*)(&tmp_vec);
+            const int local_val_offset = i * T_PER_TVec;
+#pragma unroll
+            for (int j = 0; j < T_PER_TVec; j++) {
+                float tmp = static_cast<float>(tmp_vec_ptr[j]);
+                local_sum[0] += tmp;
+                local_val[local_val_offset + j] = tmp_vec_ptr[j];
+            }
+        }
+    }
+    __shared__ float s_mean, s_variance;
+
+    // reduction for mean
+    if (bdimx <= 32) {
+        warpReduceSum<float, 1>(local_sum);
+    }
+    else {
+        blockReduceSum<float, 1>(local_sum);
+    }
+    if (tid == 0) {
+        s_mean = local_sum[0] / s_reduce_elements;
+    }
+    __syncthreads();
+
+    // reduction for std
+    local_sum[0] = 0.0f;
+#pragma unroll
+    for (int i = 0; i < TVecs_PER_THREAD; i += 1) {
+        const int current_load_start_idx = (i * bdimx + tid) * T_PER_TVec;
+        if (current_load_start_idx < s_reduce_elements) {
+            const int local_val_offset = i * T_PER_TVec;
+#pragma unroll
+            for (int j = 0; j < T_PER_TVec; j++) {
+                float tmp = static_cast<float>(local_val[local_val_offset + j]);
+                tmp -= s_mean;
+                local_sum[0] += tmp * tmp;
+            }
+        }
+    }
+    if (bdimx <= 32) {
+        warpReduceSum<float, 1>(local_sum);
+    }
+    else {
+        blockReduceSum<float, 1>(local_sum);
+    }
+    if (tid == 0) {
+        s_variance = rsqrtf(local_sum[0] / s_reduce_elements + eps);
+    }
+    __syncthreads();
+
+    // normalize
+    const int   gamma_offset_of_group = gid * v_group_stride;
+    const TVec* gamma_TVec_ptr        = (const TVec*)gamma + gamma_offset_of_group;
+    const TVec* beta_TVec_ptr         = (const TVec*)beta + gamma_offset_of_group;
+#pragma unroll
+    for (int i = 0; i < TVecs_PER_THREAD; i += 1) {
+        const int current_load_start_idx = (i * bdimx + tid) * T_PER_TVec;
+        const int offset_in_group =
+            ((current_load_start_idx / s_group_stride) * last_dim + (current_load_start_idx % s_group_stride))
+            / T_PER_TVec;
+        const int gamma_offset_in_group = (current_load_start_idx % s_group_stride) / T_PER_TVec;
+        const int local_val_offset      = i * T_PER_TVec;
+        if (current_load_start_idx < s_reduce_elements) {
+            TVec gamma_val     = gamma_TVec_ptr[gamma_offset_in_group];
+            TVec beta_val      = beta_TVec_ptr[gamma_offset_in_group];
+            T*   gamma_val_ptr = (T*)(&gamma_val);
+            T*   beta_val_ptr  = (T*)(&beta_val);
+            TVec tmp_vec;
+            T*   tmp_vec_ptr = (T*)(&tmp_vec);
+#pragma unroll
+            for (int j = 0; j < T_PER_TVec; j++) {
+                float tmp = (static_cast<float>(local_val[local_val_offset + j]) - s_mean) * s_variance
+                                * static_cast<float>(gamma_val_ptr[j])
+                            + static_cast<float>(beta_val_ptr[j]);
+                if (sizeof(T) == sizeof(half)) {
+                    tmp_vec_ptr[j] = T(__float2half_rn(tmp));
+                }
+                else {
+                    tmp_vec_ptr[j] = T(tmp);
+                }
+            }
+            output_TVec_ptr[offset_in_group] = tmp_vec;
+        }
+    }
+}
+
+// For large prod_dim1_to_last_dim/num_groups,
+// in which the data cannot be stored locally,
+// we will load from global memory multiple times,
+// grid(num_groups, dim0)
+// block(BLOCKSIZE)
+// BLOCKSIZE * TVecs_PER_THREAD <= prod_dim1_to_last_dim/num_group
+template<typename TVec, typename T, int T_PER_TVec>
+__global__ void groupnorm_twopass_multiple_load(T*          output,
+                                                const T*    input,
+                                                const T*    gamma,
+                                                const T*    beta,
+                                                int         num_groups,
+                                                int         prod_dim1_to_last_dim,
+                                                int         last_dim,
+                                                const float eps,
+                                                const int   TVecs_PER_THREAD)
+{
+    const int   bid               = blockIdx.y;   // index of batch
+    const int   gid               = blockIdx.x;   // index of group
+    const int   tid               = threadIdx.x;  // index of thread
+    const int   bdimx             = blockDim.x;
+    const int   s_reduce_elements = prod_dim1_to_last_dim / num_groups;
+    const int   v_reduce_elements = s_reduce_elements / T_PER_TVec;
+    const int   s_group_stride    = last_dim / num_groups;
+    const int   v_group_stride    = s_group_stride / T_PER_TVec;
+    const int   offset_of_group   = (bid * prod_dim1_to_last_dim + gid * s_group_stride) / T_PER_TVec;
+    const TVec* input_TVec_ptr    = (const TVec*)(input) + offset_of_group;
+    TVec*       output_TVec_ptr   = (TVec*)(output) + offset_of_group;
+    float       local_sum[1]      = {0.0f};
+
+#pragma unroll
+    for (int i = 0; i < TVecs_PER_THREAD; i += 1) {
+        const int current_load_start_idx = (i * bdimx + tid) * T_PER_TVec;
+        if (current_load_start_idx < s_reduce_elements) {
+            const int offset_in_group =
+                ((current_load_start_idx / s_group_stride) * last_dim + (current_load_start_idx % s_group_stride))
+                / T_PER_TVec;
+            TVec tmp_vec     = input_TVec_ptr[offset_in_group];
+            T*   tmp_vec_ptr = (T*)(&tmp_vec);
+#pragma unroll
+            for (int j = 0; j < T_PER_TVec; j++) {
+                float tmp = static_cast<float>(tmp_vec_ptr[j]);
+                local_sum[0] += tmp;
+            }
+        }
+    }
+    __shared__ float s_mean, s_variance;
+
+    // reduction for mean
+    if (bdimx <= 32) {
+        warpReduceSum<float, 1>(local_sum);
+    }
+    else {
+        blockReduceSum<float, 1>(local_sum);
+    }
+    if (tid == 0) {
+        s_mean = local_sum[0] / s_reduce_elements;
+    }
+    __syncthreads();
+
+    // reduction for std
+    local_sum[0] = 0.0f;
+#pragma unroll
+    for (int i = 0; i < TVecs_PER_THREAD; i += 1) {
+        const int current_load_start_idx = (i * bdimx + tid) * T_PER_TVec;
+        if (current_load_start_idx < s_reduce_elements) {
+            const int offset_in_group =
+                ((current_load_start_idx / s_group_stride) * last_dim + (current_load_start_idx % s_group_stride))
+                / T_PER_TVec;
+            TVec tmp_vec     = input_TVec_ptr[offset_in_group];
+            T*   tmp_vec_ptr = (T*)(&tmp_vec);
+#pragma unroll
+            for (int j = 0; j < T_PER_TVec; j++) {
+                float tmp = static_cast<float>(tmp_vec_ptr[j]);
+                tmp -= s_mean;
+                local_sum[0] += tmp * tmp;
+            }
+        }
+    }
+    if (bdimx <= 32) {
+        warpReduceSum<float, 1>(local_sum);
+    }
+    else {
+        blockReduceSum<float, 1>(local_sum);
+    }
+    if (tid == 0) {
+        s_variance = rsqrtf(local_sum[0] / s_reduce_elements + eps);
+    }
+    __syncthreads();
+
+    // normalize
+    const int   gamma_offset_of_group = gid * v_group_stride;
+    const TVec* gamma_TVec_ptr        = (const TVec*)gamma + gamma_offset_of_group;
+    const TVec* beta_TVec_ptr         = (const TVec*)beta + gamma_offset_of_group;
+#pragma unroll
+    for (int i = 0; i < TVecs_PER_THREAD; i += 1) {
+        const int current_load_start_idx = (i * bdimx + tid) * T_PER_TVec;
+        if (current_load_start_idx < s_reduce_elements) {
+            const int offset_in_group =
+                ((current_load_start_idx / s_group_stride) * last_dim + (current_load_start_idx % s_group_stride))
+                / T_PER_TVec;
+            const int gamma_offset_in_group = (current_load_start_idx % s_group_stride) / T_PER_TVec;
+            TVec      gamma_val             = gamma_TVec_ptr[gamma_offset_in_group];
+            TVec      beta_val              = beta_TVec_ptr[gamma_offset_in_group];
+            T*        gamma_val_ptr         = (T*)(&gamma_val);
+            T*        beta_val_ptr          = (T*)(&beta_val);
+            TVec      tmp_vec               = input_TVec_ptr[offset_in_group];
+            T*        tmp_vec_ptr           = (T*)(&tmp_vec);
+            TVec      output_tmp_vec;
+            T*        output_tmp_vec_ptr = (T*)(&output_tmp_vec);
+#pragma unroll
+            for (int j = 0; j < T_PER_TVec; j++) {
+                float tmp =
+                    (static_cast<float>(tmp_vec_ptr[j]) - s_mean) * s_variance * static_cast<float>(gamma_val_ptr[j])
+                    + static_cast<float>(beta_val_ptr[j]);
+                if (sizeof(T) == sizeof(half)) {
+                    output_tmp_vec_ptr[j] = T(__float2half_rn(tmp));
+                }
+                else {
+                    output_tmp_vec_ptr[j] = T(tmp);
+                }
+            }
+            output_TVec_ptr[offset_in_group] = output_tmp_vec;
+        }
+    }
+}
+
+//ref_input & ref_output should be [N, H, W, C]
+//ref_gamma & ref_beta should be [1, 1, 1, C]
+template <typename T>
+void groupnorm(mutlass::Tensor4DCoord input_size,
+               const int num_groups,
+               const float eps,
+               TensorRef<T, layout::TensorNHWC> ref_output,
+               TensorRef<T, layout::TensorNHWC> ref_input,
+               TensorRef<T, layout::TensorNHWC> ref_gamma,
+               TensorRef<T, layout::TensorNHWC> ref_beta,
+               musaStream_t stream){
+  const int N = input_size.n();
+  const int H = input_size.h();
+  const int W = input_size.w();
+  const int C = input_size.c();
+  if (C % num_groups != 0){
+    printf("[ERROR] C should be a multiple of num_groups.\n");
+  }
+  T* output = ref_output.data();
+  const T* input = ref_input.data();
+  const T* gamma = ref_gamma.data();
+  const T* beta = ref_beta.data();
+
+  const int dim0 = N;
+  const int last_dim = C;
+  const int prod_dim1_to_last_dim = H*W*C;
+  const int s_reduce_elements = prod_dim1_to_last_dim / num_groups;
+  const int s_group_stride = last_dim / num_groups;
+  dim3      grid(num_groups, dim0);
+  int       threadblock_size = 32;
+  if (s_group_stride % 2 == 0) {
+    const int T_PER_TVec = 2;
+    while (threadblock_size < 1024) {
+      if (s_reduce_elements / T_PER_TVec / threadblock_size <= 8)
+        break;
+        threadblock_size *= 2;
+      }
+    dim3      block(threadblock_size);
+    const int TVec_PER_THREAD = (s_reduce_elements / T_PER_TVec + threadblock_size - 1) / threadblock_size;
+    const int shm_size = T_PER_TVec * TVec_PER_THREAD * threadblock_size * sizeof(T);
+    // for small s_reduce_elements, specific case for H=W=22, C=1280, num_groups=32;
+    // the size of grid & block may have better choice for different cases.
+    // ensure shared memory is smaller than 48KB
+    if (std::is_same<T, float>::value){
+      if (shm_size < 48 * 1024) {
+        groupnorm_twopass_store_locally<float2, T, T_PER_TVec><<<grid, block, shm_size, stream>>>(
+          output, input, gamma, beta, num_groups, prod_dim1_to_last_dim, last_dim, eps, TVec_PER_THREAD);
+      }
+      else {
+        groupnorm_twopass_multiple_load<float2, T, T_PER_TVec><<<grid, block, 0, stream>>>(
+          output, input, gamma, beta, num_groups, prod_dim1_to_last_dim, last_dim, eps, TVec_PER_THREAD);
+      }
+    }
+    else{
+      if (shm_size < 48 * 1024) {
+        groupnorm_twopass_store_locally<half2, T, T_PER_TVec><<<grid, block, shm_size, stream>>>(
+          output, input, gamma, beta, num_groups, prod_dim1_to_last_dim, last_dim, eps, TVec_PER_THREAD);
+      }
+      else {
+        groupnorm_twopass_multiple_load<half2, T, T_PER_TVec><<<grid, block, 0, stream>>>(
+          output, input, gamma, beta, num_groups, prod_dim1_to_last_dim, last_dim, eps, TVec_PER_THREAD);
+      }
+    }
+  }
+  else {
+    const int T_PER_TVec = 1;
+    while (threadblock_size < 1024) {
+      if (s_reduce_elements / T_PER_TVec / threadblock_size <= 8)
+        break;
+        threadblock_size *= 2;
+      }
+    dim3      block(threadblock_size);
+    const int TVec_PER_THREAD = (s_reduce_elements / T_PER_TVec + threadblock_size - 1) / threadblock_size;
+    const int shm_size = T_PER_TVec * TVec_PER_THREAD * threadblock_size * sizeof(T);
+    if (shm_size < 48 * 1024) {
+      groupnorm_twopass_store_locally<T, T, T_PER_TVec><<<grid, block, shm_size, stream>>>(
+        output, input, gamma, beta, num_groups, prod_dim1_to_last_dim, last_dim, eps, TVec_PER_THREAD);
+    }
+    else {
+      groupnorm_twopass_multiple_load<T, T, T_PER_TVec><<<grid, block, 0, stream>>>(
+        output, input, gamma, beta, num_groups, prod_dim1_to_last_dim, last_dim, eps, TVec_PER_THREAD);
+    }
+  }
+
+}
+
+} //namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_layernorm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_layernorm.h
new file mode 100644
index 000000000..ce76537f7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_layernorm.h
@@ -0,0 +1,645 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief musa kernels to do layernorm on a device memory tensor with RowMajor layout.
+ */
+
+#include "mutlass/mutlass.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/tensor_ref.h"
+#include "device_utils.h"
+#include <float.h>
+
+namespace mutlass {
+
+/** \brief interface to do layernorm on a device memory tensor with RowMajor layout.
+ * \tparam T: data type
+ */
+template <typename T>
+void layernorm(mutlass::MatrixCoord tensor_size,
+               TensorRef<T, layout::RowMajor> ref_output,
+               TensorRef<T, layout::RowMajor> ref_input,
+               TensorRef<T, layout::RowMajor> ref_gamma,
+               TensorRef<T, layout::RowMajor> ref_beta,
+               musaStream_t stream);
+
+/**
+ * output [m, n] row-major
+ * input [m, n] row-major
+ * gamma [n]
+ * beta [n]
+ * grid(m)
+ * block(block_size) -- each block deals with n elements ; each thread deals with ITEM_PER_THREAD elements
+*/
+template<typename T, int ITEM_PER_THREAD>
+__global__ void layernorm_twoPassAlgo_stored_locally_e1(T* output,
+                                                        const T* input,
+                                                        const T* gamma,
+                                                        const T* beta,
+                                                        const int m,
+                                                        const int n)
+{
+  const int m_idx = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int bdimx = blockDim.x;
+  __shared__ float s_mean, s_variance;
+  T local_val[ITEM_PER_THREAD];
+  float local_sums[1] = {0.0f};
+  int offset = m_idx * n;
+  input += offset;
+  output += offset;
+
+  const T zero = T(0.0f);
+  #pragma unroll
+  for (int i = 0 ; i < ITEM_PER_THREAD ; i++){
+    int index = tid + i*bdimx;
+    local_val[i] = index < n ? input[index] : zero;
+    local_sums[0] += static_cast<float>(local_val[i]);
+  }
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_mean = local_sums[0] / n;
+  }
+  __syncthreads();
+
+  local_sums[0] = 0.0f;
+  #pragma unroll
+  for (int i = 0 ; i < ITEM_PER_THREAD ; i++){
+    int index = tid + i*bdimx;
+    if (index < n){
+      const float tmp = static_cast<float>(local_val[i]) - s_mean;
+      local_sums[0] += tmp * tmp;
+    }
+  }
+
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(local_sums[0] / n + 1e-5);
+  }
+  __syncthreads();
+
+  #pragma unroll
+  for (int i = 0 ; i < ITEM_PER_THREAD ; i++){
+    int index = tid + i*bdimx;
+    if (index < n) {
+      const T gamma_val = gamma[index];
+      const T beta_val = beta[index];
+      output[index] = T((static_cast<float>(local_val[i]) - s_mean) * s_variance * static_cast<float>(gamma_val) + static_cast<float>(beta_val));
+    }
+  }
+}
+
+/**
+ * output [m, n] row-major
+ * input [m, n] row-major
+ * gamma [n]
+ * beta [n]
+ * grid(m)
+ * block(block_size) -- each block deals with block_size*ITEM_PER_THREAD*2 elements;
+*/
+template<typename T2, typename T, int ITEM_PER_THREAD>
+__global__ void layernorm_twoPassAlgo_stored_locally_e2(T2* output,
+                                                        const T2* input,
+                                                        const T2* gamma,
+                                                        const T2* beta,
+                                                        const int m,
+                                                        const int n)
+{
+  const int m_idx = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int bdimx = blockDim.x;
+  __shared__ float s_mean, s_variance;
+  float local_sums[1] = {0.0f};
+  T2 local_val[ITEM_PER_THREAD];
+  const int n_2 = n / 2;
+  int offset = m_idx * n_2;
+  input += offset;
+  output += offset;
+
+  const T2 zero = {T(0.0f), T(0.0f)};
+  #pragma UNROLL
+  for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
+    const int index = i*bdimx + tid;
+    local_val[i] = index < n_2 ? input[index] : zero;
+    local_sums[0] += static_cast<float>(local_val[i].x) + static_cast<float>(local_val[i].y);
+  }
+
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_mean = local_sums[0] / n;
+  }
+  __syncthreads();
+
+  local_sums[0] = 0.0f;
+  #pragma UNROLL
+  for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
+    const int index = i*bdimx + tid;
+    if (index < n_2){
+      const float2 tmp = {static_cast<float>(local_val[i].x) - s_mean,
+                          static_cast<float>(local_val[i].y) - s_mean};
+      local_sums[0] += tmp.x * tmp.x + tmp.y * tmp.y;
+    }
+  }
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(local_sums[0] / n + 1e-5);
+  }
+  __syncthreads();
+
+  #pragma UNROLL
+  for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
+    const int index = i*bdimx + tid;
+    if (index < n_2){
+      const T2 gamma_val = gamma[index];
+      const T2 beta_val = beta[index];
+      T2 tmp;
+      tmp.x = T((static_cast<float>(local_val[i].x) - s_mean)*s_variance*static_cast<float>(gamma_val.x) + static_cast<float>(beta_val.x));
+      tmp.y = T((static_cast<float>(local_val[i].y) - s_mean)*s_variance*static_cast<float>(gamma_val.y) + static_cast<float>(beta_val.y));
+      output[index] = tmp;
+    }
+  }
+}
+
+/**
+ * output [m, n] row-major
+ * input [m, n] row-major
+ * gamma [n]
+ * beta [n]
+ * grid(m)
+ * block(block_size) -- each block deals with block_size*ITEM_PER_THREAD*4 elements;
+*/
+template<typename T4, typename T, int ITEM_PER_THREAD>
+__global__ void layernorm_twoPassAlgo_stored_locally_e4(T4* output,
+                                                        const T4* input,
+                                                        const T4* gamma,
+                                                        const T4* beta,
+                                                        const int m,
+                                                        const int n)
+{
+  const int m_idx = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int bdimx = blockDim.x;
+  __shared__ float s_mean, s_variance;
+  float local_sums[1] = {0.0f};
+  T4 local_val[ITEM_PER_THREAD];
+  const int n_4 = n / 4;
+  int offset = m_idx * n_4;
+  input += offset;
+  output += offset;
+
+  const T4 zero = {T(0.0f), T(0.0f), T(0.0f), T(0.0f)};
+  #pragma UNROLL
+  for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
+    const int index = i*bdimx + tid;
+    local_val[i] = index < n_4 ? input[index] : zero;
+    local_sums[0] += static_cast<float>(local_val[i].x) + static_cast<float>(local_val[i].y) +
+                     static_cast<float>(local_val[i].z) + static_cast<float>(local_val[i].w);
+  }
+
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_mean = local_sums[0] / n;
+  }
+  __syncthreads();
+
+  local_sums[0] = 0.0f;
+  #pragma UNROLL
+  for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
+    const int index = i*bdimx + tid;
+    if (index < n_4){
+      const float4 tmp = {static_cast<float>(local_val[i].x) - s_mean,
+                          static_cast<float>(local_val[i].y) - s_mean,
+                          static_cast<float>(local_val[i].z) - s_mean,
+                          static_cast<float>(local_val[i].w) - s_mean};
+      local_sums[0] += tmp.x * tmp.x + tmp.y * tmp.y + tmp.z * tmp.z + tmp.w * tmp.w;
+    }
+  }
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(local_sums[0] / n + 1e-5);
+  }
+  __syncthreads();
+
+  #pragma UNROLL
+  for (int i = 0; i < ITEM_PER_THREAD; i += 1) {
+    const int index = i*bdimx + tid;
+    if (index < n_4){
+      const T4 gamma_val = gamma[index];
+      const T4 beta_val = beta[index];
+      T4 tmp;
+      tmp.x = T((static_cast<float>(local_val[i].x) - s_mean)*s_variance*static_cast<float>(gamma_val.x) + static_cast<float>(beta_val.x));
+      tmp.y = T((static_cast<float>(local_val[i].y) - s_mean)*s_variance*static_cast<float>(gamma_val.y) + static_cast<float>(beta_val.y));
+      tmp.z = T((static_cast<float>(local_val[i].z) - s_mean)*s_variance*static_cast<float>(gamma_val.z) + static_cast<float>(beta_val.z));
+      tmp.w = T((static_cast<float>(local_val[i].w) - s_mean)*s_variance*static_cast<float>(gamma_val.w) + static_cast<float>(beta_val.w));
+      output[index] = tmp;
+    }
+  }
+}
+
+/**
+ * output [m, n] row-major
+ * input [m, n] row-major
+ * gamma [n]
+ * beta [n]
+ * grid(m)
+ * block(block_size) -- each block deals with n elements ; each thread deals with ITEM_PER_THREAD elements
+*/
+template<typename T>
+__global__ void layernorm_twoPassAlgo_e1(T* output,
+                                         const T* input,
+                                         const T* gamma,
+                                         const T* beta,
+                                         const int m,
+                                         const int n)
+{
+  const int m_idx = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int bdimx = blockDim.x;
+  __shared__ float s_mean, s_variance;
+  float local_sums[1] = {0.0f};
+  int offset = m_idx * n;
+  input += offset;
+  output += offset;
+
+  for (int index = tid ; index < n ; index += bdimx){
+    float local_val = static_cast<float>(input[index]);
+    local_sums[0] += local_val;
+  }
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_mean = local_sums[0] / n;
+  }
+  __syncthreads();
+
+  local_sums[0] = 0.0f;
+  for (int index = tid ; index < n ; index += bdimx){
+    float local_val = static_cast<float>(input[index]);
+    local_val = local_val - s_mean;
+    local_sums[0] += local_val * local_val;
+  }
+
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(local_sums[0] / n + 1e-5);
+  }
+  __syncthreads();
+
+  for (int index = tid ; index < n ; index += bdimx){
+    const T gamma_val = gamma[index];
+    const T beta_val = beta[index];
+    const T local_val = input[index];
+    output[index] = T((static_cast<float>(local_val) - s_mean) * s_variance * static_cast<float>(gamma_val) + static_cast<float>(beta_val));
+  }
+}
+
+/**
+ * output [m, n] row-major
+ * input [m, n] row-major
+ * gamma [n]
+ * beta [n]
+ * grid(m)
+ * block(block_size) -- each block deals with block_size*ITEM_PER_THREAD*2 elements;
+*/
+template<typename T2, typename T>
+__global__ void layernorm_twoPassAlgo_e2(T2* output,
+                                         const T2* input,
+                                         const T2* gamma,
+                                         const T2* beta,
+                                         const int m,
+                                         const int n)
+{
+  const int m_idx = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int bdimx = blockDim.x;
+  __shared__ float s_mean, s_variance;
+  float local_sums[1] = {0.0f};
+  const int n_2 = n / 2;
+  int offset = m_idx * n_2;
+  input += offset;
+  output += offset;
+
+  for (int index = tid; index < n_2; index += bdimx) {
+    const T2 local_val = input[index];
+    local_sums[0] += static_cast<float>(local_val.x) + static_cast<float>(local_val.y);
+  }
+
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_mean = local_sums[0] / n;
+  }
+  __syncthreads();
+
+  local_sums[0] = 0.0f;
+  for (int index = tid; index < n_2; index += bdimx) {
+    const T2 local_val = input[index];
+    const float2 tmp = {static_cast<float>(local_val.x) - s_mean,
+                        static_cast<float>(local_val.y) - s_mean};
+    local_sums[0] += tmp.x * tmp.x + tmp.y * tmp.y;
+  }
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(local_sums[0] / n + 1e-5);
+  }
+  __syncthreads();
+
+  for (int index = tid; index < n_2; index += bdimx) {
+    const T2 local_val = input[index];
+    const T2 gamma_val = gamma[index];
+    const T2 beta_val = beta[index];
+    T2 tmp;
+    tmp.x = T((static_cast<float>(local_val.x) - s_mean)*s_variance*static_cast<float>(gamma_val.x) + static_cast<float>(beta_val.x));
+    tmp.y = T((static_cast<float>(local_val.y) - s_mean)*s_variance*static_cast<float>(gamma_val.y) + static_cast<float>(beta_val.y));
+    output[index] = tmp;
+  }
+}
+
+template <typename T>
+void layernorm(mutlass::MatrixCoord tensor_size,
+               TensorRef<T, layout::RowMajor> ref_output,
+               TensorRef<T, layout::RowMajor> ref_input,
+               TensorRef<T, layout::RowMajor> ref_gamma,
+               TensorRef<T, layout::RowMajor> ref_beta,
+               musaStream_t stream){
+  const int m = tensor_size.row();
+  const int n = tensor_size.column();
+  T* output = ref_output.data();
+  const T* input = ref_input.data();
+  const T* gamma = ref_gamma.data();
+  const T* beta = ref_beta.data();
+  dim3 grid(m);
+  dim3 block((n + 31)/32*32);
+  if (block.x > 1024){
+    block.x = 1024;
+  }
+  // TODO : There should be better configs for different cases, we only use several samples to show how to use here
+  // TODO : using registers to store values locally can reduce the loads from global memory and speedup the kernels.
+  if ((n % 4 == 0) && (n >= 128) && (n <= 4096)) {
+    block.x = (n/4 + 31)/32*32;
+    if (std::is_same<T, float>::value) {
+      layernorm_twoPassAlgo_stored_locally_e4<float4, float, 1><<<grid, block, 0, stream>>>(
+        (float4*)output,
+        (const float4*)input,
+        (const float4*)gamma,
+        (const float4*)beta,
+        m,
+        n);
+    } // if (std::is_same<T, float>::value)
+    else {
+      layernorm_twoPassAlgo_stored_locally_e4<half4, half, 1><<<grid, block, 0, stream>>>(
+        (half4*)output,
+        (const half4*)input,
+        (const half4*)gamma,
+        (const half4*)beta,
+        m,
+        n);
+    }
+  } //if ((n % 4 == 0) && (n >= 128) && (n <= 4096))
+  else if (n % 2 == 0) {
+    if (n / 2 <= 1024) {
+      block.x = (n/2 + 31)/32*32;
+      if (std::is_same<T, float>::value) {
+        layernorm_twoPassAlgo_stored_locally_e2<float2, float, 1><<<grid, block, 0, stream>>>(
+          (float2*)output,
+          (const float2*)input,
+          (const float2*)gamma,
+          (const float2*)beta,
+          m,
+          n);
+      } //if (std::is_same<T, float>::value)
+      else {
+        layernorm_twoPassAlgo_stored_locally_e2<half2, half, 1><<<grid, block, 0, stream>>>(
+          (half2*)output,
+          (const half2*)input,
+          (const half2*)gamma,
+          (const half2*)beta,
+          m,
+          n);
+      }
+    } // if (n / 2 <= 1024)
+    else if (n <= 8192) {
+      block.x = ((n + 7)/8 + 31)/32*32;
+      if (std::is_same<T, float>::value) {
+        layernorm_twoPassAlgo_stored_locally_e2<float2, float, 4><<<grid, block, 0, stream>>>(
+          (float2*)output,
+          (const float2*)input,
+          (const float2*)gamma,
+          (const float2*)beta,
+          m,
+          n);
+      } // if (std::is_same<T, float>::value)
+      else {
+        layernorm_twoPassAlgo_stored_locally_e2<half2, half, 4><<<grid, block, 0, stream>>>(
+          (half2*)output,
+          (const half2*)input,
+          (const half2*)gamma,
+          (const half2*)beta,
+          m,
+          n);
+      }
+    } // if (n <= 8192)
+    else if (n <= 16384) {
+      block.x = ((n + 15)/ 16 + 31)/32*32;
+      if (std::is_same<T, float>::value) {
+        layernorm_twoPassAlgo_stored_locally_e2<float2, float, 8><<<grid, block, 0, stream>>>(
+          (float2*)output,
+          (const float2*)input,
+          (const float2*)gamma,
+          (const float2*)beta,
+          m,
+          n);
+      } // if (std::is_same<T, float>::value)
+      else {
+        layernorm_twoPassAlgo_stored_locally_e2<half2, half, 8><<<grid, block, 0, stream>>>(
+          (half2*)output,
+          (const half2*)input,
+          (const half2*)gamma,
+          (const half2*)beta,
+          m,
+          n);
+      }
+    } // if (n <= 16384)
+    else if (n <= 32768) {
+      block.x = ((n + 31)/32 + 31)/32*32;
+      if (std::is_same<T, float>::value) {
+        layernorm_twoPassAlgo_stored_locally_e2<float2, float, 16><<<grid, block, 0, stream>>>(
+          (float2*)output,
+          (const float2*)input,
+          (const float2*)gamma,
+          (const float2*)beta,
+          m,
+          n);
+      } // if (std::is_same<T, float>::value)
+      else {
+        layernorm_twoPassAlgo_stored_locally_e2<half2, half, 16><<<grid, block, 0, stream>>>(
+          (half2*)output,
+          (const half2*)input,
+          (const half2*)gamma,
+          (const half2*)beta,
+          m,
+          n);
+      }
+    } // if (n <= 32768)
+    else {
+      if (block.x > 512)
+        block.x = 512;
+      if (std::is_same<T, float>::value) {
+        layernorm_twoPassAlgo_e2<float2, float><<<grid, block, 0, stream>>>(
+          (float2 *)output,
+          (const float2 *)input,
+          (const float2 *)gamma,
+          (const float2 *)beta,
+          m,
+          n);
+      } // if (std::is_same<T, float>::value)
+      else {
+        layernorm_twoPassAlgo_e2<half2, half><<<grid, block, 0, stream>>>(
+          (half2 *)output,
+          (const half2 *)input,
+          (const half2 *)gamma,
+          (const half2 *)beta,
+          m,
+          n);
+      }
+    }
+  } // if (n % 2 == 0)
+  else {
+    if (n <= 1024) {
+      layernorm_twoPassAlgo_stored_locally_e1<T, 1><<<grid, block, 0, stream>>>(
+        output,
+        input,
+        gamma,
+        beta,
+        m,
+        n);
+    } // if (n <= 1024)
+    else if (n <= 8192) {
+      block.x = ((n + 7)/8 + 31)/32*32;
+      layernorm_twoPassAlgo_stored_locally_e1<T, 8><<<grid, block, 0, stream>>>(
+        output,
+        input,
+        gamma,
+        beta,
+        m,
+        n);
+    } // if (n <= 8192)
+    else if (n <= 16384) {
+      block.x = ((n + 15)/16 + 32)/32*32;
+      layernorm_twoPassAlgo_stored_locally_e1<T, 16><<<grid, block, 0, stream>>>(
+        output,
+        input,
+        gamma,
+        beta,
+        m,
+        n);
+    } // if (n <= 16384)
+    else if (n <= 32768) {
+      block.x = ((n + 31)/32 + 31)/32*32;
+      layernorm_twoPassAlgo_stored_locally_e1<T, 32><<<grid, block, 0, stream>>>(
+        output,
+        input,
+        gamma,
+        beta,
+        m,
+        n);
+    } // if (n <= 32768)
+    else{
+      if (block.x > 512) {
+        block.x = 512;
+      }
+      layernorm_twoPassAlgo_e1<<<grid, block, 0, stream>>>(
+        output,
+        input,
+        gamma,
+        beta,
+        m,
+        n);
+    }
+  }
+}
+
+} //namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_memory.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_memory.h
new file mode 100644
index 000000000..8d6517029
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_memory.h
@@ -0,0 +1,353 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief C++ interface to MUSA device memory management functions.
+ */
+
+#include <memory>
+
+#include "mutlass/platform/platform.h"
+#include "mutlass/numeric_types.h"
+#include "exceptions.h"
+
+namespace mutlass {
+namespace device_memory {
+
+/******************************************************************************
+ * Allocation lifetime
+ ******************************************************************************/
+
+/// Allocate a buffer of \p count elements of type \p T on the current MUSA device
+template <typename T>
+T* allocate(size_t count = 1) {
+
+  T* ptr = 0;
+  size_t bytes = 0;
+
+  bytes = count * sizeof(T);
+
+  musaError_t musa_error = musaMalloc((void**)&ptr, bytes);
+
+  if (musa_error != musaSuccess) {
+#if (MUTLASS_DEBUG_TRACE_LEVEL > 0)
+    std::ostringstream os;
+    os << "mutlass::device_memory::allocate: musaMalloc failed: bytes=" << bytes;
+    MUTLASS_TRACE_HOST(os.str());
+#endif
+    throw musa_exception("Failed to allocate memory", musa_error);
+  }
+
+#if (MUTLASS_DEBUG_TRACE_LEVEL > 1)
+  else {
+    std::ostringstream os;
+    os << "mutlass::device_memory::allocate: Successful musaMalloc: bytes=" << bytes;
+    MUTLASS_TRACE_HOST(os.str());
+  }
+#endif
+
+  return ptr;
+}
+
+/// Free the buffer pointed to by \p ptr
+template <typename T>
+void free(T* ptr) {
+  if (ptr) {
+    musaError_t musa_error = (musaFree(ptr));
+    if (musa_error != musaSuccess) {
+      throw musa_exception("Failed to free device memory", musa_error);
+    }
+  }
+}
+
+/******************************************************************************
+ * Data movement
+ ******************************************************************************/
+
+template <typename T>
+void copy(T* dst, T const* src, size_t count, musaMemcpyKind kind) {
+  size_t bytes = count * sizeof_bits<T>::value / 8;
+  if (bytes == 0 && count > 0) {
+    bytes = 1;
+  }
+  musaError_t musa_error = (musaMemcpy(dst, src, bytes, kind));
+  if (musa_error != musaSuccess) {
+    throw musa_exception("musaMemcpy() failed", musa_error);
+  }
+}
+
+template <typename T>
+void copy_to_device(T* dst, T const* src, size_t count = 1) {
+  copy(dst, src, count, musaMemcpyHostToDevice);
+}
+
+template <typename T>
+void copy_to_host(T* dst, T const* src, size_t count = 1) {
+  copy(dst, src, count, musaMemcpyDeviceToHost);
+}
+
+template <typename T>
+void copy_device_to_device(T* dst, T const* src, size_t count = 1) {
+  copy(dst, src, count, musaMemcpyDeviceToDevice);
+}
+
+template <typename T>
+void copy_host_to_host(T* dst, T const* src, size_t count = 1) {
+  copy(dst, src, count, musaMemcpyHostToHost);
+}
+
+/// Copies elements from device memory to host-side range
+template <typename OutputIterator, typename T>
+void insert_to_host(OutputIterator begin, OutputIterator end, T const* device_begin) {
+  size_t elements = end - begin;
+  copy_to_host(&*begin, device_begin, elements);
+}
+
+/// Copies elements to device memory from host-side range
+template <typename T, typename InputIterator>
+void insert_to_device(T* device_begin, InputIterator begin, InputIterator end) {
+  size_t elements = end - begin;
+  copy_to_device(device_begin, &*begin, elements);
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace device_memory
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename T>
+class DeviceAllocation {
+public:
+
+  /// Delete functor for MUSA device memory
+  struct deleter {
+    void operator()(T* ptr) {
+      musaError_t musa_error = (musaFree(ptr));
+      if (musa_error != musaSuccess) {
+        // noexcept
+        //                throw musa_exception("musaFree() failed", musa_error);
+        return;
+      }
+    }
+  };
+
+public:
+  //
+  // Data members
+  //
+
+  /// Number of elements of T allocated on the current MUSA device
+  size_t capacity;
+
+  /// Smart pointer
+  platform::unique_ptr<T, deleter> smart_ptr;
+
+public:
+
+  //
+  // Static methods
+  //
+
+  /// Static member to compute the number of bytes needed for a given number of elements
+  static size_t bytes(size_t elements) {
+    if (sizeof_bits<T>::value < 8) {
+      size_t const kElementsPerByte = 8 / sizeof_bits<T>::value;
+      return elements / kElementsPerByte;
+    }
+    else {
+      size_t const kBytesPerElement = sizeof_bits<T>::value / 8;
+      return elements * kBytesPerElement;
+    }
+  }
+
+public:
+
+  //
+  // Methods
+  //
+
+  /// Constructor: allocates no memory
+  DeviceAllocation() : capacity(0) {}
+
+  /// Constructor: allocates \p capacity elements on the current MUSA device
+  DeviceAllocation(size_t _capacity) :
+    smart_ptr(device_memory::allocate<T>(_capacity)), capacity(_capacity) {}
+
+  /// Constructor: allocates \p capacity elements on the current MUSA device taking ownership of the allocation
+  DeviceAllocation(T *ptr, size_t _capacity) : smart_ptr(ptr), capacity(_capacity) {}
+
+  /// Copy constructor
+  DeviceAllocation(DeviceAllocation const &p):
+    smart_ptr(device_memory::allocate<T>(p.capacity)), capacity(p.capacity) {
+
+    device_memory::copy_device_to_device(smart_ptr.get(), p.get(), capacity);
+  }
+
+  /// Move constructor
+  DeviceAllocation(DeviceAllocation &&p): capacity(0) {
+    std::swap(smart_ptr, p.smart_ptr);
+    std::swap(capacity, p.capacity);
+  }
+
+  /// Destructor
+  ~DeviceAllocation() { reset(); }
+
+  /// Returns a pointer to the managed object
+  T* get() const { return smart_ptr.get(); }
+
+  /// Releases the ownership of the managed object (without deleting) and resets capacity to zero
+  T* release() {
+    capacity = 0;
+    return smart_ptr.release();
+  }
+
+  /// Deletes the managed object and resets capacity to zero
+  void reset() {
+    capacity = 0;
+    smart_ptr.reset();
+  }
+
+  /// Deletes managed object, if owned, and allocates a new object
+  void reset(size_t _capacity) {
+    reset(device_memory::allocate<T>(_capacity), _capacity);
+  }
+
+  /// Deletes managed object, if owned, and replaces its reference with a given pointer and capacity
+  void reset(T* _ptr, size_t _capacity) {
+    smart_ptr.reset(_ptr);
+    capacity = _capacity;
+  }
+
+  /// Allocates a new buffer and copies the old buffer into it. The old buffer is then released.
+  void reallocate(size_t new_capacity) {
+
+    platform::unique_ptr<T, deleter> new_allocation(device_memory::allocate<T>(new_capacity));
+
+    device_memory::copy_device_to_device(
+      new_allocation.get(),
+      smart_ptr.get(),
+      std::min(new_capacity, capacity));
+
+    std::swap(smart_ptr, new_allocation);
+    std::swap(new_capacity, capacity);
+  }
+
+  /// Returns the number of elements
+  size_t size() const {
+    return capacity;
+  }
+
+  /// Returns the number of bytes needed to store the allocation
+  size_t bytes() const {
+    return bytes(capacity);
+  }
+
+  /// Returns a pointer to the object owned by *this
+  T* operator->() const { return smart_ptr.get(); }
+
+  /// Returns the deleter object which would be used for destruction of the managed object.
+  deleter& get_deleter() { return smart_ptr.get_deleter(); }
+
+  /// Returns the deleter object which would be used for destruction of the managed object (const)
+  const deleter& get_deleter() const { return smart_ptr.get_deleter(); }
+
+  /// Copies a device-side memory allocation
+  DeviceAllocation & operator=(DeviceAllocation const &p) {
+    if (capacity != p.capacity) {
+      smart_ptr.reset(device_memory::allocate<T>(p.capacity));
+      capacity = p.capacity;
+    }
+    device_memory::copy_device_to_device(smart_ptr.get(), p.get(), capacity);
+    return *this;
+  }
+
+  /// Move assignment
+  DeviceAllocation & operator=(DeviceAllocation && p) {
+    std::swap(smart_ptr, p.smart_ptr);
+    std::swap(capacity, p.capacity);
+    return *this;
+  }
+
+  /// Copies the entire allocation from another location in device memory.
+  void copy_from_device(T const *ptr) const {
+    copy_from_device(ptr, capacity);
+  }
+
+  /// Copies a given number of elements from device memory
+  void copy_from_device(T const *ptr, size_t elements) const {
+    device_memory::copy_device_to_device(get(), ptr, elements);
+  }
+
+  void copy_to_device(T *ptr) const {
+    copy_to_device(ptr, capacity);
+  }
+
+  void copy_to_device(T *ptr, size_t elements) const {
+    device_memory::copy_device_to_device(ptr, get(), elements);
+  }
+
+  void copy_from_host(T const *ptr) const {
+    copy_from_host(ptr, capacity);
+  }
+
+  void copy_from_host(T const *ptr, size_t elements) const {
+    device_memory::copy_to_device(get(), ptr, elements);
+  }
+
+  void copy_to_host(T *ptr) const {
+    copy_to_host(ptr, capacity);
+  }
+
+  void copy_to_host(T *ptr, size_t elements) const {
+    device_memory::copy_to_host(ptr, get(), elements);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace device_memory {
+
+/// Device allocation abstraction that tracks size and capacity
+template <typename T>
+using allocation = mutlass::DeviceAllocation<T>;
+
+}  // namespace device_memory
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nchw_to_nhwc.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nchw_to_nhwc.h
new file mode 100644
index 000000000..7f5327e1c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nchw_to_nhwc.h
@@ -0,0 +1,142 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief musa kernels to transform a device memory tensor from NCHW layout to NHWC layout.
+ */
+
+#include "mutlass/mutlass.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/tensor_ref.h"
+
+namespace mutlass {
+
+/** \brief interface to transform a device memory tensor from NCHW layout to NHWC layout.
+ * \tparam T: data type
+ */
+template <typename T>
+void nchw_to_nhwc(mutlass::Tensor4DCoord input_tensor_size,
+                  mutlass::Tensor4DCoord output_tensor_size,
+                  TensorRef<T, layout::TensorNCHW> ref_input,
+                  TensorRef<T, layout::TensorNHWC> ref_output,
+                  musaStream_t stream);
+
+template <typename T>
+__global__ void nchw_to_nhwc_kernel(T *output,
+                                    const T *input,
+                                    const int n,
+                                    const int h,
+                                    const int w,
+                                    const int c) {
+  const int hw = h*w;
+  const int chw = c*hw;
+  __shared__ T shbuf[32 * (32 + 1)];
+  const int32_t tid  = threadIdx.y*blockDim.x + threadIdx.x;
+  const int32_t wid  = tid / 32;
+  const int32_t lid  = tid % 32;
+  const int32_t ni   = blockIdx.z;
+  const int32_t ci0  = blockIdx.y * 32;
+  const int32_t hwi0 = blockIdx.x * 32;
+
+  const size_t input_idx = ni * chw + (ci0 + wid) * hw + hwi0;
+  const T *A = input + input_idx;
+  if (hwi0 + lid < hw) {
+    const int lid_x_33 = lid * 33;
+    if ((ci0 + 32) <= c) {
+      int ci = wid;  // between 0 and 7
+      MUTLASS_PRAGMA_UNROLL
+      for (int cLoopIdx = 0; cLoopIdx < 4; cLoopIdx++) {
+        shbuf[lid_x_33 + ci] = A[lid];
+        A                    = &A[8 * hw];
+        ci += 8;
+      }
+    } else {
+      for (int ci = wid; ci < 32; ci += 8) {
+        if ((ci + ci0) < c) {
+          shbuf[lid_x_33 + ci] = A[lid];
+        }
+        A = &A[8 * hw];
+      }
+    }
+  }
+  __syncthreads();
+
+  const int32_t ciOut = ci0 + lid;
+  output = &output[ni * chw + ciOut];
+  if (ciOut < c) {
+    if (hwi0 + 32 < hw) {
+      int hwI = wid;
+      MUTLASS_PRAGMA_UNROLL
+      for (int hwLoopIdx = 0; hwLoopIdx < 4; ++hwLoopIdx) {
+        output[(hwi0 + hwI) * c] = shbuf[(hwI)*33 + lid];
+        hwI += 8;
+      }
+    } else {
+      for (int hwI = wid; hwI < 32; hwI += 8) {
+        if (hwi0 + hwI < hw) {
+          output[(hwi0 + hwI) * c] = shbuf[(hwI)*33 + lid];
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+void nchw_to_nhwc(mutlass::Tensor4DCoord input_tensor_size,
+                  mutlass::Tensor4DCoord output_tensor_size,
+                  TensorRef<T, layout::TensorNCHW> ref_input,
+                  TensorRef<T, layout::TensorNHWC> ref_output,
+                  musaStream_t stream) {
+
+  assert(
+    input_tensor_size.n() == output_tensor_size.n() &&
+    input_tensor_size.c() == output_tensor_size.h() &&
+    input_tensor_size.h() == output_tensor_size.w() &&
+    input_tensor_size.w() == output_tensor_size.c());
+
+  int n = output_tensor_size.n();
+  int h = output_tensor_size.h();
+  int w = output_tensor_size.w();
+  int c = output_tensor_size.c();
+
+  dim3 grid((h*w + 31)/32, (c + 31)/32, n);
+  dim3 block(32, 8);
+  nchw_to_nhwc_kernel<<<grid, block, 0, stream>>>(ref_output.data(), ref_input.data(),
+                                                  n, h, w, c);
+}
+
+} //namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_padding.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_padding.h
new file mode 100644
index 000000000..eb511c9c7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_padding.h
@@ -0,0 +1,277 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief musa kernels for padding in device memory with NHWC layout.
+ */
+
+#include "mutlass/mutlass.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/tensor_ref.h"
+
+namespace mutlass {
+
+/** \brief interface for padding in a device memory tensor with NHWC layout
+ * \tparam T: data type
+ */
+template <typename T>
+void nhwc_padding(mutlass::Tensor4DCoord input_tensor_size,
+                  mutlass::Tensor4DCoord output_tensor_size,
+                  TensorRef<T, layout::TensorNHWC> ref_input,
+                  TensorRef<T, layout::TensorNHWC> ref_output,
+                  musaStream_t stream);
+
+
+template <typename T>
+__global__ void nhwc_padding_kernel(const int32_t n,
+                                    const int32_t h,
+                                    const int32_t w,
+                                    const int32_t c_in,
+                                    const int32_t c_out,
+                                    const T zero,
+                                    const T *input,
+                                    T *output){
+
+  const int32_t idx_jump       = blockDim.x * gridDim.x;
+  const int32_t total_elements = n * h * w * c_out;
+
+  int32_t c_idx, w_idx, h_idx, n_idx, resudial;
+
+  T value;
+  for (int32_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < total_elements; idx += idx_jump) {
+
+    c_idx = idx%c_out;
+    if (c_idx >= c_in){
+      value = zero;
+    }
+    else{
+      resudial = idx/c_out;
+      w_idx = resudial%w;
+      resudial = resudial/w;
+      h_idx = resudial%h;
+      n_idx = resudial/h;
+      resudial = ((n_idx * h + h_idx) * w + w_idx) * c_in + c_idx;
+      value = input[resudial];
+    }
+    output[idx] = value;
+  }
+}
+
+
+// fast kernel for c_in = 3 & c_out = 4
+template <typename Tio, typename Telement, int element_in_Tio>
+__global__ void nhwc_padding_channel_3To4_kernel(const int32_t n,
+                                                 const int32_t h,
+                                                 const int32_t w,
+                                                 const Tio *input,
+                                                 Tio *output,
+                                                 const int32_t max_output_element,
+                                                 const int32_t max_input_element,
+                                                 const Tio zero_io,
+                                                 const Telement zero_element){
+  __shared__ Tio shm[192];
+  const int tidx = blockIdx.x * 192 + threadIdx.x;
+  const int threadidx = threadIdx.x;
+
+  shm[threadIdx.x] = tidx >= max_input_element ? zero_io : input[tidx];
+  __syncthreads();
+
+  const int output_offset = blockIdx.x * 256;
+  const int lower_bound = max_output_element < output_offset + 256 ? max_output_element : output_offset + 256;
+  for (int i = output_offset + threadidx, j = threadidx ; i < lower_bound ; i+=192, j+=192)
+  {
+    const Telement* shm_element = (const Telement*)shm + j*3*element_in_Tio/4;
+    Telement array[element_in_Tio];
+    MUTLASS_PRAGMA_UNROLL
+    for (int k = 0 ; k < element_in_Tio ; k++)
+      array[k] = ((k+1)%4 == 0) ? zero_element : shm_element[(k > 3) ? (k - 1) : k];
+    output[i] = *((const Tio *)array);
+  }
+}
+
+// fast kernel for c_in = 3 & c_out = 8
+template <typename Tio, typename Telement, int element_in_Tio>
+__global__ void nhwc_padding_channel_3To8_kernel(const int32_t n,
+                                                 const int32_t h,
+                                                 const int32_t w,
+                                                 const Tio *input,
+                                                 Tio *output,
+                                                 const int32_t max_output_element,
+                                                 const int32_t max_input_element,
+                                                 const Tio zero_io,
+                                                 const Telement zero_element){
+  __shared__ Tio shm[192];
+  const int tidx = blockIdx.x * 192 + threadIdx.x;
+  const int threadidx = threadIdx.x;
+
+  shm[threadIdx.x] = tidx >= max_input_element ? zero_io : input[tidx];
+  __syncthreads();
+
+  const int output_offset = blockIdx.x * 512;
+  const int lower_bound = max_output_element < output_offset + 512 ? max_output_element : output_offset + 512;
+  for (int i = output_offset + threadidx, j = threadidx ; i < lower_bound ; i+=192, j+=192)
+  {
+    const Telement* shm_element = (const Telement*)shm + (element_in_Tio == 4 ? j/2 : j)*3;
+    Telement array[element_in_Tio];
+    //float
+    if (element_in_Tio == 4){
+      MUTLASS_PRAGMA_UNROLL
+      for (int k = 0 ; k < element_in_Tio ; k++)
+        array[k] = ((j % 2) == 1) ? zero_element : ((k >= 3) ? zero_element : shm_element[k]);
+    }
+    //half
+    else{
+      MUTLASS_PRAGMA_UNROLL
+      for (int k = 0 ; k < element_in_Tio ; k++)
+        array[k] = (k >= 3) ? zero_element : shm_element[k];
+    }
+    output[i] = *((const Tio *)array);
+  }
+}
+
+template <typename T>
+void nhwc_padding(mutlass::Tensor4DCoord input_tensor_size,
+                  mutlass::Tensor4DCoord output_tensor_size,
+                  TensorRef<T, layout::TensorNHWC> ref_input,
+                  TensorRef<T, layout::TensorNHWC> ref_output,
+                  musaStream_t stream){
+  assert(
+    input_tensor_size.n() == output_tensor_size.n() &&
+    input_tensor_size.h() == output_tensor_size.h() &&
+    input_tensor_size.w() == output_tensor_size.w() &&
+    input_tensor_size.c() <= output_tensor_size.c());
+
+  int n = input_tensor_size.n();
+  int h = input_tensor_size.h();
+  int w = input_tensor_size.w();
+  int c_in = input_tensor_size.c();
+  int c_out = output_tensor_size.c();
+
+  //case 1 : channel == 3 padding to 4 or 8
+  if ((c_out == 4 || c_out == 8) && c_in == 3 && (n*h*w % 8 == 0)){
+    dim3 block(192);
+    const int nhw = n*h*w;
+    const int nhwc = nhw*c_in;
+    //for half_t
+    if (mutlass::sizeof_bits<T>::value == 16){
+      const int element_in_Tio = 8;
+      const int max_input_element = nhwc/element_in_Tio;
+      const int max_output_element = nhw*c_out/element_in_Tio;
+      const int4 zero_io = {0, 0, 0, 0};
+      const half_t zero_element = static_cast<half_t>(0.0f);
+      dim3 grid((nhwc + 192*element_in_Tio - 1)/(192*element_in_Tio));
+      if (c_out == 4){
+        nhwc_padding_channel_3To4_kernel<int4, half_t, element_in_Tio><<<grid, block, 0, stream>>>
+          (n, h, w,
+          (const int4 *)ref_input.data(),
+          (int4 *)ref_output.data(),
+          max_output_element,
+          max_input_element,
+          zero_io,
+          zero_element);
+      }
+      else if (c_out == 8){
+        nhwc_padding_channel_3To8_kernel<int4, half_t, element_in_Tio><<<grid, block, 0, stream>>>
+          (n, h, w,
+          (const int4 *)ref_input.data(),
+          (int4 *)ref_output.data(),
+          max_output_element,
+          max_input_element,
+          zero_io,
+          zero_element);
+      }
+    }
+    //for float
+    else{
+      const int element_in_Tio = 4;
+      const int max_input_element = nhwc/element_in_Tio;
+      const int max_output_element = nhw*c_out/element_in_Tio;
+      const float4 zero_io = {0.0f, 0.0f, 0.0f, 0.0f};
+      const float zero_element = 0.0f;
+      dim3 grid((nhwc + 192*element_in_Tio - 1)/(192*element_in_Tio));
+      if (c_out == 4){
+        nhwc_padding_channel_3To4_kernel<float4, float, element_in_Tio><<<grid, block, 0, stream>>>
+          (n, h, w,
+          (const float4 *)ref_input.data(),
+          (float4 *)ref_output.data(),
+          max_output_element,
+          max_input_element,
+          zero_io,
+          zero_element);
+      }
+      else if (c_out == 8){
+        nhwc_padding_channel_3To8_kernel<float4, float, element_in_Tio><<<grid, block, 0, stream>>>
+          (n, h, w,
+          (const float4 *)ref_input.data(),
+          (float4 *)ref_output.data(),
+          max_output_element,
+          max_input_element,
+          zero_io,
+          zero_element);
+      }
+    }
+  }
+  //case 2 : even channel
+  else if ((c_out % 2) == 0 && (c_in % 2) == 0){
+    int32_t total_elements = n * h * w * c_out / 2;
+    int block_size = 256;
+    dim3 grid((total_elements + 255)/256);
+    dim3 block(block_size);
+    //for half_t
+    if (mutlass::sizeof_bits<T>::value == 16){
+      const __half2 zero  = {0.0f, 0.0f};
+      nhwc_padding_kernel<<<grid, block, 0, stream>>>(n, h, w, c_in/2, c_out/2, zero, (const __half2*)ref_input.data(), (__half2*)ref_output.data());
+    }
+    //for float
+    else{
+      const float2 zero  = {0.0f, 0.0f};
+      nhwc_padding_kernel<<<grid, block, 0, stream>>>(n, h, w, c_in/2, c_out/2, zero, (const float2*)ref_input.data(), (float2*)ref_output.data());
+    }
+  }
+  //case 3 : odd channel
+  else{
+    int32_t total_elements = n * h * w * c_out;
+    int block_size = 256;
+    dim3 grid((total_elements + 255)/256);
+    dim3 block(block_size);
+    const T zero = static_cast<T>(0.0f);
+    nhwc_padding_kernel<<<grid, block, 0, stream>>>(n, h, w, c_in, c_out, zero, ref_input.data(), ref_output.data());
+  }
+}
+
+
+} //namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_pooling.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_pooling.h
new file mode 100644
index 000000000..17f07139c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_pooling.h
@@ -0,0 +1,577 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief musa kernels to do avg/max pooling on a device memory tensor with NHWC layout.
+ */
+
+#include "mutlass/mutlass.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/tensor_ref.h"
+#include "device_utils.h"
+#include <float.h>
+
+namespace mutlass {
+
+/** \brief interface to do avg/max pooling on a device memory tensor with NHWC layout.
+ * \tparam T: data type
+ */
+template <typename T>
+void pooling_nhwc(mutlass::Tensor4DCoord input_tensor_size,
+                  mutlass::Tensor4DCoord filter_tensor_size,
+                  mutlass::Tensor4DCoord output_tensor_size,
+                  mutlass::MatrixCoord padding,
+                  mutlass::MatrixCoord stride,
+                  TensorRef<T, layout::TensorNHWC> ref_input,
+                  TensorRef<T, layout::TensorNHWC> ref_output,
+                  int poolingType, //0 for avg pooling ; 1 for max pooling
+                  musaStream_t stream);
+
+/** get the output size of pooling
+ */
+inline int getOutputSize(int H_W, int padding, int kernel_size, int stride)
+{
+    return (H_W + 2 * padding - kernel_size) / stride + 1;
+}
+
+/**
+ * input is [N, H, W, C]
+ * assume stride == kernel_size
+ * output_h = (H + 2*padding_H - kernel_H)/stride_H
+ * output_w = (W + 2*padding_W - kernel_W)/stride_W
+ * output is [N, output_h, output_w, C]
+ * grid(N, output_h, output_w)
+ * block(min(C, 256)) :
+ * each block deals with C elements of output when each thread deals with ((C + 255)/256 element of output)
+*/
+template<typename T, bool IS_AVG_POOLING>
+__global__ void pooling_nhwc_element1_kernel(T* output,
+                                             const T* input,
+                                             const int N,
+                                             const int H,
+                                             const int W,
+                                             const int C,
+                                             const int output_H,
+                                             const int output_W,
+                                             const int kernel_H,
+                                             const int kernel_W,
+                                             const int stride_H,
+                                             const int stride_W,
+                                             const int padding_H,
+                                             const int padding_W)
+{
+  const int tid = threadIdx.x;
+  const int n_idx = blockIdx.x;
+  const int output_h_idx = blockIdx.y;
+  const int output_w_idx = blockIdx.z;
+
+  int h_start_idx = output_h_idx * stride_H - padding_H;
+  int h_end_idx = h_start_idx + kernel_H;
+  h_start_idx = (h_start_idx < 0) ? 0 : h_start_idx;
+  h_end_idx = h_end_idx > H ? H : h_end_idx;
+
+  int w_start_idx = output_w_idx * stride_W - padding_W;
+  int w_end_idx = w_start_idx + kernel_W;
+  w_start_idx = (w_start_idx < 0) ? 0 : w_start_idx;
+  w_end_idx = w_end_idx > W ? W : w_end_idx;
+
+  input += n_idx * H * W * C;
+  output += ((n_idx * output_H + output_h_idx) * output_W + output_w_idx) * C;
+  const int kernel_size2 = kernel_H * kernel_W;
+  for (int c_idx = tid; c_idx < C; c_idx += blockDim.x) {
+    float pooling;
+    if (IS_AVG_POOLING){
+      pooling = 0.0f;
+    }
+    else{
+      pooling = -FLT_MAX;
+    }
+    for (int h = h_start_idx; h < h_end_idx; h++) {
+      for (int w = w_start_idx; w < w_end_idx; w++) {
+        const int idx = (h * W + w) * C;
+        const float tmp = static_cast<float>(input[idx + c_idx]);
+        if (IS_AVG_POOLING){
+          pooling = pooling + tmp;
+        }
+        else{
+          pooling = pooling > tmp ? pooling : tmp;
+        }
+      }
+    }
+
+    T output_val;
+    if (IS_AVG_POOLING){
+      output_val = T(pooling/kernel_size2);
+    }
+    else{
+      output_val = T(pooling);
+    }
+    output[c_idx] = output_val;
+  }
+}
+
+template<typename T2, typename T, bool IS_AVG_POOLING>
+__global__ void pooling_nhwc_element2_kernel(T2* output,
+                                             const T2* input,
+                                             const int N,
+                                             const int H,
+                                             const int W,
+                                             const int C,
+                                             const int output_H,
+                                             const int output_W,
+                                             const int kernel_H,
+                                             const int kernel_W,
+                                             const int stride_H,
+                                             const int stride_W,
+                                             const int padding_H,
+                                             const int padding_W)
+{
+  const int tid = threadIdx.x;
+  const int n_idx = blockIdx.x;
+  const int output_h_idx = blockIdx.y;
+  const int output_w_idx = blockIdx.z;
+
+  int h_start_idx = output_h_idx * stride_H - padding_H;
+  int h_end_idx = h_start_idx + kernel_H;
+  h_start_idx = (h_start_idx < 0) ? 0 : h_start_idx;
+  h_end_idx = h_end_idx > H ? H : h_end_idx;
+
+  int w_start_idx = output_w_idx * stride_W - padding_W;
+  int w_end_idx = w_start_idx + kernel_W;
+  w_start_idx = (w_start_idx < 0) ? 0 : w_start_idx;
+  w_end_idx = w_end_idx > W ? W : w_end_idx;
+
+  input += n_idx * H * W * C;
+  output += ((n_idx * output_H + output_h_idx) * output_W + output_w_idx) * C;
+  const int kernel_size2 = kernel_H * kernel_W;
+  for (int c_idx = tid; c_idx < C; c_idx += blockDim.x) {
+    float2 pooling;
+    if (IS_AVG_POOLING) {
+      pooling = {0.0f, 0.0f};
+    }
+    else {
+      pooling = {-FLT_MAX, -FLT_MAX};
+    }
+    for (int h = h_start_idx; h < h_end_idx; h++) {
+      for (int w = w_start_idx; w < w_end_idx; w++) {
+        const int idx = (h * W + w) * C;
+        const T2 tmp = input[idx + c_idx];
+        const float2 tmp_flt2 = {static_cast<float>(tmp.x), static_cast<float>(tmp.y)};
+        if (IS_AVG_POOLING) {
+          pooling.x += tmp_flt2.x;
+          pooling.y += tmp_flt2.y;
+        }
+        else {
+          pooling.x = pooling.x > tmp_flt2.x ? pooling.x : tmp_flt2.x;
+          pooling.y = pooling.y > tmp_flt2.y ? pooling.y : tmp_flt2.y;
+        }
+      }
+    }
+
+    T2 output_val;
+    if (IS_AVG_POOLING) {
+      output_val.x = T(pooling.x/kernel_size2);
+      output_val.y = T(pooling.y/kernel_size2);
+    }
+    else {
+      output_val.x = T(pooling.x);
+      output_val.y = T(pooling.y);
+    }
+    output[c_idx] = output_val;
+  }
+}
+
+/**
+ * output [N, 1, 1, C]
+ * input [N, H, W, C]
+ * grid(C, N)
+ * block(block_size) -- each block deals with H*W/block_size elements;
+*/
+template<typename T, bool IS_AVG_POOLING>
+__global__ void pooling_nxhTo1x1_element1_kernel(
+    T* output, const T* input, const int N, const int HW, const int C)
+{
+    const int c_idx = blockIdx.x;
+    const int n_idx = blockIdx.y;
+    float pooling[1];
+    if (IS_AVG_POOLING) {
+      pooling[0] = 0.0f;
+    }
+    else {
+      pooling[0] = -FLT_MAX;
+    }
+    const size_t input_offset = n_idx * HW * C + c_idx;
+    input += input_offset;
+    const size_t output_offset = n_idx * C + c_idx;
+    output += output_offset;
+    int tid = threadIdx.x;
+
+    for (int index = tid; index < HW; index += blockDim.x) {
+        float val = static_cast<float>(input[index * C]);
+        if (IS_AVG_POOLING) {
+          pooling[0] += val;
+        }
+        else {
+          pooling[0] = pooling[0] > val ? pooling[0] : val;
+        }
+    }
+    if (blockDim.x <= 32) {
+        if (IS_AVG_POOLING) {
+          warpReduceSum<float, 1>(pooling);
+        }
+        else {
+          warpReduceMax<float, 1>(pooling);
+        }
+    }
+    else {
+        if (IS_AVG_POOLING) {
+          blockReduceSum<float, 1>(pooling);
+        }
+        else {
+          blockReduceMax<float, 1>(pooling);
+        }
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) {
+        T output_val;
+        if (IS_AVG_POOLING) {
+          output_val = T(pooling[0] / HW);
+        }
+        else {
+          output_val = T(pooling[0]);
+        }
+        output[0] = output_val;
+    }
+}
+
+
+/**
+ * output [N, 1, 1, C]
+ * input [N, H, W, C]
+ * grid(C/2, N)
+ * block(block_size) -- each thread deals with H*W/block_size * 2 elements;
+*/
+template<typename T2, typename T, bool IS_AVG_POOLING>
+__global__ void pooling_nxhTo1x1_element2_kernel(
+    T2* output, const T2* input, const int N, const int HW, const int C)
+{
+    const int c_idx = blockIdx.x;
+    const int n_idx = blockIdx.y;
+    float pooling[2];
+    if (IS_AVG_POOLING) {
+      pooling[0] = pooling[1] = 0.0f;
+    }
+    else {
+      pooling[0] = pooling[1] = -FLT_MAX;
+    }
+    const int C_2 = C / 2;
+    const size_t input_offset = n_idx * HW * C_2 + c_idx;
+    input += input_offset;
+    const size_t output_offset = n_idx * C_2 + c_idx;
+    output += output_offset;
+    int tid = threadIdx.x;
+
+    for (int index = tid; index < HW; index += blockDim.x) {
+        T2 val = input[index * C_2];
+        float2 val_flt2 = {static_cast<float>(val.x), static_cast<float>(val.y)};
+        if (IS_AVG_POOLING) {
+          pooling[0] += val_flt2.x;
+          pooling[1] += val_flt2.y;
+        }
+        else {
+          pooling[0] = pooling[0] > val_flt2.x ? pooling[0] : val_flt2.x;
+          pooling[1] = pooling[1] > val_flt2.y ? pooling[1] : val_flt2.y;
+        }
+    }
+    if (blockDim.x <= 32) {
+        if (IS_AVG_POOLING) {
+          warpReduceSum<float, 2>(pooling);
+        }
+        else {
+          warpReduceMax<float, 2>(pooling);
+        }
+    }
+    else {
+        if (IS_AVG_POOLING) {
+          blockReduceSum<float, 2>(pooling);
+        }
+        else {
+          blockReduceMax<float, 2>(pooling);
+        }
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) {
+        T2 output_val;
+        if (IS_AVG_POOLING) {
+          output_val.x = T(pooling[0] / HW);
+          output_val.y = T(pooling[1] / HW);
+        }
+        else {
+          output_val.x = T(pooling[0]);
+          output_val.y = T(pooling[1]);
+        }
+        output[0] = output_val;
+    }
+}
+
+template <typename T>
+void pooling_nhwc(mutlass::Tensor4DCoord input_tensor_size,
+                  mutlass::Tensor4DCoord filter_tensor_size,
+                  mutlass::Tensor4DCoord output_tensor_size,
+                  mutlass::Tensor4DCoord padding,
+                  mutlass::MatrixCoord stride,
+                  TensorRef<T, layout::TensorNHWC> ref_input,
+                  TensorRef<T, layout::TensorNHWC> ref_output,
+                  int poolingType, //0 for avg pooling ; 1 for max pooling
+                  musaStream_t stream) {
+
+  assert(input_tensor_size.n() == output_tensor_size.n() &&
+         input_tensor_size.c() == output_tensor_size.c());
+
+  assert(filter_tensor_size.h() == stride.row() &&
+         filter_tensor_size.w() == stride.column());
+
+  const int N = input_tensor_size.n();
+  const int H = input_tensor_size.h();
+  const int W = input_tensor_size.w();
+  const int C = input_tensor_size.c();
+  const int padding_H = padding.h();
+  const int padding_W = padding.w();
+  const int kernel_H = filter_tensor_size.h();
+  const int kernel_W = filter_tensor_size.w();
+  const int stride_H = stride.row();
+  const int stride_W = stride.column();
+
+  const int output_H = getOutputSize(H, padding_H, kernel_H, stride_H);
+  const int output_W = getOutputSize(W, padding_W, kernel_W, stride_W);
+
+  assert(output_tensor_size.h() == output_H &&
+         output_tensor_size.w() == output_W);
+
+  if (C % 2 != 0) {
+    if ((H == kernel_H && padding_H == 0) && (W == kernel_W && padding_W == 0)) {
+      dim3 grid(C, N);
+      dim3 block(256);
+      if (H*W < block.x){
+        block.x = (H*W + 31)/32*32;
+      }
+      if (poolingType == 0) {
+        pooling_nxhTo1x1_element1_kernel<T, true><<<grid, block, 0, stream>>>(
+          ref_output.data(),
+          ref_input.data(),
+          N,
+          H*W,
+          C);
+      } // if (poolingType == 0)
+      else {
+        pooling_nxhTo1x1_element1_kernel<T, false><<<grid, block, 0, stream>>>(
+          ref_output.data(),
+          ref_input.data(),
+          N,
+          H*W,
+          C);
+      }
+    } // if ((H == kernel_H && padding_H == 0) && (W == kernel_W && padding_W == 0))
+    else {
+      dim3 grid(N, output_H, output_W);
+      dim3 block(256);
+      if (C < block.x) {
+        block.x = C;
+      }
+      if (poolingType == 0) {
+        pooling_nhwc_element1_kernel<T, true><<<grid, block, 0, stream>>>(
+          ref_output.data(),
+          ref_input.data(),
+          N,
+          H,
+          W,
+          C,
+          output_H,
+          output_W,
+          kernel_H,
+          kernel_W,
+          stride_H,
+          stride_W,
+          padding_H,
+          padding_W);
+      } // if (poolingType == 0)
+      else {
+        pooling_nhwc_element1_kernel<T, false><<<grid, block, 0, stream>>>(
+          ref_output.data(),
+          ref_input.data(),
+          N,
+          H,
+          W,
+          C,
+          output_H,
+          output_W,
+          kernel_H,
+          kernel_W,
+          stride_H,
+          stride_W,
+          padding_H,
+          padding_W);
+      }
+    }
+  } // if (C % 2 != 0))
+  else {
+    if ((H == kernel_H && padding_H == 0) && (W == kernel_W && padding_W == 0)) {
+      dim3 grid(C/2, N);
+      dim3 block(256);
+      if (H*W < block.x){
+        block.x = (H*W + 31)/32*32;
+      }
+      if (poolingType == 0) {
+        if (std::is_same<T, float>::value) {
+          pooling_nxhTo1x1_element2_kernel<float2, float, true><<<grid, block, 0, stream>>>(
+            (float2*)(ref_output.data()),
+            (const float2*)(ref_input.data()),
+            N,
+            H*W,
+            C);
+        } // if (std::is_same<T, float>::value)
+        else {
+          pooling_nxhTo1x1_element2_kernel<half2, half, true><<<grid, block, 0, stream>>>(
+            (half2*)(ref_output.data()),
+            (const half2*)(ref_input.data()),
+            N,
+            H*W,
+            C);
+        }
+      } // if (poolingType == 0)
+      else {
+        if (std::is_same<T, float>::value) {
+          pooling_nxhTo1x1_element2_kernel<float2, float, false><<<grid, block, 0, stream>>>(
+            (float2*)(ref_output.data()),
+            (const float2*)(ref_input.data()),
+            N,
+            H*W,
+            C);
+        } // if (std::is_same<T, float>::value)
+        else {
+          pooling_nxhTo1x1_element2_kernel<half2, half, false><<<grid, block, 0, stream>>>(
+            (half2*)(ref_output.data()),
+            (const half2*)(ref_input.data()),
+            N,
+            H*W,
+            C);
+        }
+      }
+    } // if ((H == kernel_H && padding_H == 0) && (W == kernel_W && padding_W == 0))
+    else {
+      dim3 grid(N, output_H, output_W);
+      dim3 block(256);
+      if (C/2 < block.x) {
+        block.x = C/2;
+      }
+      if (poolingType == 0) {
+        if (std::is_same<T, float>::value) {
+          pooling_nhwc_element2_kernel<float2, float, true><<<grid, block, 0, stream>>>(
+            (float2*)(ref_output.data()),
+            (const float2*)(ref_input.data()),
+            N,
+            H,
+            W,
+            C/2,
+            output_H,
+            output_W,
+            kernel_H,
+            kernel_W,
+            stride_H,
+            stride_W,
+            padding_H,
+            padding_W);
+        } // if (std::is_same<T, float>::value)
+        else {
+          pooling_nhwc_element2_kernel<half2, half, true><<<grid, block, 0, stream>>>(
+            (half2*)(ref_output.data()),
+            (const half2*)(ref_input.data()),
+            N,
+            H,
+            W,
+            C/2,
+            output_H,
+            output_W,
+            kernel_H,
+            kernel_W,
+            stride_H,
+            stride_W,
+            padding_H,
+            padding_W);
+        }
+      } // if (poolingType == 0)
+      else {
+        if (std::is_same<T, float>::value) {
+          pooling_nhwc_element2_kernel<float2, float, false><<<grid, block, 0, stream>>>(
+            (float2*)(ref_output.data()),
+            (const float2*)(ref_input.data()),
+            N,
+            H,
+            W,
+            C/2,
+            output_H,
+            output_W,
+            kernel_H,
+            kernel_W,
+            stride_H,
+            stride_W,
+            padding_H,
+            padding_W);
+        } // if (std::is_same<T, float>::value)
+        else {
+          pooling_nhwc_element2_kernel<half2, half, false><<<grid, block, 0, stream>>>(
+            (half2*)(ref_output.data()),
+            (const half2*)(ref_input.data()),
+            N,
+            H,
+            W,
+            C/2,
+            output_H,
+            output_W,
+            kernel_H,
+            kernel_W,
+            stride_H,
+            stride_W,
+            padding_H,
+            padding_W);
+        }
+      }
+    }
+  }
+}
+
+} //namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_to_nchw.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_to_nchw.h
new file mode 100644
index 000000000..41aad9943
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_nhwc_to_nchw.h
@@ -0,0 +1,145 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief musa kernels to transform a device memory tensor from NHWC layout to NCHW layout.
+ */
+
+#include "mutlass/mutlass.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/tensor_ref.h"
+
+namespace mutlass {
+
+/** \brief interface to transform a device memory tensor from NHWC layout to NCHW layout.
+ * \tparam T: data type
+ */
+template <typename T>
+void nhwc_to_nchw(mutlass::Tensor4DCoord input_tensor_size,
+                  mutlass::Tensor4DCoord output_tensor_size,
+                  TensorRef<T, layout::TensorNHWC> ref_input,
+                  TensorRef<T, layout::TensorNCHW> ref_output,
+                  musaStream_t stream);
+
+
+template <typename T>
+__global__ void nhwc_to_nchw_kernel(T *output,
+                                    const T *input,
+                                    const int n,
+                                    const int h,
+                                    const int w,
+                                    const int c) {
+
+  const int hw = h*w;
+  const int hwc = hw*c;
+  __shared__ T shbuf[32 * (32 + 1)];
+  const int32_t tid  = threadIdx.y*blockDim.x + threadIdx.x;
+  const int32_t wid  = tid / 32;
+  const int32_t lid  = tid % 32;
+  const int32_t ni   = blockIdx.z;
+  const int32_t hwi0  = blockIdx.y * 32;
+  const int32_t ci0 = blockIdx.x * 32;
+
+  const size_t input_idx = ni * hwc + (hwi0 + wid) * c + ci0;
+  const T *A = input + input_idx;
+  if (ci0 + lid < c) {
+    const int lid_x_33 = lid * 33;
+    if ((hwi0 + 32) <= hw) {
+      int hwi = wid;  // between 0 and 7
+      MUTLASS_PRAGMA_UNROLL
+      for (int cLoopIdx = 0; cLoopIdx < 4; cLoopIdx++) {
+        shbuf[lid_x_33 + hwi] = A[lid];
+        A                     = &A[8 * c];
+        hwi += 8;
+      }
+    } else {
+      for (int hwi = wid; hwi < 32; hwi += 8) {
+        if ((hwi + hwi0) < hw) {
+          shbuf[lid_x_33 + hwi] = A[lid];
+        }
+        A = &A[8 * c];
+      }
+    }
+  }
+  __syncthreads();
+
+  const int32_t hwiOut = hwi0 + lid;
+  output = &output[ni * hwc + hwiOut];
+  if (hwiOut < hw) {
+    if (ci0 + 32 < c) {
+      int cI = wid;
+      MUTLASS_PRAGMA_UNROLL
+      for (int hwLoopIdx = 0; hwLoopIdx < 4; ++hwLoopIdx) {
+        output[(ci0 + cI) * hw] = shbuf[(cI)*33 + lid];
+        cI += 8;
+      }
+    } else {
+      for (int cI = wid; cI < 32; cI += 8) {
+        if (ci0 + cI < c) {
+          output[(ci0 + cI) * hw] = shbuf[(cI)*33 + lid];
+        }
+      }
+    }
+  }
+}
+
+template <typename T>
+void nhwc_to_nchw(mutlass::Tensor4DCoord input_tensor_size,
+                  mutlass::Tensor4DCoord output_tensor_size,
+                  TensorRef<T, layout::TensorNHWC> ref_input,
+                  TensorRef<T, layout::TensorNCHW> ref_output,
+                  musaStream_t stream) {
+
+  assert(
+    input_tensor_size.n() == output_tensor_size.n() &&
+    input_tensor_size.h() == output_tensor_size.c() &&
+    input_tensor_size.w() == output_tensor_size.h() &&
+    input_tensor_size.c() == output_tensor_size.w());
+
+  int n = input_tensor_size.n();
+  int h = input_tensor_size.h();
+  int w = input_tensor_size.w();
+  int c = input_tensor_size.c();
+
+  dim3 grid((c + 31)/32, (h*w + 31)/32, n);
+  dim3 block(32, 8);
+  nhwc_to_nchw_kernel<<<grid, block, 0, stream>>>(ref_output.data(), ref_input.data(),
+                                                  n, h, w, c);
+
+}
+
+} //namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_rmsnorm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_rmsnorm.h
new file mode 100644
index 000000000..32d1e5f2e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_rmsnorm.h
@@ -0,0 +1,187 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_coord.h"
+#include "mutlass/tensor_ref.h"
+#include "mutlass/util/device_utils.h"
+#include <float.h>
+
+namespace mutlass {
+
+__global__ void rmsnorm_twoPassAlgo_e8(float4 *output, const float4 *input,
+				       const float4 *weight,
+				       const int m, const int n, float epsilon) {
+  const int m_idx = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int bdimx = blockDim.x;
+  __shared__ float s_mean;
+  float local_sums[1] = {0.0f};
+  const int n_8 = n / 8;
+  int offset = m_idx * n_8;
+  input += offset;
+  output += offset;
+
+  for (int index = tid; index < n_8; index += bdimx) {
+    const float4 local_val = input[index];
+    const half2 *h1 = (half2 *)&local_val.x;
+    const half2 *h2 = (half2 *)&local_val.y;
+    const half2 *h3 = (half2 *)&local_val.z;
+    const half2 *h4 = (half2 *)&local_val.w;
+    local_sums[0] += static_cast<float>(h1->x) * static_cast<float>(h1->x) +
+                     static_cast<float>(h1->y) * static_cast<float>(h1->y) +
+                     static_cast<float>(h2->x) * static_cast<float>(h2->x) +
+                     static_cast<float>(h2->y) * static_cast<float>(h2->y) +
+                     static_cast<float>(h3->x) * static_cast<float>(h3->x) +
+                     static_cast<float>(h3->y) * static_cast<float>(h3->y) +
+                     static_cast<float>(h4->x) * static_cast<float>(h4->x) +
+                     static_cast<float>(h4->y) * static_cast<float>(h4->y);
+  }
+
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  } else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_mean = rsqrtf(local_sums[0] / n + epsilon);
+  }
+  __syncthreads();
+
+  for (int index = tid; index < n_8; index += bdimx) {
+    const float4 local_val = input[index];
+    const float4 weight_val = weight[index];
+
+    const half2 *l1 = (half2 *)&local_val.x;
+    const half2 *l2 = (half2 *)&local_val.y;
+    const half2 *l3 = (half2 *)&local_val.z;
+    const half2 *l4 = (half2 *)&local_val.w;
+
+    const half2 *g1 = (half2 *)&weight_val.x;
+    const half2 *g2 = (half2 *)&weight_val.y;
+    const half2 *g3 = (half2 *)&weight_val.z;
+    const half2 *g4 = (half2 *)&weight_val.w;
+
+    float4 tmp;
+    half2 *h1 = (half2 *)&tmp.x;
+    half2 *h2 = (half2 *)&tmp.y;
+    half2 *h3 = (half2 *)&tmp.z;
+    half2 *h4 = (half2 *)&tmp.w;
+
+    h1->x = half(static_cast<float>(l1->x) * s_mean * static_cast<float>(g1->x));
+    h1->y = half(static_cast<float>(l1->y) * s_mean * static_cast<float>(g1->y));
+    h2->x = half(static_cast<float>(l2->x) * s_mean * static_cast<float>(g2->x));
+    h2->y = half(static_cast<float>(l2->y) * s_mean * static_cast<float>(g2->y));
+    h3->x = half(static_cast<float>(l3->x) * s_mean * static_cast<float>(g3->x));
+    h3->y = half(static_cast<float>(l3->y) * s_mean * static_cast<float>(g3->y));
+    h4->x = half(static_cast<float>(l4->x) * s_mean * static_cast<float>(g4->x));
+    h4->y = half(static_cast<float>(l4->y) * s_mean * static_cast<float>(g4->y));
+
+    output[index] = tmp;
+  }
+}
+
+template<typename T>
+__global__ void rmsnorm_twoPassAlgo_e1(T* output,
+				       const T* input,
+				       const T* weight,
+				       const int m, const int n,
+				       float epsilon)
+{
+  const int m_idx = blockIdx.x;
+  const int tid = threadIdx.x;
+  const int bdimx = blockDim.x;
+  __shared__ float s_mean;
+  float local_sums[1] = {0.0f};
+  int offset = m_idx * n;
+  input += offset;
+  output += offset;
+
+  for (int index = tid ; index < n ; index += bdimx){
+    float local_val = static_cast<float>(input[index]);
+    local_sums[0] += local_val * local_val;
+  }
+  if (blockDim.x <= 32) {
+    warpReduceSum<float, 1>(local_sums);
+  }
+  else {
+    blockReduceSum<float, 1>(local_sums);
+  }
+  if (threadIdx.x == 0) {
+    s_mean = rsqrtf(local_sums[0] / n + epsilon);
+  }
+  __syncthreads();
+
+  for (int index = tid ; index < n ; index += bdimx){
+    const T weight_val = weight[index];
+    const T local_val = input[index];
+    output[index] = T(static_cast<float>(local_val) * s_mean * static_cast<float>(weight_val));
+  }
+}
+
+template <typename T>
+void rmsnorm(mutlass::MatrixCoord tensor_size,
+             TensorRef<T, layout::RowMajor> ref_output,
+             TensorRef<T, layout::RowMajor> ref_input,
+             TensorRef<T, layout::RowMajor> ref_weight,
+             musaStream_t stream, float epsilon = 1e-5){
+  const int m = tensor_size.row();
+  const int n = tensor_size.column();
+  T* output = ref_output.data();
+  const T* input = ref_input.data();
+  const T* weight = ref_weight.data();
+  dim3 grid(m);
+
+  if (n % 8 == 0 && std::is_same<T, mutlass::half_t>::value) {
+    dim3 block(min(1024, (n / 8 + 31) / 32 * 32));
+
+    rmsnorm_twoPassAlgo_e8<<<grid, block, 0, stream>>>(
+        (float4 *)output, (const float4 *)input, (const float4 *)weight, m, n, epsilon);
+  } else {
+    dim3 block(min(1024, ((n + 31)/32 + 31)/32*32));
+
+    rmsnorm_twoPassAlgo_e1<<<grid, block, 0, stream>>>(
+        output, input, weight, m, n, epsilon);
+  }
+
+  auto result = musaGetLastError();
+  if (result != musaSuccess) {
+    std::cerr << "MUSA error: " << musaGetErrorString(result) << std::endl;
+    abort();
+  }
+}
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/device_utils.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_utils.h
new file mode 100644
index 000000000..a2001cfc9
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/device_utils.h
@@ -0,0 +1,127 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 utils code for device mutlass code
+*/
+
+#pragma once
+
+#include <musa_fp16.h>
+#include <float.h>
+#define FINAL_MASK 0xffffffff
+
+struct half4 {
+    half x, y, z, w;
+};
+
+template<typename T, int NUM>
+__inline__ __device__ T warpReduceSum(T* val)
+{
+#pragma unroll
+    for (int i = 0; i < NUM; i++) {
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1)
+            val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32);
+    }
+    return (T)(0.0f);
+}
+
+template<typename T, int NUM>
+__inline__ __device__ T blockReduceSum(T* val)
+{
+    __shared__ T shared[NUM][33];
+    int lane = threadIdx.x & 0x1f;
+    int wid = threadIdx.x >> 5;
+
+    warpReduceSum<T, NUM>(val);
+
+    if (lane == 0) {
+#pragma unroll
+        for (int i = 0; i < NUM; i++) {
+            shared[i][wid] = val[i];
+        }
+    }
+
+    __syncthreads();
+
+    bool is_mask = threadIdx.x < (blockDim.x / 32.f);
+#pragma unroll
+    for (int i = 0; i < NUM; i++) {
+        val[i] = is_mask ? shared[i][lane] : (T)(0.0f);
+    }
+    warpReduceSum<T, NUM>(val);
+    return (T)0.0f;
+}
+
+template<typename T, int NUM>
+__inline__ __device__ T warpReduceMax(T* val)
+{
+#pragma unroll
+    for (int i = 0; i < NUM; i++) {
+#pragma unroll
+        for (int mask = 16; mask > 0; mask >>= 1)
+            val[i] = max(val[i], __shfl_xor_sync(FINAL_MASK, val[i], mask, 32));
+    }
+    return (T)(0.0f);
+}
+
+template<typename T, int NUM>
+__inline__ __device__ T blockReduceMax(T* val)
+{
+    static __shared__ T shared[32][NUM];
+    int lane = threadIdx.x & 0x1f;  // in-warp idx
+    int wid = threadIdx.x >> 5;     // warp idx
+
+    warpReduceMax<T, NUM>(val);  // get maxx in each warp
+
+    if (lane == 0)  // record in-warp maxx by warp Idx
+    {
+#pragma unroll
+        for (int i = 0; i < NUM; i++) {
+            shared[wid][i] = val[i];
+        }
+    }
+
+    __syncthreads();
+
+    // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
+    // blockDim.x is not divided by 32
+    bool is_mask = threadIdx.x < (blockDim.x / 32.f);
+#pragma unroll
+    for (int i = 0; i < NUM; i++) {
+        val[i] = is_mask ? shared[lane][i] : (T)(-FLT_MAX);
+    }
+    warpReduceMax<T, NUM>(val);
+
+    return (T)0.0f;
+}
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/distribution.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/distribution.h
new file mode 100644
index 000000000..6a7d3a300
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/distribution.h
@@ -0,0 +1,155 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+/*! \file
+    \brief This header contains a class to parametrize a statistical distribution function.
+*/
+
+#include <ostream>
+
+namespace mutlass {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Distribution type
+struct Distribution {
+  /// Variant types
+  enum Kind { Invalid, Uniform, Gaussian, Identity, Sequential, AllZeros, AllOnes };
+
+  /// Distribution state
+  union {
+    /// Uniform distribution
+    struct {
+      double min;
+      double max;
+      // Percent elements set to NaN
+      double pnan;
+    } uniform;
+
+    /// Gaussian distribution
+    struct {
+      double mean;
+      double stddev;
+      double pnz;
+      double pnzA;
+      double pnzB;
+      double pnzC;
+    } gaussian;
+
+    /// Elements are linear combination of row and column index
+    struct {
+      double start;
+      double delta;
+    } sequential;
+  };
+
+  /// Active variant kind
+  Kind kind;
+
+  /// Random values are cast to integer after scaling by this power of two
+  int int_scale;
+
+  //
+  // Methods
+  //
+
+  Distribution() : kind(Invalid), int_scale(0) {}
+
+  /// Configures distribution as uniform random
+  Distribution &set_uniform(double _min, double _max, int _int_scale = 0, double _pnan = 0) {
+    kind = Uniform;
+    uniform.min = _min;
+    uniform.max = _max;
+    int_scale = _int_scale;
+    uniform.pnan = _pnan;
+    return *this;
+  }
+
+  /// Configures distribution as Gaussian distribution
+  Distribution &set_gaussian(double _mean, double _stddev, int _int_scale = 0, double _pnz = 1.0) {
+    kind = Gaussian;
+    gaussian.mean = _mean;
+    gaussian.stddev = _stddev;
+    gaussian.pnz = _pnz;
+    int_scale = _int_scale;
+    return *this;
+  }
+
+  /// Sets identity
+  Distribution &set_identity() {
+    kind = Identity;
+    return *this;
+  }
+
+  /// Sets sequential
+  Distribution &set_sequential(double start, double delta, int _int_scale = 0) {
+    kind = Sequential;
+    sequential.start = start;
+    sequential.delta = delta;
+    int_scale = _int_scale;
+    return *this;
+  }
+};
+
+}  // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Prints a Distribution to ostream
+inline std::ostream &operator<<(std::ostream &out, mutlass::Distribution const &dist) {
+  switch (dist.kind) {
+    case mutlass::Distribution::Uniform:
+      out << "uniform, min: " << dist.uniform.min << ", max: " << dist.uniform.max
+          << ", pnan: " << dist.uniform.pnan;
+      break;
+    case mutlass::Distribution::Gaussian:
+      out << "gaussian, mean: " << dist.gaussian.mean << ", stddev: " << dist.gaussian.stddev
+          << ", pnzA: " << dist.gaussian.pnzA << ", pnzB: "
+          << dist.gaussian.pnzB << ", pnzC: " << dist.gaussian.pnzC;
+      break;
+    case mutlass::Distribution::Identity:
+      out << "identity";
+      break;
+    case mutlass::Distribution::Sequential:
+      out << "sequential";
+      break;
+    default:
+      out << "unknown";
+  }
+
+  out << ", int_scale: " << dist.int_scale;
+
+  return out;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/exceptions.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/exceptions.h
new file mode 100644
index 000000000..e25f6b877
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/exceptions.h
@@ -0,0 +1,70 @@
+/******************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+/**
+ * \file
+ * \brief C++ exception semantics for MUSA error codes
+ */
+
+#include <musa_runtime.h>
+#include <iosfwd>
+#include <stdexcept>
+
+#include "mutlass/platform/platform.h"
+
+namespace mutlass {
+
+/// C++ exception wrapper for MUSA \p musaError_t
+class musa_exception : public std::exception {
+ public:
+  /// Constructor
+  musa_exception(const char* msg = "", musaError_t err = musaErrorUnknown) : msg(msg), err(err) {}
+
+  /// Returns the underlying MUSA \p musaError_t
+  musaError_t musaError() const { return err; }
+
+ protected:
+  /// Explanatory string
+  const char* msg;
+
+  /// Underlying MUSA \p musaError_t
+  musaError_t err;
+};
+
+/// Writes a musa_exception instance to an output stream
+inline std::ostream& operator<<(std::ostream& out, musa_exception const& e) {
+  return out << e.what() << ": " << musaGetErrorString(e.musaError());
+}
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/gett_commandline.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/gett_commandline.hpp
new file mode 100644
index 000000000..524d7233c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/gett_commandline.hpp
@@ -0,0 +1,370 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 GETT command line parser to gather semantic modes, their stride order, and extents.
+*/
+#pragma once
+
+#include <iostream>
+#include <iomanip>
+#include <utility>
+#include <type_traits>
+#include <vector>
+#include <map>
+#include <algorithm>
+#include <numeric>
+
+#include "mutlass/util/command_line.h"
+
+namespace mutlass {
+
+// Output shortcuts
+std::ostream& operator<<(std::ostream& os, std::vector<char> data) {
+  for (auto& a : data) os << a;
+  return os;
+}
+
+template <class T>
+std::ostream& operator<<(std::ostream& os, std::vector<T> data) {
+  for (auto& a : data) os << a << " ";
+  return os;
+}
+
+struct GettCommandLine {
+  struct GettProblem {
+    using extent_type = int;
+    using stride_type = int64_t;
+
+    // Row modes: appear in A and C/D
+    std::vector<extent_type> M;
+    std::vector<stride_type> ldAm;
+    std::vector<stride_type> ldCm;
+
+    // Column modes: appear in B and C/D
+    std::vector<extent_type> N;
+    std::vector<stride_type> ldBn;
+    std::vector<stride_type> ldCn;
+
+    // Reduction modes: appear in A and B
+    std::vector<extent_type> K;
+    std::vector<stride_type> ldAk;
+    std::vector<stride_type> ldBk;
+
+    // Batch modes: appear in all in/out tensors
+    std::vector<extent_type> L;
+    std::vector<stride_type> ldAl;
+    std::vector<stride_type> ldBl;
+    std::vector<stride_type> ldCl;
+  };
+
+  static GettProblem
+  parse(int argc, char const* argv[], bool parse_verbose = false) {
+    using extent_type = typename GettProblem::extent_type;
+    using stride_type = typename GettProblem::stride_type;
+
+    mutlass::CommandLine cmd(argc, argv);
+
+    // modeA
+    std::vector<char> a_mode;
+    cmd.get_cmd_line_arguments("modeA", a_mode);
+
+    // modeB
+    std::vector<char> b_mode;
+    cmd.get_cmd_line_arguments("modeB", b_mode);
+
+    // modeC
+    std::vector<char> c_mode;
+    cmd.get_cmd_line_arguments("modeC", c_mode);
+
+
+    // mode_sizes
+    std::map<char,extent_type> mode_size;
+    // First, initialize all modes in a, b, c to make sure they're in map
+    for (char a : a_mode) mode_size[a] = 1;
+    for (char b : b_mode) mode_size[b] = 1;
+    for (char c : c_mode) mode_size[c] = 1;
+
+    // Then, overwrite the ones in -extent
+    std::vector<std::pair<std::string, std::string> > extent_tokens;
+    cmd.get_cmd_line_argument_pairs("extents", extent_tokens);
+    for (auto e : extent_tokens) {
+      if (std::get<0>(e).size() > 1) {
+        std::cerr << "ERROR: Mode name must only be 1 character long.\n";
+        print_usage();
+        exit(1);
+      }
+      char label = std::get<0>(e)[0];
+      int  size  = std::stoi(std::get<1>(e));
+      mode_size[label] = size;
+    }
+
+    // Print out symbolic modes and their extents
+    if (parse_verbose) {
+      std::cout << "C_" << c_mode << " = A_" << a_mode << " * B_" << b_mode << "\n";
+      for (auto e : mode_size) std::cout << "     " << std::get<0>(e) << " : " << std::get<1>(e) << "\n";
+    }
+
+    //
+    // Collect/Compute strides
+    //
+
+    std::map<char,stride_type> mode_ldA;
+    std::map<char,stride_type> mode_ldB;
+    std::map<char,stride_type> mode_ldC;
+
+    {
+      stride_type current;
+
+      current = 1;
+      for (char a : a_mode) { mode_ldA[a] = current; current *= mode_size[a]; }
+
+      current = 1;
+      for (char b : b_mode) { mode_ldB[b] = current; current *= mode_size[b]; }
+
+      current = 1;
+      for (char c : c_mode) { mode_ldC[c] = current; current *= mode_size[c]; }
+    }
+
+    //
+    // Collect mode categories
+    //
+
+    std::vector<char> row_mode;  // rows
+    std::vector<char> col_mode;  // columns
+    std::vector<char> red_mode;  // reductions
+    std::vector<char> bat_mode;  // batches
+
+    {
+      std::vector<char> a_label = a_mode;
+      std::vector<char> b_label = b_mode;
+      std::vector<char> c_label = c_mode;
+
+      std::sort(std::begin(a_label), std::end(a_label));
+      std::sort(std::begin(b_label), std::end(b_label));
+      std::sort(std::begin(c_label), std::end(c_label));
+
+      // std::set_intersections to find semantic category of each symbolic mode
+      std::set_intersection(std::begin(a_label), std::end(a_label),
+                            std::begin(c_label), std::end(c_label),
+                            std::back_inserter(row_mode));
+
+      std::set_intersection(std::begin(b_label), std::end(b_label),
+                            std::begin(c_label), std::end(c_label),
+                            std::back_inserter(col_mode));
+
+      std::set_intersection(std::begin(a_label), std::end(a_label),
+                            std::begin(b_label), std::end(b_label),
+                            std::back_inserter(red_mode));
+
+      std::set_intersection(std::begin(row_mode), std::end(row_mode),
+                            std::begin(col_mode), std::end(col_mode),
+                            std::back_inserter(bat_mode));
+
+      // std::set_difference to remove batch modes from other semantic modes
+      for (char l : bat_mode) {
+        row_mode.erase(std::remove(std::begin(row_mode), std::end(row_mode), l), std::end(row_mode));
+        col_mode.erase(std::remove(std::begin(col_mode), std::end(col_mode), l), std::end(col_mode));
+        red_mode.erase(std::remove(std::begin(red_mode), std::end(red_mode), l), std::end(red_mode));
+      }
+    }
+
+    // Print out the semantic association of each symbolic mode
+    if (parse_verbose) {
+      std::cout << "  rows : " << row_mode << '\n';
+      std::cout << "  cols : " << col_mode << '\n';
+      std::cout << "  reds : " << red_mode << '\n';
+      std::cout << "  bats : " << bat_mode << '\n';
+    }
+
+    //
+    // Permute modes
+    //
+
+    // Permute the batched modes to promote coalescing
+    // Sort the batched modes by min(ldAl,ldBl) and in case of a tie by the size
+    std::sort(std::begin(bat_mode), std::end(bat_mode), [&](char l1, char l2) {
+        return std::tie(std::min(mode_ldA[l1],mode_ldB[l1]),mode_size[l1])
+             < std::tie(std::min(mode_ldA[l2],mode_ldB[l2]),mode_size[l2]);
+      });
+    // Compute sizes and strides of ordered reduction modes
+    std::vector<extent_type> L;
+    std::vector<stride_type> ldAl;
+    std::vector<stride_type> ldBl;
+    std::vector<stride_type> ldCl;
+    for (char l : bat_mode) {
+      L.push_back(mode_size[l]);
+      ldAl.push_back(mode_ldA[l]);
+      ldBl.push_back(mode_ldB[l]);
+      ldCl.push_back(mode_ldC[l]);
+    }
+
+    // Permute the reduction modes to promote coalescing
+    // Sort the reduction modes by min(ldAk,ldBk) and in case of a tie by the size
+    std::sort(std::begin(red_mode), std::end(red_mode), [&](char k1, char k2) {
+        return std::tie(std::min(mode_ldA[k1],mode_ldB[k1]),mode_size[k1])
+             < std::tie(std::min(mode_ldA[k2],mode_ldB[k2]),mode_size[k2]);
+      });
+    // Compute sizes and strides of ordered reduction modes
+    std::vector<extent_type> K;
+    std::vector<stride_type> ldAk;
+    std::vector<stride_type> ldBk;
+    for (char k : red_mode) {
+      K.push_back(mode_size[k]);
+      ldAk.push_back(mode_ldA[k]);
+      ldBk.push_back(mode_ldB[k]);
+    }
+
+    // Permute the row modes to promote coalescing
+    // Sort the row modes by min(ldAm,ldCm) and in case of a tie by ldAm
+    std::sort(std::begin(row_mode), std::end(row_mode), [&](char m1, char m2) {
+        return std::tie(std::min(mode_ldA[m1],mode_ldC[m1]),mode_ldA[m1])
+             < std::tie(std::min(mode_ldA[m2],mode_ldC[m2]),mode_ldA[m2]);
+      });
+    // Compute sizes and strides of ordered row modes
+    std::vector<extent_type> M;
+    std::vector<stride_type> ldAm;
+    std::vector<stride_type> ldCm;
+    for (char m : row_mode) {
+      M.push_back(mode_size[m]);
+      ldAm.push_back(mode_ldA[m]);
+      ldCm.push_back(mode_ldC[m]);
+    }
+
+    // Permute the col modes to promote coalescing
+    // Sort the col modes by min(ldBn,ldCn) and in case of a tie by ldBn
+    std::sort(std::begin(col_mode), std::end(col_mode), [&](char n1, char n2) {
+        return std::tie(std::min(mode_ldB[n1],mode_ldC[n1]),mode_ldB[n1])
+             < std::tie(std::min(mode_ldB[n2],mode_ldC[n2]),mode_ldB[n2]);
+      });
+    // Compute sizes and strides of ordered col modes
+    std::vector<extent_type> N;
+    std::vector<stride_type> ldBn;
+    std::vector<stride_type> ldCn;
+    for (char n : col_mode) {
+      N.push_back(mode_size[n]);
+      ldBn.push_back(mode_ldB[n]);
+      ldCn.push_back(mode_ldC[n]);
+    }
+
+    if (parse_verbose) {
+      std::cout << "C_";
+      if (! row_mode.empty()) {
+        std::cout << "(" << row_mode << ")";
+      }
+      if (! col_mode.empty()) {
+        std::cout << "(" << col_mode << ")";
+      }
+      if (! bat_mode.empty()) {
+        std::cout << "(" << bat_mode << ")";
+      }
+      std::cout << " = A_";
+      if (! row_mode.empty()) {
+        std::cout << "(" << row_mode << ")";
+      }
+      if (! red_mode.empty()) {
+        std::cout << "(" << red_mode << ")";
+      }
+      if (! bat_mode.empty()) {
+        std::cout << "(" << bat_mode << ")";
+      }
+      std::cout << " * B_";
+      if (! col_mode.empty()) {
+        std::cout << "(" << col_mode << ")";
+      }
+      if (! red_mode.empty()) {
+        std::cout << "(" << red_mode << ")";
+      }
+      if (! bat_mode.empty()) {
+        std::cout << "(" << bat_mode << ")";
+      }
+      std::cout << '\n';
+
+      int M_size = std::accumulate(std::begin(M), std::end(M), 1, std::multiplies<>{});
+      int N_size = std::accumulate(std::begin(N), std::end(N), 1, std::multiplies<>{});
+      int K_size = std::accumulate(std::begin(K), std::end(K), 1, std::multiplies<>{});
+      int L_size = std::accumulate(std::begin(L), std::end(L), 1, std::multiplies<>{});
+
+      std::cout << "     M : (" << M_size << ") ";
+      for (char m : row_mode) std::cout << m << ":" << mode_size[m] << " ";
+      std::cout << '\n';
+      std::cout << "     N : (" << N_size << ") ";
+      for (char n : col_mode) std::cout << n << ":" << mode_size[n] << " ";
+      std::cout << '\n';
+      std::cout << "     K : (" << K_size << ") ";
+      for (char k : red_mode) std::cout << k << ":" << mode_size[k] << " ";
+      std::cout << '\n';
+      std::cout << "     L : (" << L_size << ") ";
+      for (char l : bat_mode) std::cout << l << ":" << mode_size[l] << " ";
+      std::cout << '\n';
+
+      std::cout << "  ldAm : " << ldAm << '\n';
+      std::cout << "  ldAk : " << ldAk << '\n';
+      std::cout << "  ldAl : " << ldAl << '\n';
+      std::cout << "  ldBn : " << ldBn << '\n';
+      std::cout << "  ldBk : " << ldBk << '\n';
+      std::cout << "  ldBl : " << ldBl << '\n';
+      std::cout << "  ldCm : " << ldCm << '\n';
+      std::cout << "  ldCn : " << ldCn << '\n';
+      std::cout << "  ldCl : " << ldCl << '\n';
+    }
+
+    return {M, ldAm, ldCm,
+            N, ldBn, ldCn,
+            K, ldAk, ldBk,
+            L, ldAl, ldBl, ldCl};
+  }
+
+  static void
+  print_usage() {
+    std::cout <<
+      "GETT problem command line parser:\n"
+      "  --modeA=<m0,...>\n"
+      "    A comma delimited list of characters that correspond to the row, reduction, and batch modes in A tensor.\n"
+      "    The semantic association of each symbolic mode is determined automatically.\n\n"
+
+      "  --modeB=<m0,...>\n"
+      "    A comma delimited list of characters that correspond to the column, reduction, and batch modes in B tensor.\n"
+      "    The semantic association of each symbolic mode is determined automatically.\n\n"
+
+      "  --modeC=<m0,...>\n"
+      "    A comma delimited list of characters that correspond to the row, column, and batch modes in B tensor.\n"
+      "    The semantic association of each symbolic mode is determined automatically.\n\n"
+
+      "  --extents=<mode:extent,....>\n"
+      "    A command delimited list of symbolic mode and its corresponding extent.\n"
+      "    Extents are defaulted to 1 if any are not provided.\n\n"
+
+      "Example usage: gett.exe --modeC=m,n,l --modeA=m,k,l --modeB=k,n,l --extents=m:4096,n:4096,k:4096\n";
+  }
+};
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/helper_cuda.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/helper_cuda.hpp
new file mode 100644
index 000000000..727ffe3f1
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/helper_cuda.hpp
@@ -0,0 +1,107 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <musa.h>
+
+#include <mute/util/debug.hpp>
+
+namespace mute
+{
+
+void
+device_init(int device_id, bool quiet = false)
+{
+  musaDeviceProp device_prop;
+  std::size_t    device_free_physmem;
+  std::size_t    device_total_physmem;
+
+  MUTE_CHECK_ERROR(musaSetDevice(device_id));
+  MUTE_CHECK_ERROR(musaMemGetInfo(&device_free_physmem, &device_total_physmem));
+  MUTE_CHECK_ERROR(musaGetDeviceProperties(&device_prop, device_id));
+
+  if (device_prop.major < 1) {
+    fprintf(stderr, "Device does not support MUSA.\n");
+    exit(1);
+  }
+
+  //float device_giga_bandwidth = float(device_prop.memoryBusWidth) * device_prop.memoryClockRate * 2 / 8 / 1000 / 1000;
+
+  if (!quiet) {
+    printf("Using device %d: %s  (SM%d, %d SMs)\n",
+           device_id, device_prop.name,
+           device_prop.major * 10 + device_prop.minor,
+           device_prop.multiProcessorCount);
+    fflush(stdout);
+  }
+}
+
+/**
+ * Convert the SM version (e.g. v7.0, v7.5) to the physical number of cores.
+ */
+inline int
+_ConvertSMVer2Cores(int major, int minor)
+{
+  // Defines for GPU Architecture types (using the SM version to determine
+  // the # of cores per SM
+  typedef struct {
+    int SM;  // 0xMm (hexadecimal notation), M = SM Major version,
+    // and m = SM minor version
+    int Cores;
+  } sSMtoCores;
+
+  sSMtoCores nGpuArchCoresPerSM[] = {
+      {0x10, 128},
+      {0x21, 128},
+      {0x22, 128},
+      {-1, -1}};
+
+  int index = 0;
+
+  while (nGpuArchCoresPerSM[index].SM != -1) {
+    if (nGpuArchCoresPerSM[index].SM == ((major << 4) + minor)) {
+      return nGpuArchCoresPerSM[index].Cores;
+    }
+    index++;
+  }
+
+  // If we don't find the values, we default use the previous one
+  // to run properly
+  printf("MapSMtoCores for SM %d.%d is undefined."
+         "  Default to use %d Cores/SM\n",
+         major, minor, nGpuArchCoresPerSM[index - 1].Cores);
+
+  return nGpuArchCoresPerSM[index - 1].Cores;
+}
+
+} // end namespace mute
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/host_reorder.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/host_reorder.h
new file mode 100644
index 000000000..f6960e110
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/host_reorder.h
@@ -0,0 +1,112 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 reorder data from the host side
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/util/host_tensor.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/util/tensor_view_io.h"
+#include "mutlass/util/reference/host/gemm.h"
+
+namespace mutlass {
+
+/// This is needed for the interleaved integer tensor core kernels.  The purpose
+/// is to use skip the shared memory part in the epilogue.
+template <int Interleaved, typename Element, typename Layout>
+void reorder_column(TensorRef<Element, Layout> dest,
+                    TensorRef<Element, Layout> src,
+                    mutlass::gemm::GemmCoord problem_size) {
+  const int InstructionShapeCol = 8;
+  // 4 threads per Quad
+  const int ElementsPerThread = InstructionShapeCol / 4;
+  // 4 threads per Quad
+  const int ReorderedElementsPerThread =
+      Interleaved / 4;
+
+  for (int n = 0; n < problem_size.n(); n++) {
+    for (int k = 0; k < problem_size.k(); k++) {
+      dest.at({k, (n / Interleaved) * Interleaved +
+                      ((n % ReorderedElementsPerThread) / ElementsPerThread) *
+                          InstructionShapeCol +
+                      ((n % Interleaved) / ReorderedElementsPerThread) *
+                          ElementsPerThread +
+                      (n % ElementsPerThread)}) = src.at({k, n});
+    }
+  }
+}
+
+template <int ColumnInterleaved, int LayoutInterleaved = ColumnInterleaved, typename Element, typename Layout>
+void reorder_convK(TensorRef<Element, Layout> dest,
+                    TensorRef<Element, Layout> src,
+                    mutlass::gemm::GemmCoord problem_size) {
+
+    TensorRef<Element, layout::RowMajorInterleaved<LayoutInterleaved>> mappedDest(dest.data(), dest.stride(0));
+    TensorRef<Element, layout::RowMajorInterleaved<LayoutInterleaved>> mappedSrc(src.data(), src.stride(0));
+
+    reorder_column<ColumnInterleaved>(
+        mappedDest, mappedSrc, problem_size);
+}
+
+/// This is needed for the sparse tensor core kernels.  The purpose
+/// is to use ldmatrix to load from shared memory to the register file.
+template <typename Element, typename LayoutDest, typename LayoutSrc>
+void reorder_meta(TensorRef<Element, LayoutDest> dest,
+                  TensorRef<Element, LayoutSrc> src,
+                  mutlass::gemm::GemmCoord problem_size) {
+  for (int m = 0; m < problem_size.m(); m++) {
+    for (int k = 0; k < problem_size.k(); k++) {
+      // First reorder the rows.
+      int group = (sizeof(Element) == 2) ? 32 : 16;
+      int interweave = (sizeof(Element) == 2) ? 4 : 2;
+
+      int dest_row = m / group * group + (m % 8) * interweave + (m % group) / 8;
+      int dest_col = k;
+
+      // Next swizzle the 2x2 blocks from Z to N.
+      if (((dest_row % 2) == 0) && ((dest_col % 2) == 1)) {
+        ++dest_row;
+        --dest_col;
+      } else if (((dest_row % 2) == 1) && ((dest_col % 2) == 0)) {
+        --dest_row;
+        ++dest_col;
+      }
+
+      dest.at({dest_row, dest_col}) = src.at({m, k});
+    }
+  }
+}
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/host_tensor.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/host_tensor.h
new file mode 100644
index 000000000..830fc56aa
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/host_tensor.h
@@ -0,0 +1,524 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+/*! \file
+  \brief HostTensor contributes management for both host and device memory.
+
+  HostTensor allocates host and device memory upon construction. Basic element-wise operations on
+  host memory synchronize device memory automatically. Explicit copy operations provide abstractions
+  for MUSA memcpy operations.
+
+  Call {host, device}_{data, ref, view}() for accessing host or device memory.
+
+  See mutlass/tensor_ref.h and mutlass/tensor_view.h for more details.
+*/
+
+#include <vector>
+
+#include "mutlass/mutlass.h"
+#include "mutlass/tensor_ref.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/fast_math.h"
+
+#include "device_memory.h"
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Host tensor
+template <
+  /// Data type of element stored within tensor (concept: NumericType)
+  typename Element_,
+  /// Defines a mapping from logical coordinate to linear memory (concept: Layout)
+  typename Layout_
+>
+class HostTensor {
+public:
+
+  /// Data type of individual access
+  using Element = Element_;
+
+  /// Mapping function from logical coordinate to linear memory
+  using Layout = Layout_;
+
+  /// Logical rank of tensor index space
+  static int const kRank = Layout::kRank;
+
+  /// Index type
+  using Index = typename Layout::Index;
+
+  /// Long index used for pointer offsets
+  using LongIndex = typename Layout::LongIndex;
+
+  /// Coordinate in logical tensor space
+  using TensorCoord = typename Layout::TensorCoord;
+
+  /// Layout's stride vector
+  using Stride = typename Layout::Stride;
+
+  /// Tensor reference to device memory
+  using TensorRef = TensorRef<Element, Layout>;
+
+  /// Tensor reference to constant device memory
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  /// Tensor reference to device memory
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Tensor reference to constant device memory
+  using ConstTensorView = typename TensorView::ConstTensorView;
+
+  /// Reference to element in tensor
+  using Reference = typename TensorRef::Reference;
+
+  /// Constant reference to element in tensor
+  using ConstReference = typename ConstTensorRef::Reference;
+
+  /// Note: Below is used to handle packing of subbyte elements
+  /// kBitsStoredVec          : The bits of store vec that could be divisiable by the element
+  /// kElementsPerStoredVec   : The number of elements could be stored in per store vec
+  /// kNumStoragePerStoredVec : How much storage(i.e. sizeof(element storage)) the store vec needs to consume.
+  ///                           Usually the element storage of subbyte is uint8_t.
+  /// Example
+  ///  int2:  kBitsStoredVec = 8; kElementsPerStoredVec = 4; kNumStoragePerStoredVec = 1 uint8_t;
+  ///  int4:  kBitsStoredVec = 8; kElementsPerStoredVec = 2; kNumStoragePerStoredVec = 1 uint8_t;
+  static int const kBitsStoredVec        = (sizeof_bits<Element>::value < 8) ? mutlass::lcm(static_cast<int>(sizeof_bits<Element>::value), 8) : sizeof_bits<Element>::value;
+  static int const kElementsPerStoredVec = kBitsStoredVec / sizeof_bits<Element>::value;
+  static int const kNumStoragePerStoredVec = kBitsStoredVec / (sizeof(Element) * 8);
+
+ private:
+
+  //
+  // Data members
+  //
+
+  /// Extent of tensor in logical dimensions
+  TensorCoord extent_;
+
+  /// Layout object
+  Layout layout_;
+
+  /// Host-side memory allocation
+  /// avoid the std::vector<bool> specialization
+  std::vector<std::conditional_t<std::is_same_v<Element,bool>, uint8_t, Element>> host_;
+
+  /// Device-side memory
+  device_memory::allocation<Element> device_;
+
+ public:
+  //
+  // Device and Host Methods
+  //
+
+  /// Default constructor
+  HostTensor() {}
+
+  /// Constructs a tensor given an extent. Assumes a packed layout
+  HostTensor(
+    TensorCoord const &extent,
+    bool device_backed = true
+  ) {
+
+    this->reset(extent, Layout::packed(extent), device_backed);
+  }
+
+  /// Constructs a tensor given an extent and layout
+  HostTensor(
+    TensorCoord const &extent,
+    Layout const &layout,
+    bool device_backed = true
+  ) {
+
+    this->reset(extent, layout, device_backed);
+  }
+
+  ~HostTensor() { }
+
+  /// Clears the HostTensor allocation to size/capacity = 0
+  void reset() {
+    extent_ = TensorCoord();
+    layout_ = Layout::packed(extent_);
+
+    host_.clear();
+    device_.reset();
+  }
+
+  /// Resizes internal memory allocations without affecting layout or extent
+  void reserve(
+    size_t count,                                        ///< size of tensor in elements
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated
+
+    device_.reset();
+    host_.clear();
+
+    count = count / kElementsPerStoredVec * kNumStoragePerStoredVec;
+    host_.resize(count);
+
+    // Allocate memory
+    Element* device_memory = nullptr;
+    if (device_backed_) {
+      device_memory = device_memory::allocate<Element>(count);
+    }
+    device_.reset(device_memory, device_backed_ ? count : 0);
+  }
+
+  /// Updates the extent and layout of the HostTensor. Allocates memory according to the new
+  /// extent and layout.
+  void reset(
+    TensorCoord const &extent,                           ///< extent of logical tensor
+    Layout const &layout,                                ///< layout object of tensor
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated.
+
+    extent_ = extent;
+    layout_ = layout;
+
+    reserve(size_t(layout_.capacity(extent_)), device_backed_);
+  }
+
+  /// Updates the extent and layout of the HostTensor. Allocates memory according to the new
+  /// extent and layout. Assumes a packed tensor configuration.
+  void reset(
+    TensorCoord const &extent,                           ///< extent of logical tensor
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated.
+
+    reset(extent, Layout::packed(extent), device_backed_);
+  }
+
+  /// Changes the size of the logical tensor. Only allocates memory if new capacity exceeds reserved capacity.
+  /// To force allocation, call reset().
+  void resize(
+    TensorCoord const &extent,                           ///< extent of logical tensor
+    Layout const &layout,                                ///< layout object of tensor
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated.
+
+    extent_ = extent;
+    layout_ = layout;
+
+    LongIndex new_size = size_t(layout_.capacity(extent_));
+
+    if (static_cast<decltype(host_.size())>(new_size) > host_.size()) {
+      reserve(new_size, device_backed_);
+    }
+  }
+
+  /// Changes the size of the logical tensor. Only allocates memory if new capacity exceeds reserved capacity.
+  /// To force allocation, call reset(). Note, this form of resize() assumes a packed tensor configuration.
+  void resize(
+    TensorCoord const &extent,                           ///< extent of logical tensor
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated.
+
+    resize(extent, Layout::packed(extent), device_backed_);
+  }
+
+  /// Returns the number of elements stored in the host tensor
+  size_t size() const {
+    return host_.size() / kNumStoragePerStoredVec * kElementsPerStoredVec;
+  }
+
+  /// Returns the logical capacity based on extent and layout. May differ from size().
+  LongIndex capacity() const {
+    return layout_.capacity(extent_);
+  }
+
+  /// Gets pointer to host data
+  Element * host_data() { return reinterpret_cast<Element *>(host_.data()); }
+
+  /// Gets pointer to host data with a pointer offset
+  Element * host_data_ptr_offset(LongIndex ptr_element_offset) { return &ReferenceFactory<Element>::get(host_data(), ptr_element_offset); }
+
+  /// Gets a reference to an element in host memory
+  Reference host_data(LongIndex idx) {
+    return ReferenceFactory<Element>::get(host_data(), idx);
+  }
+
+  /// Gets pointer to host data
+  Element const * host_data() const { return reinterpret_cast<Element const *>(host_.data()); }
+
+  /// Gets pointer to host data with a pointer offset
+  Element const * host_data_ptr_offset(LongIndex ptr_element_offset) const { return &ReferenceFactory<Element>::get(host_data(), ptr_element_offset); }
+
+  /// Gets a constant reference to an element in host memory
+  ConstReference host_data(LongIndex idx) const {
+    return ReferenceFactory<Element const>::get(host_data(), idx);
+  }
+
+  /// Gets pointer to device data
+  Element * device_data() { return device_.get(); }
+
+  /// Gets pointer to device data
+  Element const * device_data() const { return device_.get(); }
+
+  /// Gets pointer to device data with a pointer offset
+  Element * device_data_ptr_offset(LongIndex ptr_element_offset) { return &ReferenceFactory<Element>::get(device_data(), ptr_element_offset); }
+
+  /// Gets pointer to device data with a pointer offset
+  Element const * device_data_ptr_offset(LongIndex ptr_element_offset) const { return &ReferenceFactory<Element>::get(device_data(), ptr_element_offset); }
+
+  /// Accesses the tensor reference pointing to data
+  TensorRef host_ref(LongIndex ptr_element_offset=0) { return TensorRef(host_data_ptr_offset(ptr_element_offset), layout_); }
+
+  /// Accesses the tensor reference pointing to data
+  ConstTensorRef host_ref(LongIndex ptr_element_offset=0) const { return ConstTensorRef(host_data_ptr_offset(ptr_element_offset), layout_); }
+
+  /// Accesses the tensor reference pointing to data
+  TensorRef device_ref(LongIndex ptr_element_offset=0) {
+    return TensorRef(device_data_ptr_offset(ptr_element_offset), layout_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  ConstTensorRef device_ref(LongIndex ptr_element_offset=0) const {
+    return TensorRef(device_data_ptr_offset(ptr_element_offset), layout_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  TensorView host_view(LongIndex ptr_element_offset=0) {
+    return TensorView(host_data_ptr_offset(ptr_element_offset), layout_, extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  ConstTensorView host_view(LongIndex ptr_element_offset=0) const {
+    return ConstTensorView(host_data_ptr_offset(ptr_element_offset), layout_, extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  TensorView device_view(LongIndex ptr_element_offset=0) {
+    return TensorView(device_data_ptr_offset(ptr_element_offset), layout_, extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  ConstTensorView device_view(LongIndex ptr_element_offset=0) const {
+    return ConstTensorView(device_data_ptr_offset(ptr_element_offset), layout_, extent_);
+  }
+
+  /// Returns true if device memory is allocated
+  bool device_backed() const {
+    return (device_.get() == nullptr) ? false : true;
+  }
+
+
+  /// Returns the layout object
+  Layout & layout() {
+    return layout_;
+  }
+
+  /// Returns the layout object
+  Layout layout() const {
+    return layout_;
+  }
+
+  /// Returns the layout object's stride vector
+  Stride stride() const {
+    return layout_.stride();
+  }
+
+  /// Returns the layout object's stride vector
+  Stride & stride() {
+    return layout_.stride();
+  }
+
+  /// Returns the layout object's stride in a given physical dimension
+  LongIndex stride(int dim) const {
+    return layout_.stride().at(dim);
+  }
+
+  /// Returns the layout object's stride in a given physical dimension
+  LongIndex & stride(int dim) {
+    return layout_.stride().at(dim);
+  }
+
+  /// Computes the offset of an index from the origin of the tensor
+  LongIndex offset(TensorCoord const& coord) const {
+    return layout_(coord);
+  }
+
+  /// Returns a reference to the element at the logical Coord in host memory
+  Reference at(TensorCoord const& coord) {
+    return host_data(offset(coord));
+  }
+
+  /// Returns a const reference to the element at the logical Coord in host memory
+  ConstReference at(TensorCoord const& coord) const {
+    return host_data(offset(coord));
+  }
+
+  /// Returns the extent of the tensor
+  TensorCoord extent() const {
+    return extent_;
+  }
+
+  /// Returns the extent of the tensor
+  TensorCoord & extent() {
+    return extent_;
+  }
+
+  /// Copies data from device to host
+  void sync_host() {
+    if (device_backed()) {
+      device_memory::copy_to_host(
+          host_data(), device_data(), size());
+    }
+  }
+
+  /// Copies data from host to device
+  void sync_device() {
+    if (device_backed()) {
+      device_memory::copy_to_device(
+          device_data(), host_data(), size());
+    }
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_in_device_to_host(
+    Element const* ptr_device,        ///< source device memory
+    LongIndex count = -1) {           ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __NV_STD_MIN(capacity(), count);
+    }
+    device_memory::copy_to_host(
+      host_data(), ptr_device, count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_in_device_to_device(
+    Element const* ptr_device,        ///< source device memory
+    LongIndex count = -1) {           ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __NV_STD_MIN(capacity(), count);
+    }
+    device_memory::copy_device_to_device(
+      device_data(), ptr_device, count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_in_host_to_device(
+    Element const* ptr_host,          ///< source host memory
+    LongIndex count = -1) {           ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __NV_STD_MIN(capacity(), count);
+    }
+    device_memory::copy_to_device(
+      device_data(), ptr_host, count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_in_host_to_host(
+    Element const* ptr_host,          ///< source host memory
+    LongIndex count = -1) {           ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __NV_STD_MIN(capacity(), count);
+    }
+    device_memory::copy_host_to_host(
+      host_data(), ptr_host, count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_out_device_to_host(
+    Element * ptr_host,               ///< source device memory
+    LongIndex count = -1) const {     ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __NV_STD_MIN(capacity(), count);
+    }
+    device_memory::copy_to_host(
+      ptr_host, device_data(), count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_out_device_to_device(
+    Element * ptr_device,             ///< source device memory
+    LongIndex count = -1) const {     ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __NV_STD_MIN(capacity(), count);
+    }
+    device_memory::copy_device_to_device(
+      ptr_device, device_data(), count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_out_host_to_device(
+    Element * ptr_device,             ///< source host memory
+    LongIndex count = -1) const {     ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __NV_STD_MIN(capacity(), count);
+    }
+    device_memory::copy_to_device(
+      ptr_device, host_data(), count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_out_host_to_host(
+    Element * ptr_host,               ///< source host memory
+    LongIndex count = -1) const {     ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __NV_STD_MIN(capacity(), count);
+    }
+    device_memory::copy_host_to_host(
+      ptr_host, host_data(), count);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/host_tensor_planar_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/host_tensor_planar_complex.h
new file mode 100644
index 000000000..7b22c00ff
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/host_tensor_planar_complex.h
@@ -0,0 +1,592 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+/*! \file
+  \brief HostTensor contributes management for both host and device memory.
+
+  HostTensor allocates host and device memory upon construction. Basic element-wise operations on
+  host memory synchronize device memory automatically. Explicit copy operations provide abstractions
+  for MUSA memcpy operations.
+
+  Call {host, device}_{data, ref, view}() for accessing host or device memory.
+
+  See mutlass/tensor_ref.h and mutlass/tensor_view.h for more details.
+*/
+
+#include <vector>
+
+#include "mutlass/mutlass.h"
+
+#include "mutlass/tensor_ref_planar_complex.h"
+#include "mutlass/tensor_view_planar_complex.h"
+
+#include "device_memory.h"
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Host tensor
+template <
+  /// Data type of element stored within tensor (concept: NumericType)
+  typename Element_,
+  /// Defines a mapping from logical coordinate to linear memory (concept: Layout)
+  typename Layout_
+>
+class HostTensorPlanarComplex {
+public:
+
+  /// Data type of individual access
+  using Element = Element_;
+
+  /// Mapping function from logical coordinate to linear memory
+  using Layout = Layout_;
+
+  /// Logical rank of tensor index space
+  static int const kRank = Layout::kRank;
+
+  /// Index type
+  using Index = typename Layout::Index;
+
+  /// Long index used for pointer offsets
+  using LongIndex = typename Layout::LongIndex;
+
+  /// Coordinate in logical tensor space
+  using TensorCoord = typename Layout::TensorCoord;
+
+  /// Layout's stride vector
+  using Stride = typename Layout::Stride;
+
+  /// Tensor reference to device memory
+  using TensorRef = TensorRefPlanarComplex<Element, Layout>;
+
+  /// Tensor reference to constant device memory
+  using ConstTensorRef = typename TensorRef::ConstTensorRef;
+
+  /// Tensor reference to device memory
+  using TensorView = TensorViewPlanarComplex<Element, Layout>;
+
+  /// Tensor reference to constant device memory
+  using ConstTensorView = typename TensorView::ConstTensorView;
+
+  /// Reference to element in tensor
+  using Reference = typename TensorRef::Reference;
+
+  /// Constant reference to element in tensor
+  using ConstReference = typename ConstTensorRef::Reference;
+
+ private:
+
+  //
+  // Data members
+  //
+
+  /// Extent of tensor in logical dimensions
+  TensorCoord extent_;
+
+  /// Layout object
+  Layout layout_;
+
+  /// Host-side memory allocation
+  std::vector<Element> host_;
+
+  /// Device-side memory
+  device_memory::allocation<Element> device_;
+
+ public:
+  //
+  // Device and Host Methods
+  //
+
+  /// Default constructor
+  HostTensorPlanarComplex() {}
+
+  /// Constructs a tensor given an extent. Assumes a packed layout
+  HostTensorPlanarComplex(
+    TensorCoord const &extent,
+    bool device_backed = true
+  ) {
+
+    this->reset(extent, Layout::packed(extent), device_backed);
+  }
+
+  /// Constructs a tensor given an extent and layout
+  HostTensorPlanarComplex(
+    TensorCoord const &extent,
+    Layout const &layout,
+    bool device_backed = true
+  ) {
+
+    this->reset(extent, layout, device_backed);
+  }
+
+  ~HostTensorPlanarComplex() { }
+
+  /// Clears the HostTensor allocation to size/capacity = 0
+  void reset() {
+    extent_ = TensorCoord();
+    layout_ = Layout::packed(extent_);
+
+    host_.clear();
+    device_.reset();
+  }
+
+  /// Resizes internal memory allocations without affecting layout or extent
+  void reserve(
+    size_t count,                                        ///< size of tensor in elements
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated
+
+    device_.reset();
+    host_.clear();
+
+    host_.resize(count * 2);
+
+    // Allocate memory
+    Element* device_memory = nullptr;
+    if (device_backed_) {
+      device_memory = device_memory::allocate<Element>(count * 2);
+    }
+    device_.reset(device_memory, device_backed_ ? count * 2 : 0);
+  }
+
+  /// Updates the extent and layout of the HostTensor. Allocates memory according to the new
+  /// extent and layout.
+  void reset(
+    TensorCoord const &extent,                           ///< extent of logical tensor
+    Layout const &layout,                                ///< layout object of tensor
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated.
+
+    extent_ = extent;
+    layout_ = layout;
+
+    reserve(size_t(layout_.capacity(extent_)), device_backed_);
+  }
+
+  /// Updates the extent and layout of the HostTensor. Allocates memory according to the new
+  /// extent and layout. Assumes a packed tensor configuration.
+  void reset(
+    TensorCoord const &extent,                           ///< extent of logical tensor
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated.
+
+    reset(extent, Layout::packed(extent), device_backed_);
+  }
+
+  /// Changes the size of the logical tensor. Only allocates memory if new capacity exceeds reserved capacity.
+  /// To force allocation, call reset().
+  void resize(
+    TensorCoord const &extent,                           ///< extent of logical tensor
+    Layout const &layout,                                ///< layout object of tensor
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated.
+
+    extent_ = extent;
+    layout_ = layout;
+
+    LongIndex new_size = size_t(layout_.capacity(extent_));
+
+    if (static_cast<decltype(host_.size())>(new_size * 2) > host_.size()) {
+      reserve(new_size);
+    }
+  }
+
+  /// Changes the size of the logical tensor. Only allocates memory if new capacity exceeds reserved capacity.
+  /// To force allocation, call reset(). Note, this form of resize() assumes a packed tensor configuration.
+  void resize(
+    TensorCoord const &extent,                           ///< extent of logical tensor
+    bool device_backed_ = true) {                        ///< if true, device memory is also allocated.
+
+    resize(extent, Layout::packed(extent), device_backed_);
+  }
+
+  /// Returns the number of elements stored in the host tensor
+  size_t size() const {
+    return host_.size() / 2;
+  }
+
+  /// Returns the logical capacity based on extent and layout. May differ from size().
+  LongIndex capacity() const {
+    return layout_.capacity(extent_);
+  }
+
+  /// Stride between real and imaginary parts
+  LongIndex imaginary_stride() const {
+    return host_.size() / 2;
+  }
+
+  /// Gets pointer to host data
+  Element * host_data() { return host_.data(); }
+
+  /// Gets pointer to host data imaginary part
+  Element * host_data_imag() { return host_.data() + imaginary_stride(); }
+
+  /// Gets pointer to host data with a pointer offset
+  Element * host_data_ptr_offset(LongIndex ptr_element_offset) { return host_data() + ptr_element_offset; }
+
+  /// Gets pointer to host data with a pointer offset
+  Element * host_data_imag_ptr_offset(LongIndex ptr_element_offset) { return host_data_imag() + ptr_element_offset; }
+
+  /// Gets a reference to an element in host memory
+  Reference host_data(LongIndex idx) {
+    return PlanarComplexReference<Element>(host_data() + idx, host_data_imag() + idx);
+  }
+
+  /// Gets pointer to host data
+  Element const * host_data() const { return host_.data(); }
+
+  /// Gets pointer to host data imaginary part
+  Element const * host_data_imag() const { return host_.data() + imaginary_stride(); }
+
+  /// Gets a constant reference to an element in host memory
+  ConstReference host_data(LongIndex idx) const {
+    return PlanarComplexReference<Element const>(host_data() + idx, host_data_imag() + idx);
+  }
+
+  /// Gets pointer to device data
+  Element * device_data() { return device_.get(); }
+
+  /// Gets pointer to device data with a pointer offset
+  Element * device_data_ptr_offset(LongIndex ptr_element_offset) { return device_.get() + ptr_element_offset; }
+
+  /// Gets pointer to device data
+  Element const * device_data() const { return device_.get(); }
+
+  /// Gets pointer to device data with a pointer offset
+  Element const * device_data_ptr_offset(LongIndex ptr_element_offset) const { return device_.get() + ptr_element_offset; }
+
+  /// Gets a pointer to the device data imaginary part
+  Element * device_data_imag() { return device_.get() + imaginary_stride(); }
+
+  /// Accesses the tensor reference pointing to data
+  TensorRef host_ref(LongIndex ptr_element_offset=0) {
+    return TensorRef(host_data_ptr_offset(ptr_element_offset), layout_, imaginary_stride());
+  }
+
+  /// Returns a tensor reference to the real part of the tensor
+  mutlass::TensorRef<Element, Layout> host_ref_real() {
+    return mutlass::TensorRef<Element, Layout>(host_data(), layout_);
+  }
+
+  /// Returns a tensor reference to the real part of the tensor
+  mutlass::TensorRef<Element, Layout> host_ref_imag() {
+    return mutlass::TensorRef<Element, Layout>(host_data_ptr_offset(imaginary_stride()), layout_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  ConstTensorRef host_ref(LongIndex ptr_element_offset=0) const {
+    return ConstTensorRef(host_data_ptr_offset(ptr_element_offset), layout_, imaginary_stride());
+  }
+
+  /// Accesses the tensor reference pointing to data
+  TensorRef device_ref(LongIndex ptr_element_offset=0) {
+    return TensorRef(device_data_ptr_offset(ptr_element_offset), layout_, imaginary_stride());
+  }
+
+  /// Accesses the tensor reference pointing to data
+  ConstTensorRef device_ref(LongIndex ptr_element_offset=0) const {
+    return TensorRef(device_data_ptr_offset(ptr_element_offset), layout_, imaginary_stride());
+  }
+
+  /// Returns a tensor reference to the real part of the tensor
+  mutlass::TensorRef<Element, Layout> device_ref_real() {
+    return mutlass::TensorRef<Element, Layout>(device_data(), layout_);
+  }
+
+  /// Returns a tensor reference to the real part of the tensor
+  mutlass::TensorRef<Element, Layout> device_ref_imag() {
+    return mutlass::TensorRef<Element, Layout>(device_data_ptr_offset(imaginary_stride()), layout_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  TensorView host_view(LongIndex ptr_element_offset=0) {
+    return TensorView(host_data_ptr_offset(ptr_element_offset), layout_, imaginary_stride(), extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  ConstTensorView host_view(LongIndex ptr_element_offset=0) const {
+    return ConstTensorView(host_data_ptr_offset(ptr_element_offset), layout_, imaginary_stride(), extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  mutlass::TensorView<Element, Layout> host_view_real() {
+    return mutlass::TensorView<Element, Layout>(host_data(), layout_, extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  mutlass::TensorView<Element, Layout> host_view_imag() {
+    return mutlass::TensorView<Element, Layout>(host_data_ptr_offset(imaginary_stride()), layout_, extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  TensorView device_view(LongIndex ptr_element_offset=0) {
+    return TensorView(device_data_ptr_offset(ptr_element_offset), layout_, imaginary_stride(), extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  ConstTensorView device_view(LongIndex ptr_element_offset=0) const {
+    return ConstTensorView(device_data_ptr_offset(ptr_element_offset), layout_, imaginary_stride(), extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  mutlass::TensorView<Element, Layout> device_view_real() {
+    return mutlass::TensorView<Element, Layout>(device_data(), layout_, extent_);
+  }
+
+  /// Accesses the tensor reference pointing to data
+  mutlass::TensorView<Element, Layout> device_view_imag() {
+    return mutlass::TensorView<Element, Layout>(device_data_ptr_offset(imaginary_stride()), layout_, extent_);
+  }
+
+  /// Returns true if device memory is allocated
+  bool device_backed() const {
+    return (device_.get() == nullptr) ? false : true;
+  }
+
+  /// Returns the layout object
+  Layout layout() const {
+    return layout_;
+  }
+
+  /// Returns the layout object's stride vector
+  Stride stride() const {
+    return layout_.stride();
+  }
+
+  /// Returns the layout object's stride in a given physical dimension
+  Index stride(int dim) const {
+    return layout_.stride().at(dim);
+  }
+
+  /// Computes the offset of an index from the origin of the tensor
+  LongIndex offset(TensorCoord const& coord) const {
+    return layout_(coord);
+  }
+
+  /// Returns a reference to the element at the logical Coord in host memory
+  Reference at(TensorCoord const& coord) {
+    return host_data(offset(coord));
+  }
+
+  /// Returns a const reference to the element at the logical Coord in host memory
+  ConstReference at(TensorCoord const& coord) const {
+    return host_data(offset(coord));
+  }
+
+  /// Returns the extent of the tensor
+  TensorCoord extent() const {
+    return extent_;
+  }
+
+  /// Returns the extent of the tensor
+  TensorCoord & extent() {
+    return extent_;
+  }
+
+  /// Copies data from device to host
+  void sync_host() {
+    if (device_backed()) {
+      device_memory::copy_to_host(
+          host_data(), device_data(), imaginary_stride() * 2);
+    }
+  }
+
+  /// Copies data from host to device
+  void sync_device() {
+    if (device_backed()) {
+      device_memory::copy_to_device(
+          device_data(), host_data(), imaginary_stride() * 2);
+    }
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_in_device_to_host(
+    Element const* ptr_device_real,   ///< source device memory
+    Element const* ptr_device_imag,   ///< source device memory
+    LongIndex count = -1) {           ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __MT_STD_MIN(capacity(), count);
+    }
+
+    device_memory::copy_to_host(
+      host_data(), ptr_device_real, count);
+
+    device_memory::copy_to_host(
+      host_data_imag(), ptr_device_imag, count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_in_device_to_device(
+    Element const* ptr_device_real,   ///< source device memory
+    Element const* ptr_device_imag,   ///< source device memory
+    LongIndex count = -1) {           ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __MT_STD_MIN(capacity(), count);
+    }
+
+    device_memory::copy_device_to_device(
+      device_data(), ptr_device_real, count);
+
+    device_memory::copy_device_to_device(
+      device_data_imag(), ptr_device_imag, count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_in_host_to_device(
+    Element const* ptr_host_real,      ///< source host memory
+    Element const* ptr_host_imag,      ///< source host memory
+    LongIndex count = -1) {            ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __MT_STD_MIN(capacity(), count);
+    }
+
+    device_memory::copy_to_device(
+      device_data(), ptr_host_real, count);
+
+    device_memory::copy_to_device(
+      device_data_imag(), ptr_host_imag, count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_in_host_to_host(
+    Element const* ptr_host_real,     ///< source host memory
+    Element const* ptr_host_imag,     ///< source host memory
+    LongIndex count = -1) {           ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __MT_STD_MIN(capacity(), count);
+    }
+
+    device_memory::copy_host_to_host(
+      host_data(), ptr_host_real, count);
+
+    device_memory::copy_host_to_host(
+      host_data_imag(), ptr_host_imag, count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_out_device_to_host(
+    Element * ptr_host_real,           ///< source device memory
+    Element * ptr_host_imag,           ///< source device memory
+    LongIndex count = -1) const {      ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __MT_STD_MIN(capacity(), count);
+    }
+
+    device_memory::copy_to_host(
+      ptr_host_real, device_data(), count);
+
+    device_memory::copy_to_host(
+      ptr_host_imag, device_data_imag(), count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_out_device_to_device(
+    Element * ptr_device_real,        ///< source device memory
+    Element * ptr_device_imag,        ///< source device memory
+    LongIndex count = -1) const {     ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __MT_STD_MIN(capacity(), count);
+    }
+
+    device_memory::copy_device_to_device(
+      ptr_device_real, device_data(), count);
+
+    device_memory::copy_device_to_device(
+      ptr_device_imag, device_data_imag(), count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_out_host_to_device(
+    Element * ptr_device_real,        ///< source device memory
+    Element * ptr_device_imag,        ///< source device memory
+    LongIndex count = -1) const {     ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __MT_STD_MIN(capacity(), count);
+    }
+
+    device_memory::copy_to_device(
+      ptr_device_real, host_data(), count);
+
+    device_memory::copy_to_device(
+      ptr_device_imag, host_data_imag(), count);
+  }
+
+  /// Copy data from a caller-supplied device pointer into host memory.
+  void copy_out_host_to_host(
+    Element * ptr_host_real,          ///< source host memory
+    Element * ptr_host_imag,          ///< source host memory
+    LongIndex count = -1) const {     ///< number of elements to transfer; if negative, entire tensor is overwritten.
+
+    if (count < 0) {
+      count = capacity();
+    }
+    else {
+      count = __MT_STD_MIN(capacity(), count);
+    }
+
+    device_memory::copy_host_to_host(
+      ptr_host_real, host_data(), count);
+
+    device_memory::copy_host_to_host(
+      ptr_host_imag, host_data_imag(), count);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/host_uncompress.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/host_uncompress.h
new file mode 100644
index 000000000..074433e49
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/host_uncompress.h
@@ -0,0 +1,157 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 uncompress sparse matrix from the host side
+*/
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/util/host_tensor.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/util/tensor_view_io.h"
+#include "mutlass/util/reference/host/gemm.h"
+
+namespace mutlass {
+
+// uncompress sparse tensor core A matrix
+template <typename ElementA, typename LayoutA, typename ElementE,
+          typename LayoutE>
+void uncompress(TensorRef<ElementA, LayoutA> uncompressed_tensor_a,
+                TensorRef<ElementA, LayoutA> tensor_a,
+                TensorRef<ElementE, LayoutE> tensor_e, int row, int col) {
+  // How many uncompressed data we can get with ElementE meta data
+  int DecompressedElementsPerElementE =
+      256 / mutlass::sizeof_bits<ElementA>::value;
+
+  // Process 4bit meta data a time
+  int step;
+
+  // 1:2 or 2:4 or 4:8
+  int a, b;
+
+  if (mutlass::sizeof_bits<ElementA>::value == 4) {
+    step = 8;
+    a = 4;
+    b = 8;
+  } else if (mutlass::sizeof_bits<ElementA>::value == 8) {
+    step = 4;
+    a = 2;
+    b = 4;
+  } else if (mutlass::sizeof_bits<ElementA>::value == 16) {
+    step = 4;
+    a = 2;
+    b = 4;
+  } else if (mutlass::sizeof_bits<ElementA>::value == 32) {
+    step = 2;
+    a = 1;
+    b = 2;
+  }
+
+  int ElementsPerE = (mutlass::sizeof_bits<ElementA>::value == 4) ? 2 : 1;
+
+  for (int r = 0; r < row; ++r) {
+    for (int c = 0; c < (col / DecompressedElementsPerElementE); ++c) {
+
+      ElementE meta = tensor_e.at(MatrixCoord(r, c));
+
+      for (int i = 0; i < DecompressedElementsPerElementE; i += step) {
+        int e = (meta >> (i / step * 4)) & 0xf;
+        int idx0 = e & 0x3;
+        int idx1 = e >> 2;
+
+        if (a == 1) idx0 = idx0 / 2;
+
+        for (int ii = 0; ii < step; ii += ElementsPerE) {
+          int real_col =
+              c * DecompressedElementsPerElementE + i + ii;
+          int compressed_col = (real_col / b) * a;
+
+          if (ii == (idx0 * ElementsPerE)) {
+            uncompressed_tensor_a.at(MatrixCoord(r, real_col)) =
+                tensor_a.at(MatrixCoord(r, compressed_col));
+            if (ElementsPerE == 2)
+              uncompressed_tensor_a.at(MatrixCoord(r, real_col + 1)) =
+                  tensor_a.at(MatrixCoord(r, compressed_col + 1));
+          } else if ((ii == (idx1 * ElementsPerE)) && (a != 1)) {
+            uncompressed_tensor_a.at(MatrixCoord(r, real_col)) =
+                tensor_a.at(MatrixCoord(r, compressed_col + ElementsPerE));
+            if (ElementsPerE == 2)
+              uncompressed_tensor_a.at(MatrixCoord(r, real_col + 1)) =
+                  tensor_a.at(
+                      MatrixCoord(r, compressed_col + ElementsPerE + 1));
+          } else {
+            uncompressed_tensor_a.at(MatrixCoord(r, real_col)) =
+                ElementA(0);
+            if (ElementsPerE == 2)
+              uncompressed_tensor_a.at(MatrixCoord(r, real_col + 1)) =
+                  ElementA(0);
+          }
+        }
+      }
+    }
+  }
+}
+
+// uncompress ELL block sparse matrix
+template <typename ElementA, typename LayoutA,
+          typename ElementE, typename LayoutE>
+void uncompress_ell_block_sparse(
+                TensorRef<ElementA, LayoutA> uncompressed_tensor_a,
+                TensorRef<ElementA, LayoutA> tensor_a,
+                TensorRef<ElementE, LayoutE> ell_idx,
+                int rows, int cols,
+                int ell_num_cols, int ell_blocksize) {
+
+  for (int r = 0; r < rows / ell_blocksize; ++r) {
+    for (int c = 0; c < ell_num_cols / ell_blocksize; ++c) {
+
+      ElementE idx = ell_idx.at(MatrixCoord(r, c));
+
+      if (idx != -1) {
+        int row_begin = r * ell_blocksize;
+        int col_begin_real = idx * ell_blocksize;
+        int col_begin = c * ell_blocksize;
+
+        for (int i = 0; i < ell_blocksize; ++i) {
+          for (int j = 0; j < ell_blocksize; ++j) {
+            uncompressed_tensor_a.at(MatrixCoord(row_begin + i, col_begin_real + j)) =
+                tensor_a.at(
+                    MatrixCoord(row_begin + i, col_begin +j));
+          }
+        }
+      }
+    }
+  }
+}
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/index_sequence.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/index_sequence.h
new file mode 100644
index 000000000..765900187
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/index_sequence.h
@@ -0,0 +1,38 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/numeric_types.h"
+
+// integer_sequence moved to mutlass/numeric_types.h
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/packed_stride.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/packed_stride.hpp
new file mode 100644
index 000000000..a5e874653
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/packed_stride.hpp
@@ -0,0 +1,110 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Utilities for packing constructing canonical MuTe stride types for 3.x mainloop params.
+*/
+
+#pragma once
+
+#include "mute/stride.hpp"
+#include "mute/container/array.hpp"   // mute::array
+#include "mutlass/conv/convolution.h" // mutlass::conv::Operator
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Strides without batch mode
+
+template <class IntT>
+mute::Stride<IntT, mute::Int<1>>
+make_mute_packed_stride(mute::Stride<IntT, mute::Int<1>> s, mute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  mute::get<0>(s_copy) = static_cast<IntT>(mute::get<1>(shape_MKL));
+  return s_copy;
+}
+
+template <class IntT>
+mute::Stride<mute::Int<1>, IntT>
+make_mute_packed_stride(mute::Stride<mute::Int<1>, IntT> s, mute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  mute::get<1>(s_copy) = static_cast<IntT>(mute::get<0>(shape_MKL));
+  return s_copy;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Strides with batch mode
+
+template <class IntT>
+mute::Stride<IntT, mute::Int<1>, int64_t>
+make_mute_packed_stride(mute::Stride<IntT, mute::Int<1>, int64_t> s, mute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  mute::get<0>(s_copy) = static_cast<IntT>(mute::get<1>(shape_MKL));
+  int batch_count =  mute::get<2>(shape_MKL);
+  if (batch_count > 1) {
+    mute::get<2>(s_copy) = static_cast<IntT>(mute::get<0>(shape_MKL) * mute::get<1>(shape_MKL));
+  }
+  else {
+    mute::get<2>(s_copy) = static_cast<IntT>(0);
+  }
+  return s_copy;
+}
+
+template <class IntT>
+mute::Stride<mute::Int<1>, IntT, int64_t>
+make_mute_packed_stride(mute::Stride<mute::Int<1>, IntT, int64_t> s, mute::Shape<int,int,int> shape_MKL) {
+  static_assert(std::is_integral_v<IntT>,
+    "Stride must have an integral type so it can be set dynamically. Static strides not supported.");
+  auto s_copy = s;
+  mute::get<1>(s_copy) = static_cast<IntT>(mute::get<0>(shape_MKL));
+  int batch_count =  mute::get<2>(shape_MKL);
+  if (batch_count > 1) {
+    mute::get<2>(s_copy) = static_cast<IntT>(mute::get<0>(shape_MKL) * mute::get<1>(shape_MKL));
+  }
+  else {
+    mute::get<2>(s_copy) = static_cast<IntT>(0);
+  }
+  return s_copy;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/print_error.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/print_error.hpp
new file mode 100644
index 000000000..acfa830a2
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/print_error.hpp
@@ -0,0 +1,294 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <array>
+#include <cassert>
+#include <cmath>
+#include <iostream>
+#include <type_traits>
+
+#include <mute/util/type_traits.hpp>
+#include <mute/tensor.hpp>
+
+#include <mute/numeric/half.hpp>
+#include <mute/numeric/complex.hpp>
+#include <mutlass/layout/layout.h>
+
+// The computed infinity norm does not include
+// any NaN column absolute-value sums.
+struct matrix_inf_norm_result {
+  // Accumulate errors in double, as this is generally
+  // the highest precision that the examples use.
+  double inf_norm = 0.0;
+  bool found_nan = false;
+};
+
+// In theory, mute::Tensor<ViewEngine<T*>, T> could be treated as a view type,
+// and thus passed by value (as std::span or std::string_view would be).
+// However, generic mute::Tensor are more like containers
+// and thus are best passed by reference or const reference.
+template <typename EngineType, typename LayoutType>
+matrix_inf_norm_result
+matrix_inf_norm(mute::Tensor<EngineType, LayoutType> const& host_matrix)
+{
+  using std::abs;
+  using error_type = decltype(std::declval<matrix_inf_norm_result>().inf_norm);
+  using element_type = typename EngineType::value_type;
+
+  error_type inf_norm = 0.0;
+  bool found_nan = false;
+
+  // Computing the infinity norm requires that we be able
+  // to treat the input as a matrix, with rows and columns.
+  const int64_t num_rows = mute::size<0>(host_matrix);
+  const int64_t num_cols = mute::size<1>(host_matrix);
+
+  for(int64_t i = 0; i < num_rows; ++i) {
+    error_type row_abs_sum = 0.0;
+    for(int64_t j = 0; j < num_cols; ++j) {
+      row_abs_sum += abs(host_matrix(i, j));
+    }
+    if(std::isnan(row_abs_sum)) {
+      found_nan = true;
+    } else {
+      inf_norm = row_abs_sum > inf_norm ? row_abs_sum : inf_norm;
+    }
+  }
+
+  return {inf_norm, found_nan};
+}
+
+// Infinity norm of (X - Y).
+template <typename EngineType, typename LayoutType>
+matrix_inf_norm_result
+matrix_diff_inf_norm(mute::Tensor<EngineType, LayoutType> const& X,
+                     mute::Tensor<EngineType, LayoutType> const& Y)
+{
+  using std::abs;
+  using error_type = decltype(std::declval<matrix_inf_norm_result>().inf_norm);
+  using element_type = typename EngineType::value_type;
+
+  assert(mute::size<0>(X) == mute::size<0>(Y));
+  assert(mute::size<1>(X) == mute::size<1>(Y));
+
+  // Computing the infinity norm requires that we be able
+  // to treat the input as a matrix, with rows and columns.
+  const int64_t num_rows = mute::size<0>(X);
+  const int64_t num_cols = mute::size<1>(X);
+
+  error_type inf_norm = 0.0;
+  bool found_nan = false;
+
+  for(int64_t i = 0; i < num_rows; ++i) {
+    error_type row_abs_sum = 0.0;
+    for(int64_t j = 0; j < num_cols; ++j) {
+      row_abs_sum += error_type(abs(element_type(X(i,j)) -
+                                    element_type(Y(i,j))));
+    }
+    if(std::isnan(row_abs_sum)) {
+      found_nan = true;
+    } else {
+      inf_norm = row_abs_sum > inf_norm ? row_abs_sum : inf_norm;
+    }
+  }
+
+  return {inf_norm, found_nan};
+}
+
+template <typename EngineType_A, typename LayoutType_A,
+          typename EngineType_B, typename LayoutType_B,
+          typename EngineType_C, typename LayoutType_C,
+          typename EngineType_C_ref, typename LayoutType_C_ref>
+void
+print_matrix_multiply_mollified_relative_error(
+  char const A_value_type_name[],
+  mute::Tensor<EngineType_A, LayoutType_A> const& A,
+  char const B_value_type_name[],
+  mute::Tensor<EngineType_B, LayoutType_B> const& B,
+  char const C_value_type_name[],
+  mute::Tensor<EngineType_C, LayoutType_C> const& C,
+  mute::Tensor<EngineType_C_ref, LayoutType_C_ref> const& C_ref)
+{
+  const auto [A_norm, A_has_nan] = matrix_inf_norm(A);
+  const auto [B_norm, B_has_nan] = matrix_inf_norm(B);
+  const auto [C_norm, C_has_nan] = matrix_inf_norm(C_ref);
+  const auto [diff_norm, diff_has_nan] = matrix_diff_inf_norm(C, C_ref);
+
+  const auto A_norm_times_B_norm = A_norm * B_norm;
+  const auto relative_error = A_norm_times_B_norm == 0.0 ?
+    diff_norm : (diff_norm / A_norm_times_B_norm);
+
+  // For expected error bounds, please refer to the LAPACK Users' Guide,
+  // in particular https://netlib.org/lapack/lug/node108.html .
+  // Printing the infinity norm of C is a way to check
+  // that both the function being tested (C)
+  // and the reference implementation (C_ref)
+  // don't just do nothing (or fill with zeros).
+  using std::cout;
+  using mute::shape;
+  cout << "Matrix A: " << shape<0>(A) << "x" << shape<1>(A) << " of " << A_value_type_name << '\n'
+       << "Matrix B: " << shape<0>(B) << "x" << shape<1>(B) << " of " << B_value_type_name << '\n'
+       << "Matrix C: " << shape<0>(C) << "x" << shape<1>(C) << " of " << C_value_type_name << '\n'
+       << std::scientific
+       << "Infinity norm of A: " << A_norm << '\n'
+       << "Infinity norm of B: " << B_norm << '\n'
+       << "Infinity norm of C: " << C_norm << '\n'
+       << "Infinity norm of (C - C_ref): " << diff_norm << '\n';
+
+  if(A_norm_times_B_norm == 0.0) {
+    cout << "Mollified relative error: " << relative_error << '\n';
+  } else {
+    cout << "Relative error: " << relative_error << '\n';
+  }
+
+  if (A_has_nan || B_has_nan || C_has_nan || diff_has_nan) {
+    cout << "Did we encounter NaN in A? " << (A_has_nan ? "yes" : "no") << '\n'
+         << "Did we encounter NaN in B? " << (B_has_nan ? "yes" : "no") << '\n'
+         << "Did we encounter NaN in C? " << (C_has_nan ? "yes" : "no") << '\n'
+         << "Did we encounter NaN in (C - C_ref)? " << (diff_has_nan ? "yes" : "no") << '\n';
+  }
+}
+
+template <typename EngineType, typename LayoutType>
+void
+print_matrix_multiply_mollified_relative_error(
+  const char value_type_name[],
+  const mute::Tensor<EngineType, LayoutType>& A,
+  const mute::Tensor<EngineType, LayoutType>& B,
+  const mute::Tensor<EngineType, LayoutType>& C_computed,
+  const mute::Tensor<EngineType, LayoutType>& C_expected)
+{
+  print_matrix_multiply_mollified_relative_error(value_type_name, A, value_type_name, B,
+                                                 value_type_name, C_computed, C_expected);
+}
+
+// Take a MUTLASS HostTensor (or the like) as input,
+// and return a const MuTe Tensor.
+// This is useful for use with the above error printing functions.
+// This implicitly "transposes" if the layout is RowMajor.
+// Note that the HostTensor must be captured by nonconst reference
+// in order for X.host_ref().data() to compile.
+// (MUTLASS is a bit more container-y than MuTe.)
+template<class MutlassHostTensorType>
+auto host_matrix_to_const_mute_tensor(MutlassHostTensorType& X)
+{
+  // The tensors were created with post-transposed extents.
+  const auto extents = X.extent();
+  const auto shape = mute::Shape<int, int>{extents[0], extents[1]};
+  // Both RowMajor and ColumnMajor only store one stride.
+  const int LDX = X.stride(0);
+  const auto strides = [&]() {
+      using input_layout_type = typename std::decay_t<decltype(X)>::Layout;
+      if constexpr (std::is_same_v<input_layout_type, mutlass::layout::ColumnMajor>) {
+        return mute::Stride<int, int>{1, LDX};
+      }
+      else {
+        static_assert(std::is_same_v<input_layout_type, mutlass::layout::RowMajor>);
+        return mute::Stride<int, int>{LDX, 1};
+      }
+    }();
+  const auto layout = mute::make_layout(shape, strides);
+  auto X_data = X.host_ref().data();
+  auto X_data_const = const_cast<std::add_const_t< decltype(X_data)> >(X_data);
+  return mute::make_tensor(X_data_const, layout);
+};
+
+
+// Returns EXIT_SUCCESS if the 2-norm relative error is exactly zero, else returns EXIT_FAILURE.
+// This makes the return value suitable as the return value of main().
+template <typename T1, typename T2>
+int
+print_relative_error(
+    std::size_t n,
+    T1 const& data,
+    T2 const& reference,
+    bool print_verbose = false,
+    bool print_error = true) {
+  using std::abs; using std::sqrt;
+
+  // Use either double or complex<double> for error computation
+  using value_type = mute::remove_cvref_t<decltype(reference[0])>;
+  using error_type = std::conditional_t<mute::is_complex<value_type>::value,
+                                        mute::complex<double>,
+                                        double>;
+
+  if (print_verbose) {
+    std::cout << "Idx:\t"<< "Val\t" << "RefVal\t" << "RelError" << std::endl;
+  }
+
+  double eps = 1e-200;
+
+  double tot_error_sq = 0;
+  double tot_norm_sq = 0;
+  double tot_ind_rel_err = 0;
+  double max_ind_rel_err = 0;
+  for (std::size_t i = 0; i < n; ++i)
+  {
+    error_type val = data[i];
+    error_type ref = reference[i];
+
+    double aref = abs(ref);
+    double diff = abs(ref - val);
+    double rel_error = diff / (aref + eps);
+
+    // Individual relative error
+    tot_ind_rel_err += rel_error;
+
+    // Maximum relative error
+    max_ind_rel_err  = std::max(max_ind_rel_err, rel_error);
+
+    // Total relative error
+    tot_error_sq += diff * diff;
+    tot_norm_sq  += aref * aref;
+
+    if (print_verbose) {
+      std::cout << i << ":\t" << val << "\t" << ref << "\t" << rel_error << std::endl;
+    }
+  }
+
+  printf("Vector reference  norm: [%.5e]\n", sqrt(tot_norm_sq));
+
+  double tot_rel_err = sqrt(tot_error_sq/(tot_norm_sq+eps));
+  if (print_error)
+    printf("Vector  relative error: [%.5e]\n", tot_rel_err);
+
+  double ave_rel_err = tot_ind_rel_err / double(n);
+  if (print_error)
+    printf("Average relative error: [%.5e]\n", ave_rel_err);
+
+  if (print_error)
+    printf("Maximum relative error: [%.5e]\n", max_ind_rel_err);
+
+  return (tot_rel_err == 0.0) ? EXIT_SUCCESS : EXIT_FAILURE;
+}
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/detail/inner_product.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/detail/inner_product.h
new file mode 100644
index 000000000..20922d36c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/detail/inner_product.h
@@ -0,0 +1,135 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for GEMM in host-side code.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+
+namespace mutlass {
+namespace reference {
+namespace detail {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Template function to compute an inner product.
+#pragma hd_warning_disable  // Suppresses warnings when attempting to instantiate with a
+                            // host-only type
+template <typename Atype, typename Btype, typename Ctype>
+MUTLASS_HOST_DEVICE
+Ctype inner_product(Atype a, Btype b, Ctype c) {
+  return Ctype(a) * Ctype(b) + c;
+}
+
+/// Specialization for matrix multiplication with binary operands
+template <>
+MUTLASS_HOST_DEVICE
+int inner_product<Array<bin1_t, 32>, Array<bin1_t, 32>, int>(
+    Array<bin1_t, 32> a,
+    Array<bin1_t, 32> b,
+    int c) {
+
+  int accum = 0;
+  for (int bit = 0; bit < 32; bit++) {
+    accum += a[bit] ^ b[bit];
+  }
+  return accum + c;
+}
+
+/*
+/// Specialization for matrix multiplication with signed 4-bit integer operands
+template <>
+MUTLASS_HOST_DEVICE
+int inner_product<Array<int4b_t, 8>, Array<int4b_t, 8>, int>(
+    Array<int4b_t, 8> a,
+    Array<int4b_t, 8> b,
+    int c) {
+
+  int accum = 0;
+  for (int k = 0; k < 8; k++) {
+    accum += a[k] * b[k];
+  }
+  return accum + c;
+}
+
+/// Specialization for matrix multiplication with unsigned 4-bit integer operands
+template <>
+MUTLASS_HOST_DEVICE
+int inner_product<Array<uint4b_t, 8>, Array<uint4b_t, 8>, int>(
+    Array<uint4b_t, 8> a,
+    Array<uint4b_t, 8> b,
+    int c) {
+
+  int accum = 0;
+  for (int k = 0; k < 8; k++) {
+    accum += a[k] * b[k];
+  }
+  return accum + c;
+}
+*/
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename SrcType, typename DstType>
+struct Cast {
+  // Default behavior: convert to the destination type
+#pragma hd_warning_disable  // Suppresses warnings when attempting to instantiate complex<T> with a
+                            // host-only type
+  MUTLASS_HOST_DEVICE
+  static DstType apply(SrcType src) { return static_cast<DstType>(src); };
+};
+
+template <>
+struct Cast<float, int8_t> {
+  MUTLASS_HOST_DEVICE
+  static int8_t apply(float src) {
+    // Clamp to the range of signed 8-bit integers.
+    return static_cast<int8_t>(fmaxf(-128.f, fminf(127.f, src)));
+  };
+};
+
+template <>
+struct Cast<float, uint8_t> {
+  MUTLASS_HOST_DEVICE
+  static uint8_t apply(float src) {
+    // Clamp to the range of signed 8-bit integers.
+    return static_cast<uint8_t>(fmaxf(0.f, fminf(255.f, src)));
+  };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace detail
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/detail/linear_to_coordinate.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/detail/linear_to_coordinate.h
new file mode 100644
index 000000000..385c757f2
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/detail/linear_to_coordinate.h
@@ -0,0 +1,94 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for GEMM in host-side code.
+*/
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace reference {
+namespace detail {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int Rank, int Index>
+struct LinearToCoordinateHelper {
+
+  MUTLASS_HOST_DEVICE
+  void operator()(Coord<Rank> &coord, int64_t idx, Coord<Rank> const &extent) const {
+
+    int64_t prod = 1;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = Rank - Index; i < Rank; ++i) {
+      prod *= int64_t(extent[i]);
+    }
+
+    coord[Rank - Index - 1] = int(idx / prod);
+
+    int64_t residual = idx % prod;
+    LinearToCoordinateHelper<Rank, Index - 1>()(coord, residual, extent);
+  }
+};
+
+template <int Rank>
+struct LinearToCoordinateHelper<Rank, 0> {
+
+  MUTLASS_HOST_DEVICE
+  void operator()(Coord<Rank> &coord, int64_t idx, Coord<Rank> const &extent) const {
+    coord[Rank - 1] = int(idx);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int Rank>
+struct LinearToCoordinate {
+
+  MUTLASS_HOST_DEVICE
+  void operator()(Coord<Rank> &coord, int64_t idx, Coord<Rank> const &extent) const {
+    LinearToCoordinateHelper<Rank, Rank - 1>()(coord, idx, extent);
+  }
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace detail
+} // namespace reference
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/convolution.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/convolution.h
new file mode 100644
index 000000000..87a6227b1
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/convolution.h
@@ -0,0 +1,1549 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for convolution in device-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/functional.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/matrix_shape.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_ref.h"
+#include "mutlass/conv/convolution.h"
+#include "mutlass/conv/conv2d_problem_size.h"
+#include "mutlass/conv/conv3d_problem_size.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace kernel {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+///                                   Conv2d device reference kernel
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Conv2d Fprop kernel - y = fprop(x, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>,
+  int kThreadM = 2,       // shape of a thread's tile in the GEMM M dimension
+  int kThreadN = 4,       // shape of a thread's tile in the GEMM N dimension
+  int kCtaShapeM = 16,    // shape of a threadblock in units of threads
+  int kCtaShapeN = 8      // shape of a threadblock in units of threads
+>
+__global__ void Conv2dFprop(
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_x,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_y_in,
+  TensorRef<ElementC, LayoutC> tensor_y_out,
+  ElementCompute alpha,
+  ElementCompute beta
+  ) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  ElementAccumulator element_A[kThreadM];
+  ElementAccumulator element_B[kThreadN];
+  ElementAccumulator accum[kThreadM][kThreadN];
+
+  int64_t npq_start = int64_t(blockIdx.x) * kCtaShapeM * kThreadM + threadIdx.x * kThreadM;
+  int k_start = blockIdx.y * kCtaShapeN * kThreadN + threadIdx.y * kThreadN;
+
+  int thread_n[kThreadM];
+  int thread_p[kThreadM];
+  int thread_q[kThreadM];
+
+  // Compute N, P, Q coordinates for each row of a thread's tile
+  int64_t PQ = int64_t(problem_size.P) * problem_size.Q;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+
+    int64_t npq = npq_start + m;
+
+    thread_n[m] = int(npq / PQ);
+
+    int64_t residual = npq % PQ;
+    thread_p[m] = int(residual / problem_size.Q);
+    thread_q[m] = int(residual % problem_size.Q);
+  }
+
+  // Clear accumulators
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < kThreadN; ++n) {
+      accum[m][n] = ElementAccumulator();
+    }
+  }
+
+  int c_per_group = problem_size.C / problem_size.groups;
+  int k_per_group = problem_size.K / problem_size.groups;
+
+  // Compute convolution
+  for (int R = 0; R < problem_size.R; ++R) {
+    for (int S = 0; S < problem_size.S; ++S) {
+      for (int C = 0; C < problem_size.C; ++C) {
+
+        // Get group id of currnet channel
+        int c_group_idx = C / c_per_group;
+
+        // Load from activations tensor
+        int filter_r = R;
+        int filter_s = S;
+
+        if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+          filter_r = problem_size.R - 1 - R;
+          filter_s = problem_size.S - 1 - S;
+        }
+
+        MUTLASS_PRAGMA_UNROLL
+        for (int m = 0; m < kThreadM; ++m) {
+          int h = thread_p[m] * problem_size.stride_h - problem_size.pad_h + filter_r * problem_size.dilation_h;
+          int w = thread_q[m] * problem_size.stride_w - problem_size.pad_w + filter_s * problem_size.dilation_w;
+
+          if (thread_n[m] < problem_size.N && h >= 0 && h < problem_size.H && w >= 0 && w < problem_size.W) {
+            element_A[m] = ElementAccumulator(tensor_x.at({thread_n[m], h, w, C}));
+          }
+          else {
+            element_A[m] = ElementAccumulator();
+          }
+        }
+
+        // Load from filters tensor
+        MUTLASS_PRAGMA_UNROLL
+        for (int n = 0; n < kThreadN; ++n) {
+          int thread_k = k_start + n;
+          int k_group_idx = thread_k / k_per_group;
+
+          if (thread_k < problem_size.K && k_group_idx == c_group_idx) {
+            element_B[n] = ElementAccumulator(tensor_w.at({thread_k, R, S, C % c_per_group}));
+          }
+          else {
+            element_B[n] = ElementAccumulator();
+          }
+        }
+
+        // Accumulate matrix product
+        MUTLASS_PRAGMA_UNROLL
+        for (int m = 0; m < kThreadM; ++m) {
+          MUTLASS_PRAGMA_UNROLL
+          for (int n = 0; n < kThreadN; ++n) {
+            accum[m][n] = inner_product_op(element_A[m], element_B[n], accum[m][n]);
+          }
+        }
+      }
+    }
+  }
+
+  // Write out the results
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    if (thread_n[m] < problem_size.N && thread_p[m] < problem_size.P && thread_q[m] < problem_size.Q) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < kThreadN; ++n) {
+        int thread_k = k_start + n;
+        if (thread_k < problem_size.K) {
+
+          ElementCompute c_ref = ElementCompute();
+          if (beta != ElementCompute()) {
+            c_ref = ElementCompute(tensor_y_in.at({thread_n[m], thread_p[m], thread_q[m], thread_k}));
+          }
+
+          tensor_y_out.at({thread_n[m], thread_p[m], thread_q[m], thread_k}) = convert_op(
+            alpha * ElementCompute(accum[m][n]) + beta * c_ref);
+        }
+      }
+    }
+  }
+}
+
+// Conv3d Fprop kernel - y = fprop(x, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator =  ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>,
+  int kThreadM = 2,       // shape of a thread's tile in the GEMM M dimension
+  int kThreadN = 4,       // shape of a thread's tile in the GEMM N dimension
+  int kCtaShapeM = 16,    // shape of a threadblock in units of threads
+  int kCtaShapeN = 8      // shape of a threadblock in units of threads
+>
+__global__ void Conv3dFprop(
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_x,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_y_in,
+  TensorRef<ElementC, LayoutC> tensor_y_out,
+  ElementCompute alpha,
+  ElementCompute beta
+  ) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  ElementAccumulator element_A[kThreadM];
+  ElementAccumulator element_B[kThreadN];
+  ElementAccumulator accum[kThreadM][kThreadN];
+
+  int64_t nzpq_start = int64_t(blockIdx.x) * kCtaShapeM * kThreadM + threadIdx.x * kThreadM;
+  int k_start = blockIdx.y * kCtaShapeN * kThreadN + threadIdx.y * kThreadN;
+
+  int thread_n[kThreadM];
+  int thread_z[kThreadM];
+  int thread_p[kThreadM];
+  int thread_q[kThreadM];
+
+  // Compute N, Z, P, Q coordinates for each row of a thread's tile
+  int64_t PQ = int64_t(problem_size.P) * problem_size.Q;
+  int64_t ZPQ = PQ * problem_size.Z;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+
+    int64_t nzpq = nzpq_start + m;
+
+    thread_n[m] = int(nzpq / ZPQ);
+
+    int64_t residual = nzpq % ZPQ;
+    thread_z[m] = int(residual / PQ);
+
+    residual = residual % PQ;
+    thread_p[m] = int(residual / problem_size.Q);
+    thread_q[m] = int(residual % problem_size.Q);
+  }
+
+  // Clear accumulators
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < kThreadN; ++n) {
+      accum[m][n] = ElementAccumulator();
+    }
+  }
+
+  // Compute convolution
+  for (int T = 0; T < problem_size.T; ++T) {
+    for (int R = 0; R < problem_size.R; ++R) {
+      for (int S = 0; S < problem_size.S; ++S) {
+        for (int C = 0; C < problem_size.C; ++C) {
+
+          // Load from activations tensor
+          int filter_t = T;
+          int filter_r = R;
+          int filter_s = S;
+
+          if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+            filter_t = problem_size.T - 1 - T;
+            filter_r = problem_size.R - 1 - R;
+            filter_s = problem_size.S - 1 - S;
+          }
+
+          MUTLASS_PRAGMA_UNROLL
+          for (int m = 0; m < kThreadM; ++m) {
+            int d = thread_z[m] * problem_size.stride_d - problem_size.pad_d + filter_t * problem_size.dilation_d;
+            int h = thread_p[m] * problem_size.stride_h - problem_size.pad_h + filter_r * problem_size.dilation_h;
+            int w = thread_q[m] * problem_size.stride_w - problem_size.pad_w + filter_s * problem_size.dilation_w;
+
+            if (thread_n[m] < problem_size.N &&
+              d >= 0 && d < problem_size.D &&
+              h >= 0 && h < problem_size.H &&
+              w >= 0 && w < problem_size.W) {
+
+              element_A[m] = ElementAccumulator(tensor_x.at({thread_n[m], d, h, w, C}));
+            }
+            else {
+              element_A[m] = ElementAccumulator();
+            }
+          }
+
+          // Load from filters tensor
+          MUTLASS_PRAGMA_UNROLL
+          for (int n = 0; n < kThreadN; ++n) {
+            int thread_k = k_start + n;
+
+            if (thread_k < problem_size.K) {
+              element_B[n] = ElementAccumulator(tensor_w.at({thread_k, T, R, S, C}));
+            }
+            else {
+              element_B[n] = ElementAccumulator();
+            }
+          }
+
+          // Accumulate matrix product
+          MUTLASS_PRAGMA_UNROLL
+          for (int m = 0; m < kThreadM; ++m) {
+            MUTLASS_PRAGMA_UNROLL
+            for (int n = 0; n < kThreadN; ++n) {
+              accum[m][n] = inner_product_op(element_A[m], element_B[n], accum[m][n]);
+            }
+          }
+
+        } // for (C)
+      } // for (S)
+    }  // for (R)
+  } // for (T)
+
+  // Write out the results
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+
+    if (thread_n[m] < problem_size.N &&
+      thread_z[m] < problem_size.Z &&
+      thread_p[m] < problem_size.P &&
+      thread_q[m] < problem_size.Q) {
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < kThreadN; ++n) {
+        int thread_k = k_start + n;
+        if (thread_k < problem_size.K) {
+
+          ElementCompute c_ref = ElementCompute();
+          if (beta != ElementCompute()) {
+            c_ref = ElementCompute(tensor_y_in.at({thread_n[m], thread_z[m], thread_p[m], thread_q[m], thread_k}));
+          }
+
+          tensor_y_out.at({thread_n[m], thread_z[m], thread_p[m], thread_q[m], thread_k}) = convert_op(
+            alpha * ElementCompute(accum[m][n]) + beta * c_ref);
+        }
+      } // for (n)
+
+    }
+  } // for (m)
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Conv2d dgrad kernel - dx = dgrad(dy, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>,
+  int kThreadM = 2,       // shape of a thread's tile in the GEMM M dimension
+  int kThreadN = 4,       // shape of a thread's tile in the GEMM N dimension
+  int kCtaShapeM = 16,    // shape of a threadblock in units of threads
+  int kCtaShapeN = 8      // shape of a threadblock in units of threads
+>
+__global__ void Conv2dDgrad(
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_dx_in,
+  TensorRef<ElementC, LayoutC> tensor_dx_out,
+  ElementCompute alpha,
+  ElementCompute beta
+  ) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  ElementAccumulator element_A[kThreadM];
+  ElementAccumulator element_B[kThreadN];
+  ElementAccumulator accum[kThreadM][kThreadN];
+
+  int64_t nhw_start = int64_t(blockIdx.x) * kCtaShapeM * kThreadM + threadIdx.x * kThreadM;
+  int c_start = blockIdx.y * kCtaShapeN * kThreadN + threadIdx.y * kThreadN;
+
+  int thread_n[kThreadM];
+  int thread_h[kThreadM];
+  int thread_w[kThreadM];
+
+  // Compute N, H, W coordinates for each row of a thread's tile
+  int64_t HW = int64_t(problem_size.H) * problem_size.W;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+
+    int64_t nhw = nhw_start + m;
+
+    thread_n[m] = int(nhw / HW);
+
+    int64_t residual = nhw % HW;
+    thread_h[m] = int(residual / problem_size.W);
+    thread_w[m] = int(residual % problem_size.W);
+  }
+
+  // Clear accumulators
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < kThreadN; ++n) {
+      accum[m][n] = ElementAccumulator();
+    }
+  }
+
+  // Compute convolution
+  for (int R = 0; R < problem_size.R; ++R) {
+    for (int S = 0; S < problem_size.S; ++S) {
+      for (int K = 0; K < problem_size.K; ++K) {
+
+        // Load from activations tensor
+        int filter_r = R;
+        int filter_s = S;
+
+        if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+          filter_r = problem_size.R - 1 - R;
+          filter_s = problem_size.S - 1 - S;
+        }
+
+        MUTLASS_PRAGMA_UNROLL
+        for (int m = 0; m < kThreadM; ++m) {
+
+          int p = thread_h[m] + problem_size.pad_h - filter_r * problem_size.dilation_h;
+          int q = thread_w[m] + problem_size.pad_w - filter_s * problem_size.dilation_w;
+
+          element_A[m] = ElementAccumulator();
+
+          if (p >= 0 && !(p % problem_size.stride_h) && q >= 0 && !(q % problem_size.stride_w)) {
+
+            p = p / problem_size.stride_h;
+            q = q / problem_size.stride_w;
+
+            if (thread_n[m] < problem_size.N && p < problem_size.P && q < problem_size.Q) {
+              element_A[m] = ElementAccumulator(tensor_dy.at({thread_n[m], p, q, K}));
+            }
+          }
+        }
+
+        // Load from filters tensor
+        MUTLASS_PRAGMA_UNROLL
+        for (int n = 0; n < kThreadN; ++n) {
+          int thread_c = c_start + n;
+
+          if (thread_c < problem_size.C) {
+            element_B[n] = ElementAccumulator(tensor_w.at({K, R, S, thread_c}));
+          }
+          else {
+            element_B[n] = ElementAccumulator();
+          }
+        }
+
+        // Accumulate matrix product
+        MUTLASS_PRAGMA_UNROLL
+        for (int m = 0; m < kThreadM; ++m) {
+          MUTLASS_PRAGMA_UNROLL
+          for (int n = 0; n < kThreadN; ++n) {
+            accum[m][n] = inner_product_op(element_A[m], element_B[n], accum[m][n]);
+          }
+        }
+      }
+    }
+  }
+
+  // Write out the results
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+
+    if (thread_n[m] < problem_size.N && thread_h[m] < problem_size.H && thread_w[m] < problem_size.W) {
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < kThreadN; ++n) {
+        int thread_c = c_start + n;
+        if (thread_c < problem_size.C) {
+
+          ElementCompute c_ref = ElementCompute();
+          if (beta != ElementCompute()) {
+            c_ref = ElementCompute(tensor_dx_in.at({thread_n[m], thread_h[m], thread_w[m], thread_c}));
+          }
+
+          tensor_dx_out.at({thread_n[m], thread_h[m], thread_w[m], thread_c}) = convert_op(
+            alpha * ElementCompute(accum[m][n]) + beta * c_ref);
+        }
+      }
+    }
+  }
+}
+
+// Conv3d dgrad kernel - dx = dgrad(dy, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>,
+  int kThreadM = 2,       // shape of a thread's tile in the GEMM M dimension
+  int kThreadN = 4,       // shape of a thread's tile in the GEMM N dimension
+  int kCtaShapeM = 16,    // shape of a threadblock in units of threads
+  int kCtaShapeN = 8      // shape of a threadblock in units of threads
+>
+__global__ void Conv3dDgrad(
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_dx_in,
+  TensorRef<ElementC, LayoutC> tensor_dx_out,
+  ElementCompute alpha,
+  ElementCompute beta
+  ) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  ElementAccumulator element_A[kThreadM];
+  ElementAccumulator element_B[kThreadN];
+  ElementAccumulator accum[kThreadM][kThreadN];
+
+  int64_t ndhw_start = int64_t(blockIdx.x) * kCtaShapeM * kThreadM + threadIdx.x * kThreadM;
+  int c_start = blockIdx.y * kCtaShapeN * kThreadN + threadIdx.y * kThreadN;
+
+  int thread_n[kThreadM];
+  int thread_d[kThreadM];
+  int thread_h[kThreadM];
+  int thread_w[kThreadM];
+
+  // Compute N, H, W coordinates for each row of a thread's tile
+  int64_t HW = int64_t(problem_size.H) * problem_size.W;
+  int64_t DHW = HW * problem_size.D;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+
+    int64_t ndhw = ndhw_start + m;
+
+    thread_n[m] = int(ndhw / DHW);
+
+    int64_t residual = ndhw % DHW;
+    thread_d[m] = int(residual / HW);
+
+    residual = residual % HW;
+    thread_h[m] = int(residual / problem_size.W);
+    thread_w[m] = int(residual % problem_size.W);
+  }
+
+  // Clear accumulators
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < kThreadN; ++n) {
+      accum[m][n] = ElementAccumulator();
+    }
+  }
+
+  // Compute convolution
+  for (int T = 0; T < problem_size.T; ++T) {
+    for (int R = 0; R < problem_size.R; ++R) {
+      for (int S = 0; S < problem_size.S; ++S) {
+        for (int K = 0; K < problem_size.K; ++K) {
+
+          // Load from activations tensor
+          int filter_t = T;
+          int filter_r = R;
+          int filter_s = S;
+
+          if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+            filter_t = problem_size.T - 1 - T;
+            filter_r = problem_size.R - 1 - R;
+            filter_s = problem_size.S - 1 - S;
+          }
+
+          MUTLASS_PRAGMA_UNROLL
+          for (int m = 0; m < kThreadM; ++m) {
+
+            int z = thread_d[m] + problem_size.pad_d - filter_t * problem_size.dilation_d;
+            int p = thread_h[m] + problem_size.pad_h - filter_r * problem_size.dilation_h;
+            int q = thread_w[m] + problem_size.pad_w - filter_s * problem_size.dilation_w;
+
+            element_A[m] = ElementAccumulator();
+
+            if (z >= 0 && !(z % problem_size.stride_d) &&
+              p >= 0 && !(p % problem_size.stride_h) &&
+              q >= 0 && !(q % problem_size.stride_w)) {
+
+              z = z / problem_size.stride_d;
+              p = p / problem_size.stride_h;
+              q = q / problem_size.stride_w;
+
+              if (thread_n[m] < problem_size.N && z < problem_size.Z && p < problem_size.P && q < problem_size.Q) {
+                element_A[m] = ElementAccumulator(tensor_dy.at({thread_n[m], z, p, q, K}));
+              }
+            }
+          }
+
+          // Load from filters tensor
+          MUTLASS_PRAGMA_UNROLL
+          for (int n = 0; n < kThreadN; ++n) {
+            int thread_c = c_start + n;
+
+            if (thread_c < problem_size.C) {
+              element_B[n] = ElementAccumulator(tensor_w.at({K, T, R, S, thread_c}));
+            }
+            else {
+              element_B[n] = ElementAccumulator();
+            }
+          }
+
+          // Accumulate matrix product
+          MUTLASS_PRAGMA_UNROLL
+          for (int m = 0; m < kThreadM; ++m) {
+            MUTLASS_PRAGMA_UNROLL
+            for (int n = 0; n < kThreadN; ++n) {
+              accum[m][n] = inner_product_op(element_A[m], element_B[n], accum[m][n]);
+            }
+          }
+
+        } // for (C)
+      } // for (S)
+    } // for (R)
+  } // for (T)
+
+  // Write out the results
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+
+    if (thread_n[m] < problem_size.N &&
+      thread_d[m] < problem_size.D &&
+      thread_h[m] < problem_size.H &&
+      thread_w[m] < problem_size.W) {
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < kThreadN; ++n) {
+        int thread_c = c_start + n;
+        if (thread_c < problem_size.C) {
+
+          ElementCompute c_ref = ElementCompute();
+          if (beta != ElementCompute()) {
+            c_ref = ElementCompute(tensor_dx_in.at({thread_n[m], thread_d[m], thread_h[m], thread_w[m], thread_c}));
+          }
+
+          tensor_dx_out.at({thread_n[m], thread_d[m], thread_h[m], thread_w[m], thread_c}) = convert_op(
+            alpha * ElementCompute(accum[m][n]) + beta * c_ref);
+        }
+      }
+    }
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+// Conv2d wgrad kernel - dw = wgrad(dy, x)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>,
+  int kThreadM = 2,       // shape of a thread's tile in the GEMM M dimension
+  int kThreadN = 4,       // shape of a thread's tile in the GEMM N dimension
+  int kCtaShapeM = 8,     // shape of a threadblock in units of threads
+  int kCtaShapeN = 16     // shape of a threadblock in units of threads
+>
+__global__ void Conv2dWgrad(
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_x,
+  TensorRef<ElementC, LayoutC> tensor_dw_in,
+  TensorRef<ElementC, LayoutC> tensor_dw_out,
+  ElementCompute alpha,
+  ElementCompute beta
+  ) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  ElementAccumulator element_A[kThreadM];
+  ElementAccumulator element_B[kThreadN];
+  ElementAccumulator accum[kThreadM][kThreadN];
+
+  int k_start = blockIdx.x * kCtaShapeM * kThreadM + threadIdx.x * kThreadM;
+  int64_t rsc_start = int64_t(blockIdx.y) * kCtaShapeN * kThreadN + threadIdx.y * kThreadN;
+
+  int thread_r[kThreadN];
+  int thread_s[kThreadN];
+  int thread_c[kThreadN];
+
+  // Compute R, S, C coordinates for each row of a thread's tile
+  int64_t SC = int64_t(problem_size.S) * problem_size.C;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int n = 0; n < kThreadN; ++n) {
+
+    int64_t rsc = rsc_start + n;
+    int64_t residual = rsc % SC;
+
+    thread_r[n] = int(rsc / SC);
+    thread_s[n] = int(residual / problem_size.C);
+    thread_c[n] = int(residual % problem_size.C);
+  }
+
+  // Clear accumulators
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < kThreadN; ++n) {
+      accum[m][n] = ElementAccumulator();
+    }
+  }
+
+  // Compute convolution
+  for (int N = 0; N < problem_size.N; ++N) {
+    for (int P = 0; P < problem_size.P; ++P) {
+      for (int Q = 0; Q < problem_size.Q; ++Q) {
+
+        MUTLASS_PRAGMA_UNROLL
+        for (int m = 0; m < kThreadM; ++m) {
+          int thread_k = k_start + m;
+
+          element_A[m] = ElementAccumulator();
+
+          if (thread_k < problem_size.K) {
+            element_A[m] = ElementAccumulator(tensor_dy.at({N, P, Q, thread_k}));
+          }
+        }
+
+        // Load from filters tensor
+        MUTLASS_PRAGMA_UNROLL
+        for (int n = 0; n < kThreadN; ++n) {
+
+          // Load from activations tensor
+          int filter_r = thread_r[n];
+          int filter_s = thread_s[n];
+
+          if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+            filter_r = problem_size.R - 1 - filter_r;
+            filter_s = problem_size.S - 1 - filter_s;
+          }
+
+          int h = P * problem_size.stride_h - problem_size.pad_h + filter_r * problem_size.dilation_h;
+          int w = Q * problem_size.stride_w - problem_size.pad_w + filter_s * problem_size.dilation_w;
+
+          element_B[n] = ElementAccumulator();
+
+          if (h >= 0 && h < problem_size.H && w >= 0 && w < problem_size.W && thread_c[n] < problem_size.C) {
+            element_B[n] = ElementAccumulator(tensor_x.at({N, h, w, thread_c[n]}));
+          }
+        }
+
+        // Accumulate matrix product
+        MUTLASS_PRAGMA_UNROLL
+        for (int m = 0; m < kThreadM; ++m) {
+          MUTLASS_PRAGMA_UNROLL
+          for (int n = 0; n < kThreadN; ++n) {
+            accum[m][n] = inner_product_op(element_A[m], element_B[n], accum[m][n]);
+          }
+        }
+      }
+    }
+  }
+
+  // Write out the results
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    int thread_k = k_start + m;
+
+    if (thread_k < problem_size.K) {
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < kThreadN; ++n) {
+
+        if (thread_r[n] < problem_size.R && thread_s[n] < problem_size.S && thread_c[n] < problem_size.C) {
+
+          ElementCompute c_ref = ElementCompute();
+
+          if (beta != ElementCompute()) {
+            c_ref = ElementCompute(tensor_dw_in.at({thread_k, thread_r[n], thread_s[n], thread_c[n]}));
+          }
+
+          tensor_dw_out.at({thread_k, thread_r[n], thread_s[n], thread_c[n]}) = convert_op(
+            alpha * ElementCompute(accum[m][n]) + beta * c_ref);
+        }
+      }
+    }
+  }
+}
+
+// Conv3d wgrad kernel - dw = wgrad(dy, x)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>,
+  int kThreadM = 2,       // shape of a thread's tile in the GEMM M dimension
+  int kThreadN = 4,       // shape of a thread's tile in the GEMM N dimension
+  int kCtaShapeM = 8,     // shape of a threadblock in units of threads
+  int kCtaShapeN = 16     // shape of a threadblock in units of threads
+>
+__global__ void Conv3dWgrad(
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_x,
+  TensorRef<ElementC, LayoutC> tensor_dw_in,
+  TensorRef<ElementC, LayoutC> tensor_dw_out,
+  ElementCompute alpha,
+  ElementCompute beta
+  ) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  ElementAccumulator element_A[kThreadM];
+  ElementAccumulator element_B[kThreadN];
+  ElementAccumulator accum[kThreadM][kThreadN];
+
+  int k_start = blockIdx.x * kCtaShapeM * kThreadM + threadIdx.x * kThreadM;
+  int64_t trsc_start = int64_t(blockIdx.y) * kCtaShapeN * kThreadN + threadIdx.y * kThreadN;
+
+  int thread_t[kThreadN];
+  int thread_r[kThreadN];
+  int thread_s[kThreadN];
+  int thread_c[kThreadN];
+
+  // Compute R, S, C coordinates for each row of a thread's tile
+  int64_t SC = int64_t(problem_size.S) * problem_size.C;
+  int64_t RSC = SC * problem_size.R;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int n = 0; n < kThreadN; ++n) {
+
+    int64_t trsc = trsc_start + n;
+
+    thread_t[n] = int(trsc / RSC);
+
+    int64_t residual = trsc % RSC;
+    thread_r[n] = int(residual / SC);
+
+    residual = residual % SC;
+    thread_s[n] = int(residual / problem_size.C);
+    thread_c[n] = int(residual % problem_size.C);
+  }
+
+  // Clear accumulators
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int n = 0; n < kThreadN; ++n) {
+      accum[m][n] = ElementAccumulator();
+    }
+  }
+
+  // Compute convolution
+  for (int N = 0; N < problem_size.N; ++N) {
+    for (int Z = 0; Z < problem_size.Z; ++Z) {
+      for (int P = 0; P < problem_size.P; ++P) {
+        for (int Q = 0; Q < problem_size.Q; ++Q) {
+
+          MUTLASS_PRAGMA_UNROLL
+          for (int m = 0; m < kThreadM; ++m) {
+            int thread_k = k_start + m;
+
+            element_A[m] = ElementAccumulator();
+
+            if (thread_k < problem_size.K) {
+              element_A[m] = ElementAccumulator(tensor_dy.at({N, Z, P, Q, thread_k}));
+            }
+          }
+
+          // Load from filters tensor
+          MUTLASS_PRAGMA_UNROLL
+          for (int n = 0; n < kThreadN; ++n) {
+
+            // Load from activations tensor
+            int filter_t = thread_t[n];
+            int filter_r = thread_r[n];
+            int filter_s = thread_s[n];
+
+            if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+              filter_t = problem_size.T - 1 - filter_t;
+              filter_r = problem_size.R - 1 - filter_r;
+              filter_s = problem_size.S - 1 - filter_s;
+            }
+
+            int d = Z * problem_size.stride_d - problem_size.pad_w + filter_t * problem_size.dilation_d;
+            int h = P * problem_size.stride_h - problem_size.pad_h + filter_r * problem_size.dilation_h;
+            int w = Q * problem_size.stride_w - problem_size.pad_w + filter_s * problem_size.dilation_w;
+
+            element_B[n] = ElementAccumulator();
+
+            if (d >= 0 && d < problem_size.D &&
+              h >= 0 && h < problem_size.H &&
+              w >= 0 && w < problem_size.W &&
+              thread_c[n] < problem_size.C) {
+
+              element_B[n] = ElementAccumulator(tensor_x.at({N, d, h, w, thread_c[n]}));
+            }
+          }
+
+          // Accumulate matrix product
+          MUTLASS_PRAGMA_UNROLL
+          for (int m = 0; m < kThreadM; ++m) {
+            MUTLASS_PRAGMA_UNROLL
+            for (int n = 0; n < kThreadN; ++n) {
+              accum[m][n] = inner_product_op(element_A[m], element_B[n], accum[m][n]);
+            }
+          }
+
+        } // for (Q)
+      } // for (P)
+    } // for (Z)
+  } // for (N)
+
+  // Write out the results
+  MUTLASS_PRAGMA_UNROLL
+  for (int m = 0; m < kThreadM; ++m) {
+    int thread_k = k_start + m;
+
+    if (thread_k < problem_size.K) {
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int n = 0; n < kThreadN; ++n) {
+
+        if (thread_t[n] < problem_size.T &&
+          thread_r[n] < problem_size.R &&
+          thread_s[n] < problem_size.S &&
+          thread_c[n] < problem_size.C) {
+
+          ElementCompute c_ref = ElementCompute();
+
+          if (beta != ElementCompute()) {
+            c_ref = ElementCompute(tensor_dw_in.at({thread_k, thread_t[n], thread_r[n], thread_s[n], thread_c[n]}));
+          }
+
+          tensor_dw_out.at({thread_k, thread_t[n], thread_r[n], thread_s[n], thread_c[n]}) = convert_op(
+            alpha * ElementCompute(accum[m][n]) + beta * c_ref);
+        }
+      }
+    }
+  }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace kernel
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Conv2d Fprop dispatcher - y = fprop(x, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+Status Conv2dFprop(
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_x,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_y_in,
+  TensorRef<ElementC, LayoutC> tensor_y_out,
+  ElementCompute alpha,
+  ElementCompute beta,
+  musaStream_t stream = nullptr) {
+
+  //
+  // Blocking factors improve performance of reference implementation
+  //
+
+  int const kThreadM = 4;       // shape of a thread's tile in the GEMM M dimension
+  int const kThreadN = 4;       // shape of a thread's tile in the GEMM N dimension
+  int const kCtaShapeM = 16;    // shape of a threadblock in units of threads
+  int const kCtaShapeN = 8;     // shape of a threadblock in units of threads
+
+  int64_t npq = int64_t(problem_size.N) * problem_size.P * problem_size.Q;
+  int64_t blocks_m = (npq + (kCtaShapeM * kThreadM) - 1) / (kCtaShapeM * kThreadM);
+
+  dim3 block(kCtaShapeM, kCtaShapeN);
+  dim3 grid(uint32_t(blocks_m), (problem_size.K + (kCtaShapeN * kThreadN) - 1) / (kCtaShapeN * kThreadN));
+
+  kernel::Conv2dFprop<
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    ElementCompute,
+    ElementAccumulator,
+    ConvertOp,
+    InnerProductOp,
+    kThreadM,
+    kThreadN,
+    kCtaShapeM,
+    kCtaShapeN
+  ><<< grid, block, 0, stream >>>(
+    problem_size,
+    tensor_x,
+    tensor_w,
+    tensor_y_in,
+    tensor_y_out,
+    alpha,
+    beta
+  );
+
+  musaError_t result = musaPeekAtLastError();
+  if (result != musaSuccess) {
+    return Status::kErrorInternal;
+  }
+
+  return Status::kSuccess;
+}
+
+/// Conv3d Fprop dispatcher - y = fprop(x, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+Status Conv3dFprop(
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_x,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_y_in,
+  TensorRef<ElementC, LayoutC> tensor_y_out,
+  ElementCompute alpha,
+  ElementCompute beta,
+  musaStream_t stream = nullptr) {
+
+  //
+  // Blocking factors improve performance of reference implementation
+  //
+
+  int const kThreadM = 4;       // shape of a thread's tile in the GEMM M dimension
+  int const kThreadN = 4;       // shape of a thread's tile in the GEMM N dimension
+  int const kCtaShapeM = 16;    // shape of a threadblock in units of threads
+  int const kCtaShapeN = 8;     // shape of a threadblock in units of threads
+
+  int64_t nzpq = int64_t(problem_size.N) * problem_size.Z * problem_size.P * problem_size.Q;
+  int64_t blocks_m = (nzpq + (kCtaShapeM * kThreadM) - 1) / (kCtaShapeM * kThreadM);
+
+  dim3 block(kCtaShapeM, kCtaShapeN);
+  dim3 grid(uint32_t(blocks_m), (problem_size.K + (kCtaShapeN * kThreadN) - 1) / (kCtaShapeN * kThreadN));
+
+  kernel::Conv3dFprop<
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    ElementCompute,
+    ElementAccumulator,
+    ConvertOp,
+    InnerProductOp,
+    kThreadM,
+    kThreadN,
+    kCtaShapeM,
+    kCtaShapeN
+  ><<< grid, block, 0, stream >>>(
+    problem_size,
+    tensor_x,
+    tensor_w,
+    tensor_y_in,
+    tensor_y_out,
+    alpha,
+    beta
+  );
+
+  musaError_t result = musaPeekAtLastError();
+  if (result != musaSuccess) {
+    return Status::kErrorInternal;
+  }
+
+  return Status::kSuccess;
+}
+
+/// Conv2d Dgrad dispatcher - dx = dgrad(dy, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+Status Conv2dDgrad(
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_dx_in,
+  TensorRef<ElementC, LayoutC> tensor_dx_out,
+  ElementCompute alpha,
+  ElementCompute beta,
+  musaStream_t stream = nullptr) {
+
+  //
+  // Blocking factors improve performance of reference implementation
+  //
+
+  int const kThreadM = 2;       // shape of a thread's tile in the GEMM M dimension
+  int const kThreadN = 4;       // shape of a thread's tile in the GEMM N dimension
+  int const kCtaShapeM = 16;    // shape of a threadblock in units of threads
+  int const kCtaShapeN = 8;     // shape of a threadblock in units of threads
+
+  int64_t nhw = int64_t(problem_size.N) * problem_size.H * problem_size.W;
+  int64_t blocks_m = (nhw + (kCtaShapeM * kThreadM) - 1) / (kCtaShapeM * kThreadM);
+
+  dim3 block(kCtaShapeM, kCtaShapeN);
+  dim3 grid(uint32_t(blocks_m), (problem_size.C + (kCtaShapeN * kThreadN) - 1) / (kCtaShapeN * kThreadN));
+
+  kernel::Conv2dDgrad<
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    ElementCompute,
+    ElementAccumulator,
+    ConvertOp,
+    InnerProductOp,
+    kThreadM,
+    kThreadN,
+    kCtaShapeM,
+    kCtaShapeN
+  ><<< grid, block, 0, stream >>>(
+    problem_size,
+    tensor_dy,
+    tensor_w,
+    tensor_dx_in,
+    tensor_dx_out,
+    alpha,
+    beta
+  );
+
+  musaError_t result = musaPeekAtLastError();
+  if (result != musaSuccess) {
+    return Status::kErrorInternal;
+  }
+
+  return Status::kSuccess;
+}
+
+/// Conv3d Dgrad dispatcher - dx = dgrad(dy, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+Status Conv3dDgrad(
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_dx_in,
+  TensorRef<ElementC, LayoutC> tensor_dx_out,
+  ElementCompute alpha,
+  ElementCompute beta,
+  musaStream_t stream = nullptr) {
+
+  //
+  // Blocking factors improve performance of reference implementation
+  //
+
+  int const kThreadM = 2;       // shape of a thread's tile in the GEMM M dimension
+  int const kThreadN = 4;       // shape of a thread's tile in the GEMM N dimension
+  int const kCtaShapeM = 16;    // shape of a threadblock in units of threads
+  int const kCtaShapeN = 8;     // shape of a threadblock in units of threads
+
+  int64_t ndhw = int64_t(problem_size.N) * problem_size.D * problem_size.H * problem_size.W;
+  int64_t blocks_m = (ndhw + (kCtaShapeM * kThreadM) - 1) / (kCtaShapeM * kThreadM);
+
+  dim3 block(kCtaShapeM, kCtaShapeN);
+  dim3 grid(uint32_t(blocks_m), (problem_size.C + (kCtaShapeN * kThreadN) - 1) / (kCtaShapeN * kThreadN));
+
+  kernel::Conv3dDgrad<
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    ElementCompute,
+    ElementAccumulator,
+    ConvertOp,
+    InnerProductOp,
+    kThreadM,
+    kThreadN,
+    kCtaShapeM,
+    kCtaShapeN
+  ><<< grid, block, 0, stream >>>(
+    problem_size,
+    tensor_dy,
+    tensor_w,
+    tensor_dx_in,
+    tensor_dx_out,
+    alpha,
+    beta
+  );
+
+  musaError_t result = musaPeekAtLastError();
+  if (result != musaSuccess) {
+    return Status::kErrorInternal;
+  }
+
+  return Status::kSuccess;
+}
+
+/// Conv2d Wgrad dispatcher - dw = wgrad(dy, x)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+Status Conv2dWgrad(
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_x,
+  TensorRef<ElementC, LayoutC> tensor_dw_in,
+  TensorRef<ElementC, LayoutC> tensor_dw_out,
+  ElementCompute alpha,
+  ElementCompute beta,
+  musaStream_t stream = nullptr) {
+
+  //
+  // Blocking factors improve performance of reference implementation
+  //
+
+  int const kThreadM = 2;       // shape of a thread's tile in the GEMM M dimension
+  int const kThreadN = 4;       // shape of a thread's tile in the GEMM N dimension
+  int const kCtaShapeM = 8;     // shape of a threadblock in units of threads
+  int const kCtaShapeN = 16;    // shape of a threadblock in units of threads
+
+  int64_t rsc = int64_t(problem_size.R) * problem_size.S * problem_size.C;
+  int64_t blocks_n = (rsc + (kCtaShapeN * kThreadN) - 1) / (kCtaShapeN * kThreadN);
+
+  dim3 block(kCtaShapeM, kCtaShapeN);
+  dim3 grid((problem_size.K + (kCtaShapeM * kThreadM) - 1) / (kCtaShapeM * kThreadM), uint32_t(blocks_n));
+
+  kernel::Conv2dWgrad<
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    ElementCompute,
+    ElementAccumulator,
+    ConvertOp,
+    InnerProductOp,
+    kThreadM,
+    kThreadN,
+    kCtaShapeM,
+    kCtaShapeN
+  ><<< grid, block, 0, stream >>>(
+    problem_size,
+    tensor_dy,
+    tensor_x,
+    tensor_dw_in,
+    tensor_dw_out,
+    alpha,
+    beta
+  );
+
+  musaError_t result = musaPeekAtLastError();
+  if (result != musaSuccess) {
+    return Status::kErrorInternal;
+  }
+
+  return Status::kSuccess;
+}
+
+/// Conv3d Wgrad dispatcher - dw = wgrad(dy, x)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+Status Conv3dWgrad(
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_x,
+  TensorRef<ElementC, LayoutC> tensor_dw_in,
+  TensorRef<ElementC, LayoutC> tensor_dw_out,
+  ElementCompute alpha,
+  ElementCompute beta,
+  musaStream_t stream = nullptr) {
+
+  //
+  // Blocking factors improve performance of reference implementation
+  //
+
+  int const kThreadM = 2;       // shape of a thread's tile in the GEMM M dimension
+  int const kThreadN = 4;       // shape of a thread's tile in the GEMM N dimension
+  int const kCtaShapeM = 8;     // shape of a threadblock in units of threads
+  int const kCtaShapeN = 16;    // shape of a threadblock in units of threads
+
+  int64_t trsc = int64_t(problem_size.T) * problem_size.R * problem_size.S * problem_size.C;
+  int64_t blocks_n = (trsc + (kCtaShapeN * kThreadN) - 1) / (kCtaShapeN * kThreadN);
+
+  dim3 block(kCtaShapeM, kCtaShapeN);
+  dim3 grid((problem_size.K + (kCtaShapeM * kThreadM) - 1) / (kCtaShapeM * kThreadM), uint32_t(blocks_n));
+
+  kernel::Conv3dWgrad<
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    ElementCompute,
+    ElementAccumulator,
+    ConvertOp,
+    InnerProductOp,
+    kThreadM,
+    kThreadN,
+    kCtaShapeM,
+    kCtaShapeN
+  ><<< grid, block, 0, stream >>>(
+    problem_size,
+    tensor_dy,
+    tensor_x,
+    tensor_dw_in,
+    tensor_dw_out,
+    alpha,
+    beta
+  );
+
+  musaError_t result = musaPeekAtLastError();
+  if (result != musaSuccess) {
+    return Status::kErrorInternal;
+  }
+
+  return Status::kSuccess;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Generic 2D convolution targeting Conv2dFprop, Conv2dDgrad, and Conv2dWgrad.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+Status Conv2d(
+  conv::Operator convolutional_operator,
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_A,
+  TensorRef<ElementB, LayoutB> tensor_B,
+  TensorRef<ElementC, LayoutC> tensor_C,
+  TensorRef<ElementC, LayoutC> tensor_D,
+  ElementCompute alpha,
+  ElementCompute beta,
+  musaStream_t stream = nullptr) {
+
+  switch (convolutional_operator) {
+  case conv::Operator::kFprop:
+    return Conv2dFprop<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, stream);
+    break;
+
+  case conv::Operator::kDgrad:
+    return Conv2dDgrad<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, stream);
+    break;
+
+  case conv::Operator::kWgrad:
+    return Conv2dWgrad<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, stream);
+    break;
+
+  default: break;
+  }
+
+  return Status::kErrorNotSupported;
+}
+
+/// Generic 3D convolution targeting Conv3dFprop, Conv3dDgrad, and Conv3dWgrad.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+Status Conv3d(
+  conv::Operator convolutional_operator,
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_A,
+  TensorRef<ElementB, LayoutB> tensor_B,
+  TensorRef<ElementC, LayoutC> tensor_C,
+  TensorRef<ElementC, LayoutC> tensor_D,
+  ElementCompute alpha,
+  ElementCompute beta,
+  musaStream_t stream = nullptr) {
+
+  switch (convolutional_operator) {
+  case conv::Operator::kFprop:
+    return Conv3dFprop<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, stream);
+
+  case conv::Operator::kDgrad:
+    return Conv3dDgrad<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, stream);
+
+  case conv::Operator::kWgrad:
+    return Conv3dWgrad<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta, stream);
+
+  default: break;
+  }
+
+  return Status::kErrorNotSupported;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace device
+}  // namespace reference
+}  // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm.h
new file mode 100644
index 000000000..c83f2ef52
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm.h
@@ -0,0 +1,386 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for GEMM in device-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+
+#include "mutlass/numeric_types.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+
+#include "mutlass/util/reference/device/kernel/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename AccumulatorType,
+  typename InnerProductOp = multiply_add<AccumulatorType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_gemm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  AccumulatorType initial_accum) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  // Blocking structure potentially improves performance of reference implementation
+  // with a minor increase in complexity.
+  //
+  // Note, this reference implementation is NOT expected to approach peak performance.
+  using OutputTile = MatrixShape<4, 4>;
+
+  dim3 block(16, 8);
+
+  dim3 grid(
+    (problem_size.m() + block.x * OutputTile::kRow - 1) / (block.x * OutputTile::kRow),
+    (problem_size.n() + block.y * OutputTile::kColumn - 1) / (block.y * OutputTile::kColumn)
+  );
+
+  // Launch a GEMM kernel
+  kernel::Gemm<
+    TensorRef<ElementA, LayoutA>,
+    TensorRef<ElementB, LayoutB>,
+    TensorRef<ElementC, LayoutC>,
+    ScalarType,
+    AccumulatorType,
+    OutputTile,
+    InnerProductOp,
+    ConvertOp
+  ><<< grid, block >>>(
+    problem_size,
+    alpha,
+    tensor_a,
+    tensor_b,
+    beta,
+    tensor_c,
+    tensor_d,
+    initial_accum
+  );
+}
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// This assumes the accumulator type is the same type as the scalars.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename AccumulatorType,
+  typename InnerProductOp = multiply_add<AccumulatorType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_gemm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  AccumulatorType initial_accum) {
+
+  compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                ScalarType, AccumulatorType, InnerProductOp, ConvertOp>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_c,
+        initial_accum);
+}
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename AccumulatorType,
+  typename InnerProductOp = mutlass::arch::OpMultiplyAdd
+>
+struct Gemm;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename AccumulatorType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+            ScalarType, AccumulatorType, arch::OpMultiplyAdd> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  AccumulatorType initial_accum = AccumulatorType(0)) {
+
+    static_assert(
+      LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+      "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                  ScalarType, AccumulatorType, multiply_add<AccumulatorType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  AccumulatorType initial_accum = AccumulatorType(0)) {
+    static_assert(
+      LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+      "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                ScalarType, AccumulatorType, multiply_add<AccumulatorType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add-saturate
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename AccumulatorType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            AccumulatorType, arch::OpMultiplyAddSaturate> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  AccumulatorType initial_accum = AccumulatorType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, AccumulatorType, multiply_add<AccumulatorType>,
+                 NumericConverterClamp<ElementC, ScalarType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  AccumulatorType initial_accum = AccumulatorType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, AccumulatorType, multiply_add<AccumulatorType>,
+                 NumericConverterClamp<ElementC, ScalarType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for XOR-popc
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename AccumulatorType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            AccumulatorType, arch::OpXorPopc> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  AccumulatorType initial_accum = AccumulatorType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, AccumulatorType, xor_add<AccumulatorType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  AccumulatorType initial_accum = AccumulatorType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, AccumulatorType, xor_add<AccumulatorType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Batched GEMM
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a batch of GEMMs over a set of matrices of common dimension.
+//
+// TensorRefCollection* is a type satisfying the TensorRefCollection concept.
+//
+template <
+  typename TensorRefCollectionA,
+  typename TensorRefCollectionB,
+  typename TensorRefCollectionC,
+  typename ScalarType,
+  typename AccumulatorType,
+  typename InnerProductOp,
+  typename ConvertOp
+>
+void BatchedGemm(
+  gemm::GemmCoord problem_size,
+  int batch_count,
+  ScalarType alpha,
+  TensorRefCollectionA const& tensor_a,
+  TensorRefCollectionB const& tensor_b,
+  ScalarType beta,
+  TensorRefCollectionC &tensor_c,
+  AccumulatorType initial_accum) {
+
+  static_assert(
+    TensorRefCollectionA::kRank == 2 &&
+    TensorRefCollectionB::kRank == 2 &&
+    TensorRefCollectionC::kRank == 2, "Tensors must be of rank 2");
+
+  // Blocking structure potentially improves performance of reference implementation
+  // with a minor increase in complexity.
+  //
+  // Note, this reference implementation is NOT expected to approach peak performance.
+  using OutputTile = MatrixShape<4, 4>;
+
+  dim3 block(16, 8);
+  dim3 grid(
+    (problem_size.m() + block.x * OutputTile::kRow - 1) / (block.x * OutputTile::kRow),
+    (problem_size.n() + block.y * OutputTile::kColumn - 1) / (block.y * OutputTile::kColumn),
+    batch_count
+  );
+
+  // Launch a GEMM kernel
+  kernel::BatchedGemm<
+    TensorRefCollectionA,
+    TensorRefCollectionB,
+    TensorRefCollectionC,
+    ScalarType,
+    AccumulatorType,
+    OutputTile,
+    InnerProductOp,
+    ConvertOp
+  ><<< grid, block >>>(
+    problem_size,
+    alpha,
+    tensor_a,
+    tensor_b,
+    beta,
+    tensor_c,
+    initial_accum
+  );
+}
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+//
+// TensorRefCollection* is a type satisfying the TensorRefCollection concept.
+//
+template <
+  typename TensorRefCollectionA,
+  typename TensorRefCollectionB,
+  typename TensorRefCollectionC,
+  typename ScalarType,
+  typename AccumulatorType
+>
+void BatchedGemm(
+  gemm::GemmCoord problem_size,
+  int batch_count,
+  ScalarType alpha,
+  TensorRefCollectionA const& tensor_a,
+  TensorRefCollectionB const& tensor_b,
+  ScalarType beta,
+  TensorRefCollectionC &tensor_c) {
+
+  BatchedGemm(problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, ScalarType(0));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm_complex.h
new file mode 100644
index 000000000..2a8cda7fb
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm_complex.h
@@ -0,0 +1,351 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued GEMM in device-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/complex.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace kernel {
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ElementD = ElementC,
+  typename ConvertOp = NumericConverter<ElementD, ScalarType>,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  int kMblock = 4,
+  int kNblock = 4
+>
+__global__ void GemmComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementD, LayoutC> tensor_d,
+  ComputeType initial_accum,
+  int batch_count = 1,
+  int64_t batch_stride_A = 0,
+  int64_t batch_stride_B = 0,
+  int64_t batch_stride_C = 0,
+  int64_t batch_stride_D = 0) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  int row_block = (blockIdx.x * blockDim.x + threadIdx.x) * kMblock;
+  int col_block = (blockIdx.y * blockDim.y + threadIdx.y) * kNblock;
+  int batch_idx = blockIdx.z;
+
+  tensor_a.add_pointer_offset(batch_idx * batch_stride_A);
+  tensor_b.add_pointer_offset(batch_idx * batch_stride_B);
+  tensor_c.add_pointer_offset(batch_idx * batch_stride_C);
+  tensor_d.add_pointer_offset(batch_idx * batch_stride_D);
+
+  for (; batch_idx < batch_count; batch_idx += gridDim.z) {
+
+    // Compute matrix product using blocks
+    ComputeType accum[kMblock][kNblock];
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < kNblock; j++) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kMblock; i++) {
+        accum[i][j] = initial_accum;
+      }
+    }
+
+    for (int k_block = 0; k_block < K; ++k_block) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < kNblock; j++) {
+        MUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kMblock; i++) {
+          int row = row_block + i;
+          int col = col_block + j;
+
+          if (row < M && col < N) {
+            ElementA a = tensor_a.at(MatrixCoord(row, k_block));
+            ElementB b = tensor_b.at(MatrixCoord(k_block, col));
+
+            ComputeType a_ik = ComputeType(a);
+            ComputeType b_kj = ComputeType(b);
+
+            if (transform_a == ComplexTransform::kConjugate) {
+              a_ik = conj(a_ik);
+            }
+
+            if (transform_b == ComplexTransform::kConjugate) {
+              b_kj = conj(b_kj);
+            }
+
+            accum[i][j] = inner_product_op(a_ik, b_kj,  accum[i][j]);
+          }
+        }
+      }
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < kNblock; j++) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kMblock; i++) {
+        int row = row_block + i;
+        int col = col_block + j;
+
+        MatrixCoord coord = MatrixCoord(row, col);
+
+        if (row < M && col < N) {
+
+          tensor_d.at(coord) = convert_op(
+            alpha * ScalarType(accum[i][j]) +
+            beta * ScalarType(tensor_c.at(coord)));
+        }
+      }
+    }
+
+    tensor_a.add_pointer_offset(batch_stride_A * gridDim.z);
+    tensor_b.add_pointer_offset(batch_stride_B * gridDim.z);
+    tensor_c.add_pointer_offset(batch_stride_C * gridDim.z);
+    tensor_d.add_pointer_offset(batch_stride_D * gridDim.z);
+
+  } // for (batch_idx)
+}
+
+} // namespace kernel
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ElementD = ElementC,
+  typename ConvertOp = NumericConverter<ElementD, ScalarType>,
+  typename InnerProductOp = multiply_add<ComputeType>
+>
+void GemmComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementD, LayoutC> tensor_d,
+  ComputeType initial_accum,
+  int batch_count = 1,
+  int64_t batch_stride_A = 0,
+  int64_t batch_stride_B = 0,
+  int64_t batch_stride_C = 0,
+  int64_t batch_stride_D = 0) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  int const kMblock = 4;
+  int const kNblock = 4;
+
+  dim3 block(16, 8);
+  dim3 grid(
+    (problem_size.m() + block.x * kMblock - 1) / (block.x * kMblock),
+    (problem_size.n() + block.y * kNblock - 1) / (block.y * kNblock),
+    batch_count % std::numeric_limits<uint16_t>::max()
+  );
+
+  if (grid.y <= std::numeric_limits<uint16_t>::max()) {
+    kernel::GemmComplex<
+      ElementA,
+      LayoutA,
+      ElementB,
+      LayoutB,
+      ElementC,
+      LayoutC,
+      ScalarType,
+      ComputeType,
+      ElementD,
+      ConvertOp,
+      InnerProductOp,
+      kMblock,
+      kNblock
+    ><<< grid, block >>>(
+      problem_size,
+      alpha,
+      tensor_a,
+      transform_a,
+      tensor_b,
+      transform_b,
+      beta,
+      tensor_c,
+      tensor_d,
+      initial_accum,
+      batch_count,
+      batch_stride_A,
+      batch_stride_B,
+      batch_stride_C,
+      batch_stride_D
+    );
+  } else {
+    // Using bigger thread tile size
+    int const kBigMblock = 4;
+    int const kBigNblock = 16;
+
+    dim3 Bigblock(16, 8);
+    dim3 Biggrid(
+      (problem_size.m() + block.x * kBigMblock - 1) / (block.x * kBigMblock),
+      (problem_size.n() + block.y * kBigNblock - 1) / (block.y * kBigNblock),
+      batch_count % std::numeric_limits<uint16_t>::max()
+    );
+
+    kernel::GemmComplex<
+      ElementA,
+      LayoutA,
+      ElementB,
+      LayoutB,
+      ElementC,
+      LayoutC,
+      ScalarType,
+      ComputeType,
+      ElementD,
+      ConvertOp,
+      InnerProductOp,
+      kBigMblock,
+      kBigNblock
+    ><<< Biggrid, Bigblock >>>(
+      problem_size,
+      alpha,
+      tensor_a,
+      transform_a,
+      tensor_b,
+      transform_b,
+      beta,
+      tensor_c,
+      tensor_d,
+      initial_accum,
+      batch_count,
+      batch_stride_A,
+      batch_stride_B,
+      batch_stride_C,
+      batch_stride_D
+    );
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// This assumes the accumulator type is the same type as the scalars.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ElementD = ElementC
+>
+void GemmComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementD, LayoutC> tensor_d) {
+
+  GemmComplex(problem_size, alpha, tensor_a, transform_a, tensor_b, transform_b, beta, tensor_c, tensor_d, ScalarType(0));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm_planar_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm_planar_complex.h
new file mode 100644
index 000000000..098277649
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gemm_planar_complex.h
@@ -0,0 +1,312 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued GEMM in device code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/complex.h"
+#include "mutlass/matrix_coord.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_ref_planar_complex.h"
+
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace kernel {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static int const kGemmPlanarComplexBlockSize = 4;
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>,
+  typename InnerProductOp = multiply_add<complex<ComputeType>>
+>
+__global__ void GemmPlanarComplex(
+  gemm::GemmCoord problem_size,
+  complex<ScalarType> alpha,
+  TensorRefPlanarComplex<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRefPlanarComplex<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  complex<ScalarType> beta,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_c,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_d,
+  complex<ComputeType> initial_accum) {
+
+  int const kMblock = kGemmPlanarComplexBlockSize;
+  int const kNblock = kGemmPlanarComplexBlockSize;
+
+  using ComplexA = typename TensorRefPlanarComplex<ElementA, LayoutA>::ComplexElement;
+  using ComplexB = typename TensorRefPlanarComplex<ElementB, LayoutB>::ComplexElement;
+  using ComplexC = typename TensorRefPlanarComplex<ElementC, LayoutC>::ComplexElement;
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  complex<ComputeType> accum[kMblock][kNblock];
+
+  int row_block = (blockIdx.x * blockDim.x + threadIdx.x) * kMblock;
+  int col_block = (blockIdx.y * blockDim.y + threadIdx.y) * kNblock;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int j = 0; j < kNblock; j++) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kMblock; i++) {
+      accum[i][j] = initial_accum;
+    }
+  }
+
+  MUTLASS_PRAGMA_NO_UNROLL
+  for (int k_block = 0; k_block < K; ++k_block) {
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < kNblock; j++) {
+
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kMblock; i++) {
+
+        int row = row_block + i;
+        int col = col_block + j;
+
+        if (row < M && col < N) {
+
+          ComplexA a_ik = tensor_a.at(MatrixCoord(row, k_block));
+          ComplexB b_kj = tensor_b.at(MatrixCoord(k_block, col));
+
+          complex<ComputeType> a = complex<ComputeType>{
+            ComputeType(a_ik.real()),
+            ComputeType(a_ik.imag())
+          };
+
+          complex<ComputeType> b = complex<ComputeType>{
+            ComputeType(b_kj.real()),
+            ComputeType(b_kj.imag())
+          };
+
+          if (transform_a == ComplexTransform::kConjugate) {
+            a = conj(a);
+          }
+
+          if (transform_b == ComplexTransform::kConjugate) {
+            b = conj(b);
+          }
+
+          accum[i][j] = inner_product_op(a, b,  accum[i][j]);
+        }
+      }
+    }
+  }
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int j = 0; j < kNblock; j++) {
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < kMblock; i++) {
+
+      int row = row_block + i;
+      int col = col_block + j;
+
+      MatrixCoord coord = MatrixCoord(row, col);
+
+      if (row < M && col < N) {
+
+        complex<ScalarType> acc{
+          ScalarType(accum[i][j].real()),
+          ScalarType(accum[i][j].imag())
+        };
+
+        ComplexC c_ij = ComplexC();
+
+        if (beta.real() != ScalarType() || beta.imag() != ScalarType()) {
+          c_ij = tensor_c.at(coord);
+        }
+
+        complex<ScalarType> src{
+          ScalarType(c_ij.real()),
+          ScalarType(c_ij.imag())
+        };
+
+        complex<ScalarType> result = alpha * acc + beta * src;
+
+        ComplexC d_ij;
+
+        d_ij.real() = convert_op(result.real());
+        d_ij.imag() = convert_op(result.imag());
+
+        tensor_d.at(coord) = d_ij;
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace kernel
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>,
+  typename InnerProductOp = multiply_add<complex<ComputeType>>
+>
+void GemmPlanarComplex(
+  gemm::GemmCoord problem_size,
+  complex<ScalarType> alpha,
+  TensorRefPlanarComplex<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRefPlanarComplex<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  complex<ScalarType> beta,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_c,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_d,
+  complex<ComputeType> initial_accum) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  int const kMblock = kernel::kGemmPlanarComplexBlockSize;
+  int const kNblock = kernel::kGemmPlanarComplexBlockSize;
+
+  dim3 block(16, 8);
+
+  dim3 grid(
+    (problem_size.m() + block.x * kMblock - 1) / (block.x * kMblock),
+    (problem_size.n() + block.y * kNblock - 1) / (block.y * kNblock),
+    1);
+
+  kernel::GemmPlanarComplex<
+    ElementA, LayoutA,
+    ElementB, LayoutB,
+    ElementC, LayoutC,
+    ScalarType,
+    ComputeType,
+    ConvertOp,
+    InnerProductOp
+  ><<< grid, block >>>(
+    problem_size,
+    alpha,
+    tensor_a,
+    transform_a,
+    tensor_b,
+    transform_b,
+    beta,
+    tensor_c,
+    tensor_d,
+    initial_accum
+  );
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// This assumes the accumulator type is the same type as the scalars.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType
+>
+void GemmPlanarComplex(
+  gemm::GemmCoord problem_size,
+  complex<ScalarType> alpha,
+  TensorRefPlanarComplex<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRefPlanarComplex<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  complex<ScalarType> beta,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_c,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_d) {
+
+  GemmPlanarComplex(
+    problem_size,
+    alpha,
+    tensor_a, transform_a,
+    tensor_b, transform_b,
+    beta,
+    tensor_c,
+    tensor_d,
+    complex<ScalarType>());
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace reference
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gett.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gett.hpp
new file mode 100644
index 000000000..6effc63e8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/gett.hpp
@@ -0,0 +1,147 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 GETT device reference code
+*/
+#pragma once
+
+#include <mute/tensor.hpp>
+
+namespace mutlass::reference::device {
+
+template <
+  class ATensor,
+  class BTensor,
+  class CTensor,
+  class DTensor,
+  class ElementAccumulator,
+  class ElementEpilogue>
+__global__ static
+void
+gett_kernel(
+  DTensor       D,
+  ATensor const A,
+  BTensor const B,
+  CTensor const C,
+  ElementEpilogue alpha, ElementEpilogue beta,
+  ElementAccumulator acc_init)
+{
+  using namespace mute;
+
+  static_assert(DTensor::rank == 3, "(M,N,L)");
+  static_assert(ATensor::rank == 3, "(M,K,L)");
+  static_assert(BTensor::rank == 3, "(N,K,L)");
+  static_assert(CTensor::rank == 3, "(M,N,L)");
+
+  assert(size<0>(A) == size<0>(D));  // M
+  assert(size<0>(C) == size<0>(D));  // M
+  assert(size<0>(B) == size<1>(D));  // N
+  assert(size<1>(C) == size<1>(D));  // N
+  assert(size<1>(A) == size<1>(B));  // K
+  assert(size<2>(A) == size<2>(D));  // L
+  assert(size<2>(B) == size<2>(D));  // L
+  assert(size<2>(C) == size<2>(D));  // L
+
+  NumericConverter<ElementAccumulator, typename ATensor::value_type> a_converter;
+  NumericConverter<ElementAccumulator, typename BTensor::value_type> b_converter;
+  NumericConverter<ElementEpilogue, ElementAccumulator> acc_converter;
+  NumericConverter<ElementEpilogue, typename CTensor::value_type> source_converter;
+  NumericConverter<typename DTensor::value_type, ElementEpilogue> output_converter;
+
+  // Thread id to each element of D
+  for (int tid = threadIdx.x + blockDim.x * blockIdx.x;
+       tid < size(D);
+       tid += blockDim.x * gridDim.x) {
+    // (m,n,l) coordinate
+    auto mnl_coord = idx2crd(tid, product_each(shape(D)));
+    auto m = get<0>(mnl_coord);
+    auto n = get<1>(mnl_coord);
+    auto l = get<2>(mnl_coord);
+
+    auto A_ml = A(m,_,l);
+    auto B_nl = B(n,_,l);
+
+    ElementAccumulator accum = ElementAccumulator(0);
+    for (int k = 0; k < size<1>(A); ++k) {
+      ElementAccumulator a = a_converter(A_ml(k));
+      ElementAccumulator b = b_converter(B_nl(k));
+      accum += a * b;
+    }
+
+    ElementEpilogue scaled_output = (alpha * acc_converter(accum)) + (beta * source_converter(C(m,n,l)));
+    D(m,n,l) = output_converter(scaled_output);
+  }
+}
+
+// Most general version
+template <
+  class ProblemShapeMNKL,
+  class ElementA,
+  class StrideA,
+  class ElementB,
+  class StrideB,
+  class ElementAccumulator,
+  class ElementC,
+  class StrideC,
+  class ElementD,
+  class StrideD,
+  class ElementEpilogue>
+void
+gett(
+    ProblemShapeMNKL problem_shape_mnkl,
+    ElementA const* ptr_A, StrideA stride_a_mkl,
+    ElementB const* ptr_B, StrideB stride_b_nkl,
+    ElementAccumulator _,
+    ElementC const* ptr_C, StrideC stride_c_mnl,
+    ElementD      * ptr_D, StrideD stride_d_mnl,
+    ElementEpilogue alpha, ElementEpilogue beta,
+    musaStream_t stream = 0) {
+  using namespace mute;
+
+  static_assert(rank(ProblemShapeMNKL{}) == 4);
+  auto M = get<0>(problem_shape_mnkl);
+  auto N = get<1>(problem_shape_mnkl);
+  auto K = get<2>(problem_shape_mnkl);
+  auto L = get<3>(problem_shape_mnkl);
+
+  // Represent the full tensors
+  auto A = make_tensor(make_gmem_ptr(ptr_A), make_shape(M,K,L), stride_a_mkl); // (M,K,L)
+  auto B = make_tensor(make_gmem_ptr(ptr_B), make_shape(N,K,L), stride_b_nkl); // (N,K,L)
+  auto C = make_tensor(make_gmem_ptr(ptr_C), make_shape(M,N,L), stride_c_mnl); // (M,N,L)
+  auto D = make_tensor(make_gmem_ptr(ptr_D), make_shape(M,N,L), stride_d_mnl); // (M,N,L)
+
+  dim3 dimBlock(256);
+  dim3 dimGrid(240);
+  gett_kernel<<< dimGrid, dimBlock, 0, stream >>>(D, A, B, C, alpha, beta, ElementAccumulator(0));
+}
+
+} // namespace mutlass::reference::device
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/gemm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/gemm.h
new file mode 100644
index 000000000..b0541e764
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/gemm.h
@@ -0,0 +1,163 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for GEMM in host-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+
+#include "mutlass/util/reference/device/thread/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+namespace kernel {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename TensorRefA,
+  typename TensorRefB,
+  typename TensorRefC,
+  typename ScalarType,
+  typename AccumulatorType,
+  typename OutputTile,
+  typename InnerProductOp,
+  typename ConvertOp
+>
+__global__ void Gemm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRefA tensor_a,
+  TensorRefB tensor_b,
+  ScalarType beta,
+  TensorRefC tensor_c,
+  TensorRefC tensor_d,
+  AccumulatorType initial_accum) {
+
+  // Map each thread to a unique tile of the output matrix
+  MatrixCoord output_coord(
+    MatrixCoord::Index((threadIdx.x + blockIdx.x * blockDim.x) * OutputTile::kRow),
+    MatrixCoord::Index((threadIdx.y + blockIdx.y * blockDim.y) * OutputTile::kColumn)
+  );
+
+  // Compute the general matrix product
+  thread::Gemm<
+    TensorRefA,
+    TensorRefB,
+    TensorRefC,
+    ScalarType,
+    AccumulatorType,
+    OutputTile,
+    InnerProductOp,
+    ConvertOp
+  > gemm(initial_accum);
+
+  gemm.multiply_add(
+    problem_size,
+    tensor_a,
+    tensor_b,
+    output_coord);
+
+  gemm.epilogue(problem_size, alpha, beta, tensor_c, tensor_d, output_coord);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename TensorRefCollectionA,
+  typename TensorRefCollectionB,
+  typename TensorRefCollectionC,
+  typename ScalarType,
+  typename AccumulatorType,
+  typename OutputTile,
+  typename InnerProductOp,
+  typename ConvertOp
+>
+__global__ void BatchedGemm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRefCollectionA tensor_collection_a,
+  TensorRefCollectionB tensor_collection_b,
+  ScalarType beta,
+  TensorRefCollectionC tensor_collection_c,
+  AccumulatorType initial_accum) {
+
+  // Obtain batch ID
+  int batch_id = blockIdx.z;
+
+  // Dereference based on batch_id
+  typename TensorRefCollectionA::TensorRef tensor_a = tensor_collection_a.at(batch_id);
+  typename TensorRefCollectionB::TensorRef tensor_b = tensor_collection_b.at(batch_id);
+  typename TensorRefCollectionC::TensorRef tensor_c = tensor_collection_c.at(batch_id);
+
+  // Map each thread to a unique tile of the output matrix
+  MatrixCoord output_coord(
+    (threadIdx.x + blockIdx.x * blockDim.x) * OutputTile::kColumn,
+    (threadIdx.y + blockIdx.y * blockDim.y) * OutputTile::kRow
+  );
+
+  // Compute the general matrix product
+  thread::Gemm<
+    typename TensorRefCollectionA::TensorRef,
+    typename TensorRefCollectionB::TensorRef,
+    typename TensorRefCollectionC::TensorRef,
+    ScalarType,
+    AccumulatorType,
+    OutputTile,
+    InnerProductOp,
+    ConvertOp
+  > gemm(initial_accum);
+
+  gemm.multiply_add(
+    problem_size,
+    tensor_a,
+    tensor_b,
+    output_coord);
+
+  gemm.epilogue(problem_size, alpha, beta, tensor_c, output_coord);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace kernel
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/tensor_elementwise.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/tensor_elementwise.h
new file mode 100644
index 000000000..392ea60d7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/tensor_elementwise.h
@@ -0,0 +1,169 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <curand_kernel.h>
+
+#include "mutlass/mutlass.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+namespace kernel {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Kernel to initialize tensor to uniform random distribution
+template <typename T>
+__global__ void TensorInitializeUniform(
+    Distribution dist, int64_t seed, int dim_contiguous, int dim_strided, T *tensor, int ldm) {
+  __shared__ curandState_t rng_state[1024];
+
+  uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x + blockIdx.y * gridDim.x * blockDim.x;
+
+  curand_init(seed, gtid, 0, &rng_state[threadIdx.x]);
+
+  int c_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int s_idx = blockIdx.y * blockDim.x;
+
+  tensor += s_idx * ldm + c_idx;
+
+  for (int s_offset = 0; s_offset < blockDim.x; ++s_offset, ++s_idx) {
+    if (s_idx < dim_strided && c_idx < dim_contiguous) {
+      double range = dist.uniform.max - dist.uniform.min;
+
+      double rnd = curand_uniform(&rng_state[threadIdx.x]);
+
+      rnd = dist.uniform.min + range * rnd;
+
+      // Random values are cast to integer after scaling by a power of two to facilitate error
+      // testing
+      if (dist.int_scale >= 0) {
+        rnd = double(int(rnd * double(1 << dist.int_scale)));
+        *tensor = T(rnd / double(1 << dist.int_scale));
+      } else {
+        *tensor = T(rnd);
+      }
+
+      tensor += ldm;
+    }
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Kernel to initialize tensor to uniform distribution
+template <typename T>
+__global__ void TensorInitializeGaussian(
+    Distribution dist, int64_t seed, int dim_contiguous, int dim_strided, T *tensor, int ldm) {
+  __shared__ curandState_t rng_state[1024];
+
+  uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x + blockIdx.y * gridDim.x * blockDim.x;
+
+  curand_init(seed, gtid, 0, &rng_state[threadIdx.x]);
+
+  int c_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int s_idx = blockIdx.y * blockDim.x;
+
+  tensor += s_idx * ldm + c_idx;
+
+  for (int s_offset = 0; s_offset < blockDim.x; ++s_offset, ++s_idx) {
+    if (s_idx < dim_strided && c_idx < dim_contiguous) {
+      // Random values are cast to integer after scaling by a power of two to facilitate error
+      // testing
+
+      double rnd = curand_normal(&rng_state[threadIdx.x]);
+
+      rnd = dist.gaussian.mean + dist.gaussian.stddev * rnd;
+
+      if (dist.int_scale >= 0) {
+        rnd = double(int(rnd * double(1 << dist.int_scale)));
+        *tensor = T(rnd / double(1 << dist.int_scale));
+      } else {
+        *tensor = T(rnd);
+      }
+    }
+  }
+}
+
+/// Kernel to initialize tensor to an identity matrix
+template <typename T>
+__global__ void TensorInitializeLinear(
+    Distribution dist, int64_t seed, int dim_contiguous, int dim_strided, T *tensor, int ldm) {
+  __shared__ curandState_t rng_state[1024];
+
+  uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x + blockIdx.y * gridDim.x * blockDim.x;
+
+  curand_init(seed, gtid, 0, &rng_state[threadIdx.x]);
+
+  int c_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int s_idx = blockIdx.y * blockDim.x;
+
+  tensor += s_idx * ldm + c_idx;
+
+  for (int s_offset = 0; s_offset < blockDim.x; ++s_offset, ++s_idx) {
+    if (s_idx < dim_strided && c_idx < dim_contiguous) {
+      *tensor =
+          dist.linear.offset + dist.linear.delta_row * c_idx + dist.linear.delta_column * s_idx;
+    }
+  }
+}
+
+/// Kernel to initialize tensor to an identity matrix
+template <typename T>
+__global__ void TensorInitializeIdentity(
+    Distribution dist, int64_t seed, int dim_contiguous, int dim_strided, T *tensor, int ldm) {
+  __shared__ curandState_t rng_state[1024];
+
+  uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x + blockIdx.y * gridDim.x * blockDim.x;
+
+  curand_init(seed, gtid, 0, &rng_state[threadIdx.x]);
+
+  int c_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int s_idx = blockIdx.y * blockDim.x;
+
+  tensor += s_idx * ldm + c_idx;
+
+  for (int s_offset = 0; s_offset < blockDim.x; ++s_offset, ++s_idx) {
+    if (s_idx < dim_strided && c_idx < dim_contiguous) {
+      *tensor = (c_idx == s_idx ? T(1) : T(0));
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace kernel
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/tensor_foreach.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/tensor_foreach.h
new file mode 100644
index 000000000..3d1991570
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/kernel/tensor_foreach.h
@@ -0,0 +1,159 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include "mutlass/mutlass.h"
+#include "mutlass/coord.h"
+#include "mutlass/subbyte_reference.h"
+#include "mutlass/fast_math.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+namespace kernel {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines several helpers
+namespace detail {
+
+/// Helper to perform for-each operation
+template <typename Func, int Rank, int RankRemaining>
+struct TensorForEachHelper {
+
+  /// Constructor for general rank
+  __inline__ __device__
+  TensorForEachHelper(Func &func, Coord<Rank> const &size, Coord<Rank> &coord, int64_t index) {
+
+    int64_t product = 1;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = Rank - RankRemaining; i < Rank; ++i) {
+      product *= size[i];
+    }
+
+    coord[Rank - 1 - RankRemaining] = index / product;
+    int64_t remaining = index % product;
+
+    TensorForEachHelper<Func, Rank, RankRemaining-1>(func, size, coord, remaining);
+  }
+};
+
+/// Helper to perform for-each operation
+template <typename Func, int Rank>
+struct TensorForEachHelper<Func, Rank, 0> {
+
+  /// Constructor for fastest changing rank
+  __inline__ __device__
+  TensorForEachHelper(Func &func, Coord<Rank> const &size, Coord<Rank> &coord, int64_t index) {
+
+    coord[Rank - 1] = index;
+
+    if (coord < size) {
+      func(coord);
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Kernel calls a functor for each element in a tensor's index space
+template <typename Func, int Rank, typename Params>
+__global__ void TensorForEach(Coord<Rank> size, Params params = Params()) {
+
+  Func func(params);
+
+  int64_t index = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t max_index = 1;
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < Rank; ++i) {
+    max_index *= size[i];
+  }
+
+  MUTLASS_PRAGMA_NO_UNROLL
+  while  (index < max_index) {
+    Coord<Rank> coord;
+
+    detail::TensorForEachHelper<Func, Rank, Rank - 1>(func, size, coord, index);
+    index += blockDim.x * gridDim.x;
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Kernel calls a functor for each element along a tensor's diagonal
+template <typename Func, int Rank, typename Params>
+__global__ void TensorDiagonalForEach(Coord<Rank> size, Params params, int start, int end) {
+
+  Func func(params);
+
+  int64_t index = threadIdx.x + blockIdx.x * blockDim.x + start;
+
+  if (index < end) {
+    Coord<Rank> coord;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Rank; ++i) {
+      coord[i] = index;
+    }
+
+    func(coord);
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Element, typename Func>
+__global__ void BlockForEach(
+  Element *ptr,
+  size_t capacity,
+  typename Func::Params params) {
+
+  Func func(params);
+
+  size_t index = threadIdx.x + blockIdx.x * blockDim.x;
+
+  for (; index < capacity; index += blockDim.x * gridDim.x) {
+    ReferenceFactory<Element>::get(ptr, index) = func();
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace kernel
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/rank_2k_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/rank_2k_complex.h
new file mode 100644
index 000000000..0c6c4205e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/rank_2k_complex.h
@@ -0,0 +1,356 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued GEMM in device-side code.
+*/
+
+#pragma once
+
+#include "mutlass/blas3.h"
+#include "mutlass/complex.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace kernel {
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  int kMblock = 4,
+  int kNblock = 4
+>
+__global__ void Rank2KComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum,
+  FillMode fill_mode_c,
+  BlasMode blas_mode,
+  int batch_count = 1,
+  int64_t batch_stride_A = 0,
+  int64_t batch_stride_B = 0,
+  int64_t batch_stride_C = 0,
+  int64_t batch_stride_D = 0) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  assert(M=N);
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  int row_block = (blockIdx.x * blockDim.x + threadIdx.x) * kMblock;
+  int col_block = (blockIdx.y * blockDim.y + threadIdx.y) * kNblock;
+  int batch_idx = blockIdx.z;
+
+  tensor_a.add_pointer_offset(batch_idx * batch_stride_A);
+  tensor_b.add_pointer_offset(batch_idx * batch_stride_B);
+  tensor_c.add_pointer_offset(batch_idx * batch_stride_C);
+  tensor_d.add_pointer_offset(batch_idx * batch_stride_D);
+
+  for (; batch_idx < batch_count; batch_idx += gridDim.z) {
+
+    // Compute matrix product using blocks
+    ComputeType accum[kMblock][kNblock];
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < kNblock; j++) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kMblock; i++) {
+        accum[i][j] = initial_accum;
+      }
+    }
+
+    for (int k_block = 0; k_block < K; ++k_block) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < kNblock; j++) {
+        MUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < kMblock; i++) {
+          int row = row_block + i;
+          int col = col_block + j;
+
+          if (row < M && col < N &&
+             ( (fill_mode_c == FillMode::kLower && row >= col) ||
+              (fill_mode_c == FillMode::kUpper && row <= col) )
+            ) {
+
+            // A x B^T (Symmetric) or A x B^H (Hermitian)
+            // complex conjugation on operandB (b_t) is function of blas3 computation
+            ElementA a = tensor_a.at(MatrixCoord(row, k_block));
+            ElementB b_t = (blas_mode == BlasMode::kHermitian) ?
+                          conj(tensor_b.at(MatrixCoord(col, k_block))) :
+                          tensor_b.at(MatrixCoord(col, k_block));
+
+            ComputeType a_ik = ComputeType(a);
+            ComputeType b_jk = ComputeType(b_t);
+
+            // complex conjugation is a function of operand layouts
+            if (transform_a == ComplexTransform::kConjugate) {
+              a_ik = conj(a_ik);
+            }
+            // complex conjugation is a function of operand layouts
+            if (transform_b == ComplexTransform::kConjugate) {
+              b_jk = conj(b_jk);
+            }
+
+            accum[i][j] = inner_product_op(a_ik, b_jk,  accum[i][j]);
+
+            // B x A^T (Symmetric) or B x A^H (Hermitian)
+            // complex conjugation on operandB (a_t) is function of blas3 computation
+            ElementB b = tensor_b.at(MatrixCoord(row, k_block));
+            ElementA a_t = (blas_mode == BlasMode::kHermitian) ?
+                            conj(tensor_a.at(MatrixCoord(col, k_block))):
+                            tensor_a.at(MatrixCoord(col, k_block));
+
+            ComputeType b_ik = ComputeType(b);
+            ComputeType a_jk = ComputeType(a_t);
+
+            // complex conjugation here is a function of operand layouts
+            if (transform_b == ComplexTransform::kConjugate) {
+              b_ik = conj(b_ik);
+            }
+            // complex conjugation here is a function of operand layouts
+            if (transform_a == ComplexTransform::kConjugate) {
+              a_jk = conj(a_jk);
+            }
+
+            accum[i][j] = inner_product_op(a_ik, b_kj,  accum[i][j]);
+          }
+        }
+      }
+    }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < kNblock; j++) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < kMblock; i++) {
+        int row = row_block + i;
+        int col = col_block + j;
+
+        MatrixCoord coord = MatrixCoord(row, col);
+
+        if (row < M && col < N &&
+            ((fill_mode_c == FillMode::kLower && row >= col) ||
+             (fill_mode_c == FillMode::kUpper && row <= col))
+          ) {
+
+          ScalarType c = tensor_c.at(coord);
+          // The imaginary parts of the diagonal elements of
+          // a complex data type are assumed and set to zero
+          if (blas_mode == BlasMode::kHermitian) {
+            c = (row == col) ? real(c) : c;
+          }
+
+          tensor_d.at(coord) = convert_op(
+            alpha * ScalarType(accum[i][j]) +
+            beta * c);
+        }
+      }
+    }
+
+    tensor_a.add_pointer_offset(batch_stride_A * gridDim.z);
+    tensor_b.add_pointer_offset(batch_stride_B * gridDim.z);
+    tensor_c.add_pointer_offset(batch_stride_C * gridDim.z);
+    tensor_d.add_pointer_offset(batch_stride_D * gridDim.z);
+
+  } // for (batch_idx)
+}
+
+} // namespace kernel
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>,
+  typename InnerProductOp = multiply_add<ComputeType>
+>
+void Rank2KComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum,
+  FillMode fill_mode_c,
+  BlasMode blas_mode,
+  int batch_count = 1,
+  int64_t batch_stride_A = 0,
+  int64_t batch_stride_B = 0,
+  int64_t batch_stride_C = 0,
+  int64_t batch_stride_D = 0) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  int const kMblock = 4;
+  int const kNblock = 4;
+
+  dim3 block(16, 8);
+  dim3 grid(
+    (problem_size.m() + block.x * kMblock - 1) / (block.x * kMblock),
+    (problem_size.n() + block.y * kNblock - 1) / (block.y * kNblock),
+    batch_count % std::numeric_limits<uint16_t>::max()
+  );
+
+  kernel::Rank2KComplex<
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    ElementC,
+    LayoutC,
+    ScalarType,
+    ComputeType,
+    ConvertOp,
+    InnerProductOp,
+    kMblock,
+    kNblock
+  ><<< grid, block >>>(
+    problem_size,
+    alpha,
+    tensor_a,
+    transform_a,
+    tensor_b,
+    transform_b,
+    beta,
+    tensor_c,
+    tensor_d,
+    initial_accum,
+    fill_mode_c,
+    blas_mode,
+    batch_count,
+    batch_stride_A,
+    batch_stride_B,
+    batch_stride_C,
+    batch_stride_D
+  );
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// This assumes the accumulator type is the same type as the scalars.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType
+>
+void Rank2KComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  FillMode fill_mode_c,
+  BlasMode blas_mode) {
+
+  Rank2KComplex(
+    problem_size, alpha,
+    tensor_a, transform_a,
+    tensor_b, transform_b,
+    beta, tensor_c, tensor_d,
+    ScalarType(0),
+    fill_mode_c,
+    blas_mode);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_compare.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_compare.h
new file mode 100644
index 000000000..c55308f1c
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_compare.h
@@ -0,0 +1,247 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 host-side elementwise operations on TensorView.
+*/
+
+#pragma once
+// Standard Library includes
+#include <utility>
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/relatively_equal.h"
+
+#include "mutlass/util/distribution.h"
+
+#include "tensor_foreach.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace kernel {
+
+template <typename Element>
+__global__ void BlockCompareEqual(
+  int *equal,
+  Element const *ptr_A,
+  Element const *ptr_B,
+  size_t capacity) {
+
+  size_t idx = threadIdx.x + blockDim.x * blockIdx.x;
+
+  for (; idx < capacity; idx += gridDim.x * blockDim.x) {
+
+    Element a = mutlass::ReferenceFactory<Element>::get(ptr_A, idx);
+    Element b = mutlass::ReferenceFactory<Element>::get(ptr_B, idx);
+
+    if (a != b) {
+      *equal = 0;
+
+      return;
+    }
+  }
+}
+
+template <typename Element>
+__global__ void BlockCompareRelativelyEqual(
+  int *equal,
+  Element const *ptr_A,
+  Element const *ptr_B,
+  size_t capacity,
+  Element epsilon,
+  Element nonzero_floor) {
+
+  size_t idx = threadIdx.x + blockDim.x * blockIdx.x;
+
+  for (; idx < capacity; idx += gridDim.x * blockDim.x) {
+
+    Element a = mutlass::ReferenceFactory<Element>::get(ptr_A, idx);
+    Element b = mutlass::ReferenceFactory<Element>::get(ptr_B, idx);
+
+    if (!relatively_equal(a, b, epsilon, nonzero_floor)) {
+      *equal = 0;
+      return;
+    }
+  }
+}
+
+} // namespace kernel
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Performs a bit-level equality check between two blocks
+template <typename Element>
+bool BlockCompareEqual(
+  Element const *ptr_A,
+  Element const *ptr_B,
+  size_t capacity,
+  int grid_size = 0,
+  int block_size = 0) {
+
+  int equal_flag = 1;
+  int *device_equal_flag = nullptr;
+
+  if (musaMalloc((void **)&device_equal_flag, sizeof(int)) != musaSuccess) {
+    throw std::runtime_error("Failed to allocate device flag.");
+  }
+
+  if (musaMemcpy(
+    device_equal_flag,
+    &equal_flag,
+    sizeof(int),
+    musaMemcpyHostToDevice) != musaSuccess) {
+
+    throw std::runtime_error("Failed to copy equality flag to device.");
+  }
+
+  if (!grid_size || !block_size) {
+
+    // if grid_size or block_size are zero, query occupancy using the MUSA Occupancy API
+    musaError_t result = musaOccupancyMaxPotentialBlockSize(
+      &grid_size,
+      &block_size,
+      reinterpret_cast<void const *>(kernel::BlockCompareEqual<Element>));
+
+    if (result != musaSuccess) {
+      throw std::runtime_error("Failed to query occupancy.");
+    }
+
+    // Limit block size. This has the effect of increasing the number of items processed by a
+    // single thread and reduces the impact of initialization overhead.
+    block_size = (block_size < 128 ? block_size : 128);
+  }
+
+  dim3 grid(grid_size, 1, 1);
+  dim3 block(block_size, 1, 1);
+
+  kernel::BlockCompareEqual<Element><<< grid, block >>>(device_equal_flag, ptr_A, ptr_B, capacity);
+
+  if (musaMemcpy(
+    &equal_flag,
+    device_equal_flag,
+    sizeof(int),
+    musaMemcpyDeviceToHost) != musaSuccess) {
+
+    musaFree(device_equal_flag);
+
+    throw std::runtime_error("Failed to copy equality flag from device.");
+  }
+
+  musaFree(device_equal_flag);
+
+  return equal_flag;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Performs a bit-level equality check between two blocks
+template <typename Element>
+bool BlockCompareRelativelyEqual(
+  Element const *ptr_A,
+  Element const *ptr_B,
+  size_t capacity,
+  Element epsilon,
+  Element nonzero_floor,
+  int grid_size = 0,
+  int block_size = 0) {
+
+  int equal_flag = 1;
+  int *device_equal_flag = nullptr;
+
+  if (musaMalloc((void **)&device_equal_flag, sizeof(int)) != musaSuccess) {
+    throw std::runtime_error("Failed to allocate device flag.");
+  }
+
+  if (musaMemcpy(
+    device_equal_flag,
+    &equal_flag,
+    sizeof(int),
+    musaMemcpyHostToDevice) != musaSuccess) {
+
+    throw std::runtime_error("Failed to copy equality flag to device.");
+  }
+
+  if (!grid_size || !block_size) {
+
+    // if grid_size or block_size are zero, query occupancy using the MUSA Occupancy API
+    musaError_t result = musaOccupancyMaxPotentialBlockSize(
+      &grid_size,
+      &block_size,
+      reinterpret_cast<void const *>(kernel::BlockCompareRelativelyEqual<Element>));
+
+    if (result != musaSuccess) {
+      throw std::runtime_error("Failed to query occupancy.");
+    }
+
+    // Limit block size. This has the effect of increasing the number of items processed by a
+    // single thread and reduces the impact of initialization overhead.
+    block_size = (block_size < 128 ? block_size : 128);
+  }
+
+  dim3 grid(grid_size, 1, 1);
+  dim3 block(block_size, 1, 1);
+
+  kernel::BlockCompareRelativelyEqual<Element><<< grid, block >>>(
+    device_equal_flag,
+    ptr_A,
+    ptr_B,
+    capacity,
+    epsilon,
+    nonzero_floor
+  );
+
+  if (musaMemcpy(
+    &equal_flag,
+    device_equal_flag,
+    sizeof(int),
+    musaMemcpyDeviceToHost) != musaSuccess) {
+
+    musaFree(device_equal_flag);
+
+    throw std::runtime_error("Failed to copy equality flag from device.");
+  }
+
+  musaFree(device_equal_flag);
+
+  return equal_flag;
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // device
+} // reference
+} // mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_fill.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_fill.h
new file mode 100644
index 000000000..2d4f3cb53
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_fill.h
@@ -0,0 +1,2078 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright notice,
+ * this list of conditions and the following disclaimer in the domumentation
+ * and/or other materials provided with the distribution.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 device-side elementwise operations on TensorView. Note, the operations defined
+    in this header are not specialized for any partimular data layout and are therefore not
+    intended to offer the best possible performance. Rather, they are intended to be generic
+    reference implementations to support the MUTLASS unit tests.
+*/
+
+#pragma once
+
+#if !defined(__MUSACC_RTC__)
+
+// Standard Library includes
+#include <utility>
+#include <cstdlib>
+#include <cmath>
+#include <type_traits>
+#include <cstdint>
+
+#endif
+
+// MUSA includes
+#include <murand_kernel.h>
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/array.h"
+#include "mutlass/complex.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/blas3.h"
+#include "mutlass/numeric_types.h"
+
+#include "mutlass/layout/vector.h"
+
+#include "mutlass/util/reference/device/tensor_foreach.h"
+#include "mutlass/util/distribution.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <typename FloatType>
+MUTLASS_DEVICE
+FloatType random_normal_float(murandState_t *state) {
+  return murand_normal(state);
+}
+
+template <>
+MUTLASS_DEVICE
+double random_normal_float<double>(murandState_t *state) {
+  return murand_normal_double(state);
+}
+
+template <typename FloatType>
+MUTLASS_DEVICE
+FloatType random_uniform_float(murandState_t *state) {
+  return murand_uniform(state);
+}
+
+template <>
+MUTLASS_DEVICE
+double random_uniform_float<double>(murandState_t *state) {
+  return murand_uniform_double(state);
+}
+
+template <typename Element>
+struct RandomGaussianFunc {
+
+  using FloatType = typename std::conditional<(sizeof(Element) > 4), double, float>::type;
+  using IntType = typename std::conditional<(sizeof(Element) > 4), int64_t, int>::type;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    uint64_t seed;
+    FloatType mean;
+    FloatType stddev;
+    int int_scale;
+    FloatType float_scale_up;
+    FloatType float_scale_down;
+    int exclude_zero;           ///< If non-negative, excludes zeros
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      uint64_t seed_ = 0,
+      Element mean_ = 0,
+      Element stddev_ = 1,
+      int int_scale_ = -1,
+      int exclude_zero_ = -1
+    ):
+      seed(seed_),
+      mean(static_cast<FloatType>(mean_)),
+      stddev(static_cast<FloatType>(stddev_)),
+      int_scale(int_scale_),
+      exclude_zero(exclude_zero_) {
+
+      float_scale_up = FloatType(IntType(2) << int_scale); // scale up to clamp low order bits
+      float_scale_down = FloatType(1) / FloatType(IntType(2) << int_scale);
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  /// RNG state object
+  murandState_t rng_state;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  RandomGaussianFunc(Params const &params): params(params) {
+
+    uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x;
+
+    murand_init(params.seed, gtid, 0, &rng_state);
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  Element operator()() {
+
+    FloatType rnd = random_normal_float<FloatType>(&rng_state);
+    rnd = params.mean + params.stddev * rnd;
+
+    Element result;
+    if (params.int_scale >= 0) {
+      rnd = FloatType(IntType(rnd * params.float_scale_up));
+      result = Element(rnd * params.float_scale_down);
+    }
+    else {
+      result = Element(rnd);
+    }
+
+    if (params.exclude_zero >=0 && result == Element(0.0)) {
+      if (rnd > FloatType(0)) {
+        rnd += FloatType(1);
+      } else {
+        rnd -= FloatType(1);
+      }
+      result = Element(rnd);
+    }
+
+    return result;
+  }
+};
+
+
+template <typename Real>
+struct RandomGaussianFunc<complex<Real>> {
+
+  using Element = complex<Real>;
+  using FloatType = typename std::conditional<(sizeof(Real) > 4), double, float>::type;
+  using IntType = typename std::conditional<(sizeof(Real) > 4), int64_t, int>::type;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    uint64_t seed;
+    FloatType mean;
+    FloatType stddev;
+    int int_scale;
+    FloatType float_scale_up;
+    FloatType float_scale_down;
+    int exclude_zero;           ///< If non-negative, excludes zeros
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      uint64_t seed_ = 0,
+      Real mean_ = 0,
+      Real stddev_ = 1,
+      int int_scale_ = -1,
+      int exclude_zero_ = -1
+    ):
+      seed(seed_),
+      mean(static_cast<FloatType>(mean_)),
+      stddev(static_cast<FloatType>(stddev_)),
+      int_scale(int_scale_),
+      exclude_zero(exclude_zero_) {
+
+      float_scale_up = FloatType(IntType(1) << int_scale);
+      float_scale_up += FloatType(0.5) * float_scale_up;
+      float_scale_down = FloatType(1) / FloatType(IntType(1) << int_scale);
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  /// RNG state object
+  murandState_t rng_state;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  RandomGaussianFunc(Params const &params): params(params) {
+
+    uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x;
+
+    murand_init(params.seed, gtid, 0, &rng_state);
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  Element operator()() {
+
+    FloatType rnd_r = random_normal_float<FloatType>(&rng_state);
+    FloatType rnd_i = random_normal_float<FloatType>(&rng_state);
+    rnd_r = params.mean + params.stddev * rnd_r;
+    rnd_i = params.mean + params.stddev * rnd_i;
+
+    Element result;
+    if (params.int_scale >= 0) {
+      rnd_r = FloatType(IntType(rnd_r * params.float_scale_up));
+      rnd_i = FloatType(IntType(rnd_i * params.float_scale_down));
+
+      result = {
+        Real(rnd_r * params.float_scale_down),
+        Real(rnd_i * params.float_scale_down)
+      };
+    }
+    else {
+      result = Element(Real(rnd_r), Real(rnd_i));
+    }
+
+    if (params.exclude_zero >= 0 &&
+        result.real() == Real(0.0) &&
+        result.imag() == Real(0.0)) {
+
+      if (rnd_r > FloatType(0)) {
+        rnd_r += FloatType(1);
+      } else {
+        rnd_r -= FloatType(1);
+      }
+      result = Element(Real(rnd_r), Real(rnd_i));
+    }
+
+    return result;
+  }
+};
+
+/// Computes a random Gaussian distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillRandomGaussianFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  using RandomFunc = RandomGaussianFunc<Element>;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    typename RandomFunc::Params random;
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_ = TensorView(),
+      typename RandomFunc::Params random_ = typename RandomFunc::Params()
+    ):
+      view(view_), random(random_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  Params params;
+  RandomFunc random;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorFillRandomGaussianFunc(Params const &params): params(params), random(params.random) {
+
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    params.view.at(coord) = random();
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a Gaussian distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandomGaussian(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  typename RealType<Element>::Type mean = Element(0),   ///< Gaussian distribution's mean
+  typename RealType<Element>::Type stddev = Element(1), ///< Gaussian distribution's standard deviation
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+  int exclude_zero = -1,                  ///< If non-negative, excludes zeros from tensor init
+  musaStream_t stream = nullptr) {
+
+  using RandomFunc = detail::RandomGaussianFunc<Element>;
+  using Func = detail::TensorFillRandomGaussianFunc<Element, Layout>;
+  using Params = typename Func::Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, typename RandomFunc::Params(seed, mean, stddev, bits, exclude_zero)),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a Gaussian distribution.
+template <typename Element>               ///< Element type
+void BlockFillRandomGaussian(
+  Element *ptr,
+  size_t capacity,
+  uint64_t seed,                              ///< seed for RNG
+  typename RealType<Element>::Type mean,      ///< Gaussian distribution's mean
+  typename RealType<Element>::Type stddev,    ///< Gaussian distribution's standard deviation
+  int bits = -1,                              ///< If non-negative, specifies number of fractional bits that
+                                              ///  are not truncated to zero. Permits reducing precision of
+                                              ///  data.
+  musaStream_t stream = nullptr) {
+
+  using RandomFunc = detail::RandomGaussianFunc<Element>;
+
+  typename RandomFunc::Params params(seed, mean, stddev, bits);
+
+  BlockForEach<Element, RandomFunc>(ptr, capacity, params, /*grid_size*/0, /*block_size*/0, stream);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Computes a random uniform distribution
+template <typename Element>                ///< Element type
+struct RandomUniformFunc {
+
+  using FloatType = typename std::conditional<
+    (sizeof(Element) > 4),
+    double,
+    float>::type;
+
+  using IntType = typename std::conditional<
+    (sizeof(Element) > 4),
+    int64_t,
+    int>::type;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    uint64_t seed;
+    FloatType range;
+    FloatType max;
+    int int_scale;
+    double pnan;
+    FloatType float_scale_up;
+    FloatType float_scale_down;
+    int exclude_zero;           ///< If non-negative, excludes zeros
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      uint64_t seed_ = 0,
+      Element max_ = 1,
+      Element min = 0,
+      int int_scale_ = -1,
+      double pnan_ = 0,
+      int exclude_zero_ = -1
+    ):
+      seed(seed_),
+      range(static_cast<FloatType>(max_) - static_cast<FloatType>(min)),
+      max(static_cast<FloatType>(max_)),
+      int_scale(int_scale_),
+      pnan(pnan_),
+      exclude_zero(exclude_zero_) {
+
+      float_scale_up = FloatType(IntType(2) << int_scale); // scale up to clamp low order bits
+      float_scale_down = FloatType(1) / FloatType(IntType(2) << int_scale);
+
+      // Handle cases where min = 0 or max = 0 for excluding zeros
+      if (exclude_zero >= 0) {
+        range = (min == Element(0)) ? range - FloatType(1): range;
+        max = (max_ == Element(0)) ? max - FloatType(1): max;
+      }
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  /// RNG state object
+  murandState_t rng_state;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  RandomUniformFunc(Params const &params): params(params) {
+
+    uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x;
+
+    murand_init(params.seed, gtid, 0, &rng_state);
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  Element operator()() {
+
+    // Draw random float in [0.0, 1.0] to determine if element should be NaN.
+    if constexpr (std::numeric_limits<Element>::has_quiet_NaN) {
+      if (params.pnan > 0 && (murand_uniform(&rng_state) < (params.pnan))) {
+        return Element(NAN);
+      }
+    }
+
+    FloatType rnd = random_uniform_float<FloatType>(&rng_state);
+    rnd = params.max - params.range * rnd;
+
+    // Random values are cast to integer after scaling by a power of two to facilitate error
+    // testing
+    Element result;
+
+    if (params.int_scale >= 0) {
+      rnd = FloatType(IntType(rnd * params.float_scale_up));
+      result = Element(rnd * params.float_scale_down);
+    }
+    else {
+      result = Element(rnd);
+    }
+
+    if (params.exclude_zero >=0 && result == Element(0.0)) {
+      if (rnd > FloatType(0)) {
+        rnd = std::min(params.max, rnd + FloatType(1));
+      } else {
+        rnd = std::max((params.max - params.range), rnd - FloatType(1));
+      }
+      result = Element(rnd);
+    }
+
+    return result;
+  }
+};
+
+/// Computes a random Gaussian distribution
+template <typename Real>
+struct RandomUniformFunc<complex<Real>> {
+
+  using Element = complex<Real>;
+
+  using FloatType = typename std::conditional<
+    (sizeof(Real) > 4),
+    double,
+    float>::type;
+
+  using IntType = typename std::conditional<
+    (sizeof(Real) > 4),
+    int64_t,
+    int>::type;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    uint64_t seed;
+    FloatType range;
+    FloatType min;
+    int int_scale;
+    double pnan;
+    FloatType float_scale_up;
+    FloatType float_scale_down;
+    int exclude_zero;           ///< If non-negative, excludes zeros
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      uint64_t seed_ = 0,
+      FloatType max = 1,
+      FloatType min_ = 0,
+      int int_scale_ = -1,
+      double pnan_ = 0,
+      int exclude_zero_ = -1
+    ):
+      seed(seed_),
+      range(static_cast<FloatType>(max - min_)),
+      min(static_cast<FloatType>(min_)),
+      int_scale(int_scale_),
+      pnan(pnan_),
+      exclude_zero(exclude_zero_) {
+
+      float_scale_up = FloatType(IntType(1) << int_scale);
+      float_scale_up += FloatType(0.5) * float_scale_up;
+      float_scale_down = FloatType(1) / FloatType(IntType(1) << int_scale);
+
+      // Handle cases where min = 0 or max = 0 for excluding zeros
+      if (exclude_zero >= 0) {
+        min = (min == FloatType(0)) ? min + FloatType(1): min;
+        range = (max == FloatType(0)) ? range - FloatType(1): range;
+      }
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  /// RNG state object
+  murandState_t rng_state;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  RandomUniformFunc(Params const &params): params(params) {
+
+    uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x;
+
+    murand_init(params.seed, gtid, 0, &rng_state);
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  Element operator()() {
+
+    // Draw random float in [0.0, 1.0] to determine if element should be NaN.
+    if constexpr (std::numeric_limits<Element>::has_quiet_NaN) {
+      if (params.pnan > 0 && (murand_uniform(&rng_state) < (params.pnan))) {
+        return Element(Real(NAN), Real(NAN));
+      }
+    }
+
+    FloatType rnd_r = random_uniform_float<FloatType>(&rng_state);
+    FloatType rnd_i = random_uniform_float<FloatType>(&rng_state);
+
+    rnd_r = params.min + params.range * rnd_r;
+    rnd_i = params.min + params.range * rnd_i;
+
+    // Random values are cast to integer after scaling by a power of two to facilitate error
+    // testing
+    Element result;
+
+    if (params.int_scale >= 0) {
+      rnd_r = FloatType(IntType(rnd_r * params.float_scale_up));
+      rnd_i = FloatType(IntType(rnd_i * params.float_scale_up));
+
+      result = {
+        Real(rnd_r * params.float_scale_down),
+        Real(rnd_i * params.float_scale_down)
+      };
+    }
+    else {
+      result = Element(Real(rnd_r), Real(rnd_i));
+    }
+
+    if (params.exclude_zero >= 0 &&
+        result.real() == Real(0.0) &&
+        result.imag() == Real(0.0)) {
+
+      if (rnd_r > FloatType(0)) {
+        rnd_r = std::min(params.min + params.range, rnd_r + FloatType(1));
+      } else {
+        rnd_r = std::max((params.min), rnd_r - FloatType(1));
+      }
+      result = Element(Real(rnd_r), Real(rnd_i));
+    }
+
+    return result;
+  }
+};
+
+/// Computes a random uniform distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillRandomUniformFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  using RandomFunc = RandomUniformFunc<Element>;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    typename RandomFunc::Params random;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_ = TensorView(),
+      typename RandomFunc::Params random_ = RandomFunc::Params()
+    ):
+      view(view_), random(random_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  Params params;
+  RandomFunc random;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorFillRandomUniformFunc(Params const &params): params(params), random(params.random) {
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    params.view.at(coord) = random();
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandomUniform(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  typename RealType<Element>::Type max = Element(1), ///< upper bound of distribution
+  typename RealType<Element>::Type min = Element(0), ///< lower bound for distribution
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+  double pnan = 0,                        ///< Percentage of NaN elements.
+  int exclude_zero = -1,               ///< If non-negative, excludes zeros from tensor init
+  musaStream_t stream = nullptr) {
+
+  using RandomFunc = detail::RandomUniformFunc<Element>;
+  using Func = detail::TensorFillRandomUniformFunc<Element, Layout>;
+  using Params = typename Func::Params;
+
+  typename RandomFunc::Params random(seed, max, min, bits, pnan, exclude_zero);
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, random),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <typename Element>
+void BlockFillRandomUniform(
+  Element *ptr,
+  size_t capacity,
+  uint64_t seed,                          ///< seed for RNG
+  typename RealType<Element>::Type max,   ///< upper bound of distribution
+  typename RealType<Element>::Type min,   ///< lower bound for distribution
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+  double pnan = 0,                        ///< Percentage of NaN elements.
+  musaStream_t stream = nullptr) {
+
+  using RandomFunc = detail::RandomUniformFunc<Element>;
+
+  typename RandomFunc::Params params(seed, max, min, bits, pnan);
+
+  BlockForEach<Element, RandomFunc>(ptr, capacity, params, /*grid_size*/0, /*block_size*/0, stream);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Computes a random sparse meta
+template <typename Element>               ///< Element type
+struct RandomSparseMetaFunc {
+
+  using FloatType = float;
+
+  using IntType = int32_t;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    uint64_t seed;
+    FloatType range;
+    int MetaSizeInBits;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      uint64_t seed_ = 0,
+      int MetaSizeInBits_ = 2
+    ):
+      seed(seed_),
+      MetaSizeInBits(MetaSizeInBits_) {
+      if (MetaSizeInBits_ == 2) {
+        range = 6;
+      }
+      else if (MetaSizeInBits_ == 4) {
+        range = 2;
+      }
+      else {
+        throw std::invalid_argument("Invalid MetaSizeInBits");
+      }
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  /// RNG state object
+  murandState_t rng_state;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  RandomSparseMetaFunc(Params const &params): params(params) {
+
+    uint64_t gtid = threadIdx.x + blockIdx.x * blockDim.x;
+
+    murand_init(params.seed, gtid, 0, &rng_state);
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  Element operator()() {
+    Element FourToTwoMeta[6] = {0x4, 0x8, 0x9, 0xc, 0xd, 0xe};
+    Element TwoToOneMeta[2] = {0x4, 0xe};
+
+    Element *MetaArray =
+        (params.MetaSizeInBits == 2) ? FourToTwoMeta : TwoToOneMeta;
+
+    Element result = 0x0;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < mutlass::sizeof_bits<Element>::value / 4; ++i) {
+      FloatType rnd = random_uniform_float<FloatType>(&rng_state);
+      rnd = params.range * rnd;
+      Element meta = MetaArray[(int)rnd];
+
+      result = (Element)(result | ((Element)(meta << (i * 4))));
+    }
+
+    return result;
+  }
+};
+
+/// Computes a random Gaussian distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillRandomSparseMetaFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  using RandomFunc = RandomSparseMetaFunc<Element>;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    typename RandomFunc::Params random;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_ = TensorView(),
+      typename RandomFunc::Params random_ = RandomFunc::Params()
+    ):
+      view(view_), random(random_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  Params params;
+  RandomFunc random;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorFillRandomSparseMetaFunc(Params const &params): params(params), random(params.random) {
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    params.view.at(coord) = random();
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandomSparseMeta(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  int MetaSizeInBits = 2,                 ///< meta data size
+  musaStream_t stream = nullptr) {
+
+  using RandomFunc = detail::RandomSparseMetaFunc<Element>;
+  using Func = detail::TensorFillRandomUniformFunc<Element, Layout>;
+  using Params = typename Func::Params;
+
+  typename RandomFunc::Params random(seed, MetaSizeInBits);
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, random),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <typename Element>
+void BlockFillRandomSparseMeta(
+  Element *ptr,
+  size_t capacity,
+  uint64_t seed,                          ///< seed for RNG
+  int MetaSizeInBits = 2,                 ///< meta data size
+  musaStream_t stream = nullptr) {
+
+  using RandomFunc = detail::RandomSparseMetaFunc<Element>;
+
+  typename RandomFunc::Params params(seed, MetaSizeInBits);
+
+  BlockForEach<Element, RandomFunc>(ptr, capacity, params, /*grid_size*/0, /*block_size*/0, stream);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Functor to fill a tensor with zeros off the diagonal and a uniform value on the diagonal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillDiagonalFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    Element diag;
+    Element other;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    Params(
+      TensorView view_ = TensorView(),
+      Element diag_ = Element(1),
+      Element other_ = Element(0)
+    ):
+      view(view_), diag(diag_), other(other_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorFillDiagonalFunc(Params const &params): params(params) {
+
+  }
+
+  /// Updates the tensor
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    bool is_diag = true;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < Layout::kRank; ++i) {
+      if (coord[i] != coord[i - 1]) {
+        is_diag = false;
+        break;
+      }
+    }
+
+    params.view.at(coord) = (is_diag ? params.diag : params.other);
+  }
+};
+
+// Overwrites the elements of a tensor with a uniform value depending on fill mode
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillPartialFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    Element element;
+    FillMode fill_mode;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params(): fill_mode(FillMode::kNone) { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_,
+      Element element_,
+      FillMode fill_mode_
+    ):
+      view(view_), element(element_), fill_mode(fill_mode_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  //
+  // Methods
+  //
+
+  MUTLASS_DEVICE
+  TensorFillPartialFunc(Params const &params): params(params) {
+
+  }
+
+  /// Overwrites the element if it is within the covered region.
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    bool predicate = true;
+
+    switch (params.fill_mode) {
+    case FillMode::kFull:
+      predicate = true;
+      break;
+
+    case FillMode::kLower:
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 1; i < Layout::kRank; ++i) {
+        if (coord[i - 1] < coord[i]) {
+          predicate = false;
+          break;
+        }
+      }
+      break;
+
+    case FillMode::kUpper:
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 1; i < Layout::kRank; ++i) {
+        if (coord[i - 1] > coord[i]) {
+          predicate = false;
+          break;
+        }
+      }
+      break;
+
+    case FillMode::kDiagonal:
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 1; i < Layout::kRank; ++i) {
+        if (coord[i - 1] != coord[i]) {
+          predicate = false;
+          break;
+        }
+      }
+      break;
+
+    case FillMode::kNone: // fall-through
+
+    default:
+      predicate = false;
+      break;
+    }
+
+    if (predicate) {
+      params.view.at(coord) = params.element;
+    }
+  }
+};
+
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorClearPartialFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  ///
+  static_assert((Layout::kRank == 2), "TensorClearPartial is only supported for matrices");
+
+  /// Parameters structure
+  struct Params {
+    TensorView view{};
+    Element element{};
+    FillMode fill_mode{FillMode::kNone};
+    int alignment{0};
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  //
+  // Methods
+  //
+
+  MUTLASS_DEVICE
+  TensorClearPartialFunc(Params const &params): params(params) {
+
+  }
+
+  /// Overwrites the element if it is within the covered region.
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    bool predicate = true;
+
+    switch (params.fill_mode) {
+
+    case FillMode::kLower:
+      if ((coord[0] >= coord[1]) ||
+          ((coord[1] - coord[0]) >= params.alignment))  {
+          predicate = false;
+        break;
+      }
+      break;
+
+    case FillMode::kUpper:
+      if ((coord[0] <= coord[1]) ||
+          ((coord[0] - coord[1]) >= params.alignment))  {
+          predicate = false;
+        break;
+      }
+      break;
+
+    case FillMode::kNone: // fall-through
+
+    default:
+      predicate = false;
+      break;
+    }
+
+    if (predicate) {
+      params.view.at(coord) = params.element;
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor everywhere with a unique value for its diagonal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillDiagonal(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  Element diag = Element(1),              ///< value to write in the diagonal
+  Element other = Element(0),             ///< value to write off the diagonal
+  musaStream_t stream = nullptr) {
+
+  typedef detail::TensorFillDiagonalFunc<Element, Layout> Func;
+  typedef typename Func::Params Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, diag, other),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+/// Fills a tensor partially depending on fill mode. Elements not covered by the fillmode are
+/// not written.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillPartial(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  Element element,
+  FillMode fill_mode,
+  musaStream_t stream = nullptr) {
+
+  typedef detail::TensorFillPartialFunc<Element, Layout> Func;
+  typedef typename Func::Params Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, element, fill_mode),
+    stream
+  );
+}
+
+/// Clears a tensor partially depending on fill mode and alignment. Elements on the wrong-side
+/// of fillmode (upto the alignment) are overwritten with the user supplied element (typically zeros)
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorClearPartial(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  Element element,
+  FillMode fill_mode,
+  int alignment,
+  musaStream_t stream = nullptr) {
+
+  typedef detail::TensorClearPartialFunc<Element, Layout> Func;
+  typedef typename Func::Params Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params{view, element, fill_mode, alignment},
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with a uniform value
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFill(
+  TensorView<Element, Layout> view,         ///< destination tensor
+  Element val = Element(0),                 ///< value to uniformly fill it with
+  musaStream_t stream = nullptr) {
+
+  TensorFillDiagonal(view, val, val, stream);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor's diagonal with 1 and 0 everywhere else.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillIdentity(
+  TensorView<Element, Layout> view,                 ///< destination tensor
+  musaStream_t stream = nullptr) {
+
+  TensorFillDiagonal(view, Element(1), Element(0), stream);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Computes a random Gaussian distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorUpdateDiagonalFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    Element diag;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_ = TensorView(),
+      Element diag_ = Element(1)
+    ):
+      view(view_), diag(diag_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorUpdateDiagonalFunc(Params const &params): params(params) {
+
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    bool is_diag = true;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < Layout::kRank; ++i) {
+      if (coord[i] != coord[i - 1]) {
+        is_diag = false;
+        break;
+      }
+    }
+
+    if (is_diag) {
+      params.view.at(coord) = params.diag;
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Writes a uniform value to the diagonal of a tensor without modifying off-diagonal elements.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorUpdateDiagonal(
+  TensorView<Element, Layout> view,                 ///< destination tensor
+  Element diag = Element(1),
+  musaStream_t stream = nullptr) {
+
+  typedef detail::TensorUpdateDiagonalFunc<Element, Layout> Func;
+  typedef typename Func::Params Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, diag),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Computes a random Gaussian distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorUpdateOffDiagonalFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    Element other;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_ = TensorView(),
+      Element other_ = Element(0)
+    ):
+      view(view_), other(other_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorUpdateOffDiagonalFunc(Params const &params): params(params) {
+
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    bool is_diag = true;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < Layout::kRank; ++i) {
+      if (coord[i] != coord[i - 1]) {
+        is_diag = false;
+        break;
+      }
+    }
+
+    if (!is_diag) {
+      params.view.at(coord) = params.other;
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Writes a uniform value to all elements in the tensor without modifying diagonal elements.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorUpdateOffDiagonal(
+  TensorView<Element, Layout> view,      ///< destination tensor
+  Element other = Element(1),
+  musaStream_t stream = nullptr) {
+
+  typedef detail::TensorUpdateOffDiagonalFunc<Element, Layout> Func;
+  typedef typename Func::Params Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, other),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Computes a random Gaussian distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillLinearFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    Array<Element, Layout::kRank> v;
+    Element s;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_,      ///< destination tensor
+      Array<Element, Layout::kRank> const & v_,
+      Element s_ = Element(0)
+    ):
+      view(view_), v(v_), s(s_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorFillLinearFunc(Params const &params): params(params) {
+
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    Element sum = params.s;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Layout::kRank; ++i) {
+      if constexpr (is_complex<Element>::value) {
+        if constexpr (sizeof_bits<Element>::value <= 32) {
+          sum = Element(static_cast<complex<float>>(sum) +
+                  static_cast<complex<float>>(params.v[i]) * static_cast<complex<float>>(coord[i]));
+        }
+      }
+      else if constexpr (sizeof_bits<Element>::value <= 32) {
+        if constexpr (std::numeric_limits<Element>::is_integer) {
+          sum = Element(static_cast<int32_t>(sum) +
+                  static_cast<int32_t>(params.v[i]) * static_cast<int32_t>(coord[i]));
+        }
+        else {
+          sum = Element(static_cast<float>(sum) +
+                  static_cast<float>(params.v[i]) * static_cast<float>(coord[i]));
+        }
+      }
+      else {
+        sum += params.v[i] * coord[i];
+      }
+    }
+
+    params.view.at(coord) = sum;
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills tensor with a linear combination of its coordinate and another vector
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillLinear(
+  TensorView<Element, Layout> view,      ///< destination tensor
+  Array<Element, Layout::kRank> const & v,
+  Element s = Element(0),
+  musaStream_t stream = nullptr) {
+
+  using Func = detail::TensorFillLinearFunc<Element, Layout>;
+  using Params = typename Func::Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, v, s),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values from a distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandom(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  uint64_t seed,
+  Distribution dist,
+  musaStream_t stream = nullptr,
+  int exclude_zero = -1                   ///< If non-negative, excludes 0.
+                                          ///  Note that setting this flag will result in more 1's,
+                                          ///  as we use a simple mechanism to replace 0's by adding/subtracting 1's.
+  ) {
+
+  using Real = typename RealType<Element>::Type;
+
+  if (dist.kind == Distribution::Gaussian) {
+    TensorFillRandomGaussian<Element, Layout>(
+      view,
+      seed,
+      static_cast<Real>(dist.gaussian.mean),
+      static_cast<Real>(dist.gaussian.stddev),
+      dist.int_scale,
+      exclude_zero,
+      stream);
+  } else if (dist.kind == Distribution::Uniform) {
+    TensorFillRandomUniform<Element, Layout>(
+      view,
+      seed,
+      static_cast<Real>(dist.uniform.max),
+      static_cast<Real>(dist.uniform.min),
+      dist.int_scale,
+      dist.uniform.pnan,
+      exclude_zero,
+      stream);
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a block of data with sequential elements
+template <
+  typename Element
+>
+void BlockFillSequential(
+  Element *ptr,
+  int64_t capacity,
+  Element v = Element(1),
+  Element s = Element(0)) {
+
+  using Layout = layout::PackedVectorLayout;
+  Layout::TensorCoord size(static_cast<Layout::Index>(capacity)); // -Wconversion
+  Layout layout = Layout::packed(size);
+  TensorView<Element, Layout> view(ptr, layout, size);
+
+  Array<Element, Layout::kRank> c{};
+  c[0] = v;
+
+  TensorFillLinear(view, c, s);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a block of data with sequential elements
+template <
+  typename Element
+>
+void BlockFillRandom(
+  Element *ptr,
+  size_t capacity,
+  uint64_t seed,
+  Distribution dist,
+  musaStream_t stream = nullptr) {
+
+  using Real = typename RealType<Element>::Type;
+
+  if (dist.kind == Distribution::Gaussian) {
+    BlockFillRandomGaussian<Element>(
+      ptr,
+      capacity,
+      seed,
+      static_cast<Real>(dist.gaussian.mean),
+      static_cast<Real>(dist.gaussian.stddev),
+      dist.int_scale,
+      stream);
+  }
+  else if (dist.kind == Distribution::Uniform) {
+    BlockFillRandomUniform<Element>(
+      ptr,
+      capacity,
+      seed,
+      static_cast<Real>(dist.uniform.max),
+      static_cast<Real>(dist.uniform.min),
+      dist.int_scale,
+      dist.uniform.pnan,
+      stream);
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Computes a random Gaussian distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorCopyDiagonalInFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    Element const *ptr;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_,      ///< destination tensor
+      Element const *ptr_
+    ):
+      view(view_), ptr(ptr_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorCopyDiagonalInFunc(Params const &params): params(params) {
+
+  }
+
+  /// Only update the diagonal element
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+    bool is_diagonal = true;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < Layout::kRank; ++i) {
+      if (coord[i] != coord[0]) {
+        is_diagonal = false;
+      }
+    }
+    if (is_diagonal) {
+      params.view.at(coord) = params.ptr[coord[0]];
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies a diagonal in from host memory without modifying off-diagonal elements.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorCopyDiagonalIn(
+  TensorView<Element, Layout> view,   ///< destination tensor
+  Element const *ptr,                        ///< dense buffer of elements
+  musaStream_t stream = nullptr) {
+
+  using Func = detail::TensorCopyDiagonalInFunc<Element, Layout>;
+  using Params = typename Func::Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, ptr),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+namespace detail {
+
+/// Computes a random Gaussian distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorCopyDiagonalOutFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Scalar type
+  typedef typename TensorView::Element T;
+
+  /// Coordinate in tensor's index space
+  typedef typename TensorView::TensorCoord TensorCoord;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    Element *ptr;
+
+    /// Default ctor
+    MUTLASS_HOST_DEVICE
+    Params() { }
+
+    //
+    // Methods
+    //
+
+    /// Construction of Gaussian RNG functor.
+    Params(
+      TensorView view_,      ///< destination tensor
+      Element *ptr_
+    ):
+      view(view_), ptr(ptr_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  /// Parameters object
+  Params params;
+
+  //
+  // Methods
+  //
+
+  /// Device-side initialization of RNG
+  MUTLASS_DEVICE
+  TensorCopyDiagonalOutFunc(Params const &params): params(params) {
+
+  }
+
+  /// Compute random value and update RNG state
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+    bool is_diagonal = true;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < Layout::kRank; ++i) {
+      if (coord[i] != coord[0]) {
+        is_diagonal = false;
+      }
+    }
+    if (is_diagonal) {
+      params.ptr[coord[0]] = params.view.at(coord);
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies the diagonal of a tensor into a dense buffer in host memory.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorCopyDiagonalOut(
+  Element *ptr,                               ///< dense buffer of elements
+  TensorView<Element, Layout> view,      ///< source tensor
+  musaStream_t stream = nullptr) {
+
+  using Func = detail::TensorCopyDiagonalOutFunc<Element, Layout>;
+  using Params = typename Func::Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, ptr),
+    /*grid_size*/0, /*block_size*/0,
+    stream
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_foreach.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_foreach.h
new file mode 100644
index 000000000..da2d1f603
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_foreach.h
@@ -0,0 +1,145 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <stdexcept>
+#include "mutlass/mutlass.h"
+#include "mutlass/util/reference/device/kernel/tensor_foreach.h"
+
+namespace mutlass  {
+namespace reference {
+namespace device {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Launches a kernel calling a functor for each element in a tensor's index space.
+template <typename Func, int Rank, typename Params>
+struct TensorForEach {
+
+  /// Constructor performs the operation.
+  TensorForEach(
+    Coord<Rank> size, Params params = Params(),
+    int grid_size = 0, int block_size = 0,
+    musaStream_t stream = nullptr) {
+
+    if (!grid_size || !block_size) {
+
+      // if grid_size or block_size are zero, query occupancy using the MUSA Occupancy API
+      musaError_t result = musaOccupancyMaxPotentialBlockSize(
+        &grid_size,
+        &block_size,
+        reinterpret_cast<void const *>(kernel::TensorForEach<Func, Rank, Params>));
+
+      if (result != musaSuccess) {
+        throw std::runtime_error("Failed to query occupancy.");
+      }
+
+      // Limit block size. This has the effect of increasing the number of items processed by a
+      // single thread and reduces the impact of initialization overhead.
+      block_size = (block_size < 128 ? block_size : 128);
+    }
+
+    dim3 grid(grid_size, 1, 1);
+    dim3 block(block_size, 1, 1);
+
+    kernel::TensorForEach<Func, Rank, Params><<< grid, block, 0, stream >>>(size, params);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Launches a kernel calling a functor for each element along a tensor's diagonal
+template <typename Func, int Rank, typename Params>
+struct TensorDiagonalForEach {
+
+  /// Constructor performs the operation
+  TensorDiagonalForEach(
+    Coord<Rank> size, Params params = Params(),
+    int start = 0, int end = -1,
+    int block_size = 128, musaStream_t stream = nullptr) {
+
+    if (end < 0) {
+      end = size.min();
+    }
+
+    dim3 block(block_size, 1, 1);
+    dim3 grid((end - start + block_size - 1) / block_size, 1, 1);
+
+    kernel::TensorDiagonalForEach<Func, Rank, Params><<< grid, block, 0, stream >>>(
+      size, params, start, end);
+  }
+};
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Element, typename Func>
+struct BlockForEach {
+
+  /// Constructor performs the operation.
+  BlockForEach(
+    Element *ptr,
+    size_t capacity,
+    typename Func::Params params = typename Func::Params(),
+    int grid_size = 0,
+    int block_size = 0,
+    musaStream_t stream = nullptr) {
+
+    if (!grid_size || !block_size) {
+
+      // if grid_size or block_size are zero, query occupancy using the MUSA Occupancy API
+      musaError_t result = musaOccupancyMaxPotentialBlockSize(
+        &grid_size,
+        &block_size,
+        reinterpret_cast<void const *>(kernel::BlockForEach<Element, Func>));
+
+      if (result != musaSuccess) {
+        throw std::runtime_error("Failed to query occupancy.");
+      }
+
+      // Limit block size. This has the effect of increasing the number of items processed by a
+      // single thread and reduces the impact of initialization overhead.
+      block_size = (block_size < 128 ? block_size : 128);
+    }
+
+    dim3 grid(grid_size, 1, 1);
+    dim3 block(block_size, 1, 1);
+
+    kernel::BlockForEach<Element, Func><<< grid, block, 0, stream >>>(ptr, capacity, params);
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_reduce.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_reduce.h
new file mode 100644
index 000000000..3cc45e604
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_reduce.h
@@ -0,0 +1,511 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <cmath>
+
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/util/device_memory.h"
+#include "mutlass/util/reference/detail/linear_to_coordinate.h"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace kernel {
+
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp,
+  int kBlockSize = 128
+>
+__global__ void TensorTransformReducePartial(
+  TensorView<Element, Layout> view,     /// View of the tensor to reduce over
+  ComputeType identity,                 /// Identity element of the reduction operation
+  ReduceOp reduce,                      /// Reduces an accumulated value with a transformed element: f(ComputeType, ComputeType) => ComputeType
+  TransformOp transform,                /// Transforms the tensor element to ComputeType: g(Element) => ComputeType
+  ComputeType *workspace) {             /// Device-side workspace for accumulating partial results. The reduced element is stored in workspace[0]
+
+  int64_t idx = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t size = view.size();
+
+  __shared__ ComputeType scratchpad[kBlockSize];
+
+  for (; idx < size; idx += blockDim.x * gridDim.x) {
+
+    // Map linear thread ID onto tensor coordinate
+    typename Layout::TensorCoord coord;
+
+    mutlass::reference::detail::LinearToCoordinate<Layout::kRank>()(coord, idx, view.extent());
+
+    if (view.contains(coord)) {
+
+      // Fetch element
+      Element x = view.at(coord);
+
+      // Transform
+      identity = reduce(identity, transform(x));
+    }
+  }
+
+  scratchpad[threadIdx.x] = identity;
+
+  __syncthreads();
+
+  // One thread performs the final reduction and stores out. This could be enhanced via
+  // a tree reduction and pipelining.
+  if (threadIdx.x == 0) {
+
+    for (int i = 1; i < kBlockSize; ++i) {
+      identity = reduce(identity, scratchpad[i]);
+    }
+
+    workspace[blockIdx.x] = identity;
+  }
+}
+
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp,
+  int kBlockSize = 128
+>
+__global__ void TensorTransformReducePartial(
+  TensorView<Element, Layout> view_A,   /// View of the tensor to reduce over
+  TensorView<Element, Layout> view_B,   /// View of the tensor to reduce over
+  ComputeType identity,                 /// Identity element of the reduction operation
+  ReduceOp reduce,                      /// Reduces an accumulated value with a transformed element: f(ComputeType, ComputeType) => ComputeType
+  TransformOp transform,                /// Transforms the tensor element to ComputeType: g(Element) => ComputeType
+  ComputeType *workspace) {             /// Device-side workspace for accumulating partial results. The reduced element is stored in workspace[0]
+
+  int64_t idx = threadIdx.x + blockIdx.x * blockDim.x;
+  int64_t size = view_A.size();
+
+  __shared__ ComputeType scratchpad[kBlockSize];
+
+  for (; idx < size; idx += blockDim.x * gridDim.x) {
+
+    // Map linear thread ID onto tensor coordinate
+    typename Layout::TensorCoord coord;
+
+    mutlass::reference::detail::LinearToCoordinate<Layout::kRank>()(coord, idx, view_A.extent());
+
+    if (view_A.contains(coord)) {
+
+      // Fetch element
+      Element a = view_A.at(coord);
+      Element b = view_B.at(coord);
+
+      // Transform
+      identity = reduce(identity, transform(a, b));
+    }
+  }
+
+  scratchpad[threadIdx.x] = identity;
+
+  __syncthreads();
+
+  // One thread performs the final reduction and stores out. This could be enhanced via
+  // a tree reduction and pipelining.
+  if (threadIdx.x == 0) {
+
+    for (int i = 1; i < kBlockSize; ++i) {
+      identity = reduce(identity, scratchpad[i]);
+    }
+
+    workspace[blockIdx.x] = identity;
+  }
+}
+
+
+template <
+  typename ComputeType,
+  typename ReduceOp,
+  int kBlockSize = 32
+>
+__global__ void TensorTransformReduceFinalize(
+  ComputeType *workspace,
+  ComputeType identity,
+  int workspace_size,
+  ReduceOp reduce) {
+
+  __shared__ ComputeType scratchpad[kBlockSize];
+
+  for (int idx = threadIdx.x; idx < workspace_size; idx += kBlockSize) {
+    identity = reduce(identity, workspace[idx]);
+  }
+
+  scratchpad[threadIdx.x] = identity;
+
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+
+    for (int i = 1; i < kBlockSize; ++i) {
+      identity = reduce(identity, scratchpad[i]);
+    }
+
+    workspace[0] = identity;
+  }
+}
+
+} // namespace kernel
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Transform-reduce operation over the elements of a tensor
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp
+>
+ComputeType TensorTransformReduce(
+  TensorView<Element, Layout> view,     /// View of the tensor to reduce over
+  ComputeType identity,                 /// Identity element of the reduction operation
+  ReduceOp reduce,                      /// Reduces an accumulated value with a transformed element: f(ComputeType, ComputeType) => ComputeType
+  TransformOp transform,                /// Transforms the tensor element to ComputeType: g(Element) => ComputeType
+  ComputeType *workspace,               /// Device-side workspace for accumulating partial results. The reduced element is stored in workspace[0]
+  int workspace_size,                   /// Number of elements in workspace
+  musaStream_t stream = nullptr,        /// MUSA stream to launch into
+  bool copy_out = true                  /// If true, the value of workspace[0] is copied to host and returned. Otherwise, `identity` is returned.
+) {
+
+  int const kBlockSize = 128;
+
+  dim3 block(kBlockSize, 1);
+  dim3 grid(workspace_size, 1);
+
+  kernel::TensorTransformReducePartial<
+    Element, Layout, ComputeType, ReduceOp, TransformOp, kBlockSize
+  ><<< grid, block, 0, stream >>>(
+    view, identity, reduce, transform, workspace
+  );
+
+  int const kFinalizeBlockSize = 32;
+
+  kernel::TensorTransformReduceFinalize<
+    ComputeType, ReduceOp, kFinalizeBlockSize
+  ><<< dim3(1, 1), dim3(kFinalizeBlockSize, 1), 0, stream >>>(
+    workspace, identity, workspace_size, reduce
+  );
+
+  if (copy_out) {
+    musaError_t result = musaMemcpy(&identity, workspace, sizeof(identity), musaMemcpyDeviceToHost);
+    if (result != musaSuccess) {
+      throw std::runtime_error("musaMemcpy() failed");
+    }
+  }
+
+  return identity;
+}
+
+/// Transform-reduce operation over the elements of two tensors, zipped together
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp
+>
+ComputeType TensorTransformReduce(
+  TensorView<Element, Layout> view_A,   /// View of the tensor to reduce over
+  TensorView<Element, Layout> view_B,   /// View of the tensor to reduce over
+  ComputeType identity,                 /// Identity element of the reduction operation
+  ReduceOp reduce,                      /// Reduces an accumulated value with a transformed element: f(ComputeType, ComputeType) => ComputeType
+  TransformOp transform,                /// Transforms the tensor element to ComputeType: g(Element) => ComputeType
+  ComputeType *workspace,               /// Device-side workspace for accumulating partial results. The reduced element is stored in workspace[0]
+  int workspace_size,                   /// Number of elements in workspace
+  musaStream_t stream = nullptr,        /// MUSA stream to launch into
+  bool copy_out = true                  /// If true, the value of workspace[0] is copied to host and returned. Otherwise, `identity` is returned.
+) {
+
+  if (view_A.extent() != view_B.extent()) {
+    throw std::runtime_error("Extents must be equal.");
+  }
+
+  int const kBlockSize = 128;
+
+  dim3 block(kBlockSize, 1);
+  dim3 grid(workspace_size, 1);
+
+  kernel::TensorTransformReducePartial<
+    Element, Layout, ComputeType, ReduceOp, TransformOp, kBlockSize
+  ><<< grid, block, 0, stream >>>(
+    view_A, view_B, identity, reduce, transform, workspace
+  );
+
+  int const kFinalizeBlockSize = 32;
+
+  kernel::TensorTransformReduceFinalize<
+    ComputeType, ReduceOp, kFinalizeBlockSize
+  ><<< dim3(1, 1), dim3(kFinalizeBlockSize, 1), 0, stream >>>(
+    workspace, identity, workspace_size, reduce
+  );
+
+  if (copy_out) {
+    musaError_t result = musaMemcpy(&identity, workspace, sizeof(identity), musaMemcpyDeviceToHost);
+    if (result != musaSuccess) {
+      throw std::runtime_error("musaMemcpy() failed");
+    }
+  }
+
+  return identity;
+}
+
+/// Transform-reduce operation over the elements of a tensor. This helper allocates the device-side
+/// workspace
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp
+>
+ComputeType TensorTransformReduce(
+  TensorView<Element, Layout> view,
+  ComputeType identity,
+  ReduceOp reduce,
+  TransformOp transform,
+  musaStream_t stream = nullptr,
+  int workspace_size = 0
+) {
+
+  // Optionally query for the SM count to size the workspace.
+  if (!workspace_size) {
+
+    int device_idx = 0;
+    musaDeviceProp prop;
+
+    musaError_t result = musaGetDevice(&device_idx);
+    if (result != musaSuccess) {
+      throw std::runtime_error("musaGetDevice() failed");
+    }
+
+    result = musaGetDeviceProperties(&prop, device_idx);
+    if (result != musaSuccess) {
+      throw std::runtime_error("musaGetDeviceProp() failed");
+    }
+
+    workspace_size = int(prop.multiProcessorCount);
+  }
+
+  DeviceAllocation<ComputeType> workspace(workspace_size);
+
+  ComputeType output = TensorTransformReduce(
+    view,
+    identity,
+    reduce,
+    transform,
+    workspace.get(),
+    workspace_size,
+    stream,
+    true);
+
+  return output;
+}
+
+
+/// Transform-reduce operation over the elements of a tensor. This helper allocates the device-side
+/// workspace
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp
+>
+ComputeType TensorTransformReduce(
+  TensorView<Element, Layout> view_A,
+  TensorView<Element, Layout> view_B,
+  ComputeType identity,
+  ReduceOp reduce,
+  TransformOp transform,
+  musaStream_t stream = nullptr,
+  int workspace_size = 0
+) {
+
+  // Optionally query for the SM count to size the workspace.
+  if (!workspace_size) {
+
+    int device_idx = 0;
+    musaDeviceProp prop;
+
+    musaError_t result = musaGetDevice(&device_idx);
+    if (result != musaSuccess) {
+      throw std::runtime_error("musaGetDevice() failed");
+    }
+
+    result = musaGetDeviceProperties(&prop, device_idx);
+    if (result != musaSuccess) {
+      throw std::runtime_error("musaGetDeviceProp() failed");
+    }
+
+    workspace_size = int(prop.multiProcessorCount);
+  }
+
+  DeviceAllocation<ComputeType> workspace(workspace_size);
+
+  ComputeType output = TensorTransformReduce(
+    view_A,
+    view_B,
+    identity,
+    reduce,
+    transform,
+    workspace.get(),
+    workspace_size,
+    stream,
+    true);
+
+  return output;
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper to compute the sum of the elements of a tensor
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = Element
+>
+ComputeType TensorSum(
+  TensorView<Element, Layout> view,
+  ComputeType identity = ComputeType(),
+  musaStream_t stream = nullptr,
+  int workspace_size = 0
+) {
+
+  plus<ComputeType> reduce;
+  NumericConverter<ComputeType, Element> transform;
+
+  return TensorTransformReduce(
+    view, identity, reduce, transform, stream, workspace_size);
+}
+
+/// Helper to compute the sum of the squares of the elements of a tensor
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = Element
+>
+ComputeType TensorSumSq(
+  TensorView<Element, Layout> view,
+  ComputeType identity = ComputeType(),
+  musaStream_t stream = nullptr,
+  int workspace_size = 0
+) {
+
+  plus<ComputeType> reduce;
+  magnitude_squared<Element, ComputeType> transform;
+
+  return TensorTransformReduce(
+    view, identity, reduce, transform, stream, workspace_size);
+}
+
+/// Helper to compute the norm of the elements of a tensor.
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = double
+>
+ComputeType TensorNorm(
+  TensorView<Element, Layout> view,
+  ComputeType identity = ComputeType(),
+  musaStream_t stream = nullptr,
+  int workspace_size = 0
+) {
+
+  return std::sqrt(TensorSumSq(view, identity, stream, workspace_size));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper to compute the sum of the squares of the differences of two tensors
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = double
+>
+ComputeType TensorSumSqDiff(
+  TensorView<Element, Layout> view_A,
+  TensorView<Element, Layout> view_B,
+  ComputeType identity = ComputeType(),
+  musaStream_t stream = nullptr,
+  int workspace_size = 0
+) {
+
+  plus<ComputeType> reduce;
+  magnitude_squared_difference<Element, ComputeType> transform;
+
+  return TensorTransformReduce(
+    view_A, view_B, identity, reduce, transform, stream, workspace_size);
+}
+
+
+/// Helper to compute the norm of the tensor computed as the difference of two tensors in memory
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = double
+>
+ComputeType TensorNormDiff(
+  TensorView<Element, Layout> view_A,
+  TensorView<Element, Layout> view_B,
+  ComputeType identity = ComputeType(),
+  musaStream_t stream = nullptr,
+  int workspace_size = 0
+) {
+
+  return std::sqrt(TensorSumSqDiff(view_A, view_B, identity, stream, workspace_size));
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace reference
+} // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_relu.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_relu.h
new file mode 100644
index 000000000..5dc35b258
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/tensor_relu.h
@@ -0,0 +1,142 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 device-side elementwise operations on TensorView. Note, the operations defined
+    in this header are not specialized for any particular data layout and are therefore not
+    intended to offer the best possible performance. Rather, they are intended to be generic
+    reference implementations to support the MUTLASS unit tests.
+*/
+
+#pragma once
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/tensor_view.h"
+
+#include "mutlass/util/reference/device/tensor_foreach.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace reference {
+namespace device {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorReLuFunc {
+
+  /// View type
+  using TensorView = TensorView<Element, Layout>;
+
+  /// Coordinate in tensor's index space
+  using TensorCoord = typename TensorView::TensorCoord;
+
+  /// Parameters structure
+  struct Params {
+
+    //
+    // Data members
+    //
+
+    TensorView view;
+    Element threshold;
+
+
+    //
+    // Methods
+    //
+
+    Params(
+      TensorView view_ = TensorView(),
+      Element threshold_ = Element(0)
+    ):
+      view(view_), threshold(threshold_) {
+
+    }
+  };
+
+  //
+  // Data members
+  //
+
+  Params params;
+
+  //
+  // Methods
+  //
+
+  MUTLASS_DEVICE
+  TensorReLuFunc(Params const &params): params(params) {
+
+  }
+
+  MUTLASS_DEVICE
+  void operator()(TensorCoord const &coord) {
+
+    Element const & value = params.view.at(coord);
+    params.view.at(coord) = (value < params.threshold) ? params.threshold : value;
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Apply ReLu on a tensor
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorReLu(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  Element threshold = Element(0)) {         ///< ReLu threshold
+
+  using Func = detail::TensorReLuFunc<Element, Layout>;
+  using Params = typename Func::Params;
+
+  TensorForEach<Func, Layout::kRank, Params>(
+    view.extent(),
+    Params(view, threshold)
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/thread/gemm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/thread/gemm.h
new file mode 100644
index 000000000..4333ad519
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/device/thread/gemm.h
@@ -0,0 +1,187 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for GEMM in host-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace device {
+namespace thread {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Thread-level blocked general matrix product.
+//
+// Note, this is a reference implementation. Performance is not expected to approach peak.
+//
+template <
+  typename TensorRefA,
+  typename TensorRefB,
+  typename TensorRefC,
+  typename ScalarType,
+  typename AccumulatorType,
+  typename OutputTile,
+  typename InnerProductOp = multiply_add<AccumulatorType>,
+  typename ConvertOp = NumericConverter<typename TensorRefC::Element, ScalarType>
+>
+struct Gemm {
+
+  using ElementA = typename TensorRefA::Element;
+  using ElementB = typename TensorRefB::Element;
+  using ElementC = typename TensorRefC::Element;
+
+  //
+  // Data members
+  //
+
+  /// Tile for A operand
+  ElementA A_tile[OutputTile::kColumn];
+
+  /// Tile for B operand
+  ElementB B_tile[OutputTile::kRow];
+
+  /// Tile for Accumulator
+  AccumulatorType accum[OutputTile::kColumn][OutputTile::kRow];
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  MUTLASS_HOST_DEVICE
+  Gemm(AccumulatorType initial_accum = AccumulatorType(0)) {
+
+    // Clear fetch registers
+    for (int i = 0; i < OutputTile::kColumn; ++i) {
+      A_tile[i] = ElementA(0);
+    }
+
+    for (int j = 0; j < OutputTile::kColumn; ++j) {
+      B_tile[j] = ElementB(0);
+    }
+
+    // Clear accumulators
+    MUTLASS_PRAGMA_UNROLL
+    for (int j = 0; j < OutputTile::kColumn; ++j) {
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < OutputTile::kRow; ++i) {
+        accum[j][i] = initial_accum;
+      }
+    }
+  }
+
+  /// Computes a matrix product
+  MUTLASS_HOST_DEVICE
+  Gemm & multiply_add(
+    gemm::GemmCoord problem_size,
+    TensorRefA tensor_a,
+    TensorRefB tensor_b,
+    MatrixCoord output_coord = MatrixCoord()) {
+
+    InnerProductOp inner_product_op;
+
+    // Loop over the GEMM K dimension
+    MUTLASS_PRAGMA_NO_UNROLL
+    for (int k = 0; k < problem_size.k(); ++k) {
+
+      // Fetch a slice of the A matrix
+      MUTLASS_PRAGMA_UNROLL
+      for (int i = 0; i < OutputTile::kColumn; ++i) {
+        if (output_coord.row() + i < problem_size.m()) {
+          A_tile[i] = tensor_a.at(make_Coord(output_coord.row() + i, k));
+        }
+      }
+
+      // Fetch a slice of the B matrix
+      MUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < OutputTile::kRow; ++j) {
+        if (output_coord.column() + j < problem_size.n()) {
+          B_tile[j] = tensor_b.at(make_Coord(k, output_coord.column() + j));
+        }
+      }
+
+      // Compute an accumulated matrix product
+      MUTLASS_PRAGMA_UNROLL
+      for (int j = 0; j < OutputTile::kRow; ++j) {
+        MUTLASS_PRAGMA_UNROLL
+        for (int i = 0; i < OutputTile::kColumn; ++i) {
+          accum[j][i] = inner_product_op(A_tile[i], B_tile[j], accum[j][i]);
+        }
+      }
+    }
+
+    return *this;
+  }
+
+  /// Performs linear scaling of matrix product and updates output tensor
+  MUTLASS_HOST_DEVICE
+  Gemm & epilogue(
+    gemm::GemmCoord problem_size,
+    ScalarType alpha,
+    ScalarType beta,
+    TensorRefC tensor_c,
+    TensorRefC tensor_d,
+    MatrixCoord output_coord = MatrixCoord()) {
+
+    ConvertOp convert_op;
+
+    // Update the output tensor
+    for (int j = 0; j < OutputTile::kRow; ++j) {
+      for (int i = 0; i < OutputTile::kColumn; ++i) {
+        MatrixCoord coord = output_coord + MatrixCoord(i, j);
+        if (coord.row() < problem_size.m() && coord.column() < problem_size.n()) {
+
+          tensor_d.at(coord) = convert_op(
+            alpha * ScalarType(accum[j][i]) +
+            beta * ScalarType(tensor_c.at(coord))
+          );
+        }
+      }
+    }
+
+    return *this;
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace thread
+} // namespace device
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/convolution.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/convolution.h
new file mode 100644
index 000000000..01ac0e8c9
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/convolution.h
@@ -0,0 +1,797 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for convolution in host-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/functional.h"
+#include "mutlass/layout/tensor.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/tensor_ref.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/conv/convolution.h"
+#include "mutlass/conv/conv2d_problem_size.h"
+#include "mutlass/conv/conv3d_problem_size.h"
+#include <iostream>
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// Forward propagation
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// y = conv2d(x, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ElementD = ElementC,
+  typename ConvertOp = NumericConverter<ElementD, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+void Conv2dFprop(
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_x,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_y_in,
+  TensorRef<ElementD, LayoutC> tensor_y_out,
+  ElementCompute alpha,
+  ElementCompute beta) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  // Apply MMA and accumulate ElementAccumulator
+  for (int n = 0; n < problem_size.N; ++n) {
+    for (int p = 0; p < problem_size.P; ++p) {
+      for (int q = 0; q < problem_size.Q; ++q) {
+        for (int k = 0; k < problem_size.K; ++k) {
+
+          int group_idx = k / (problem_size.K / problem_size.groups);
+          int channels_per_group = problem_size.C / problem_size.groups;
+
+          ElementAccumulator acc = ElementAccumulator();
+
+          for (int r = 0; r < problem_size.R; ++r) {
+            for (int s = 0; s < problem_size.S; ++s) {
+              for (int c = 0; c < channels_per_group; ++c) {
+
+                int filter_r = r;
+                int filter_s = s;
+
+                if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+                  filter_r = problem_size.R - 1 - r;
+                  filter_s = problem_size.S - 1 - s;
+                }
+
+                int h = p * problem_size.stride_h - problem_size.pad_h + filter_r * problem_size.dilation_h;
+                int w = q * problem_size.stride_w - problem_size.pad_w + filter_s * problem_size.dilation_w;
+
+                if (h >= 0 && h < problem_size.H && w >= 0 && w < problem_size.W) {
+
+                  ElementA a = tensor_x.at({n, h, w, c + group_idx * channels_per_group});
+                  ElementB b = tensor_w.at({k, r, s, c});
+
+                  acc = inner_product_op(ElementAccumulator(a), ElementAccumulator(b), acc);
+
+                }
+              }
+            }
+          }
+
+          // Apply Epilogue, compute ElementCompute, convert and store ElementC
+          ElementC c_ref = ElementC();
+
+          if (beta != ElementCompute()) {
+            c_ref = tensor_y_in.at(mutlass::make_Coord(n, p, q, k));
+          }
+
+          tensor_y_out.at(mutlass::make_Coord(n, p, q, k)) =
+              convert_op(alpha * ElementCompute(acc) + beta * ElementCompute(c_ref));
+        }
+      }
+    }
+  }
+}
+
+/// Depthwise-separable convolution
+template <typename ElementA,
+          typename LayoutA,
+          typename ElementB,
+          typename LayoutB,
+          typename ElementC,
+          typename LayoutC,
+          typename ElementCompute,
+          typename ElementAccumulator = ElementCompute,
+          typename ElementD = ElementC,
+          typename ConvertOp = NumericConverter<ElementD, ElementCompute>,
+          typename InnerProductOp = multiply_add<ElementAccumulator>>
+void Depsep_Fprop(mutlass::TensorView<ElementA, LayoutA> tensor_A,
+                  mutlass::TensorView<ElementB, LayoutB> tensor_B,
+                  mutlass::TensorView<ElementC, LayoutC> tensor_C,
+                  mutlass::TensorView<ElementD, LayoutC> tensor_D,
+                  ElementCompute alpha,
+                  ElementCompute beta,
+                  mutlass::Tensor4DCoord padding = mutlass::Tensor4DCoord(),
+                  mutlass::Coord<2> conv_stride = mutlass::Coord<2>(),
+                  mutlass::Coord<2> dilation = mutlass::Coord<2>(),
+                  mutlass::conv::Mode mode = mutlass::conv::Mode::kCrossCorrelation) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  // Apply MMA and accumulate ElementAccumulator
+  for (int n = 0; n < tensor_C.extent().n(); ++n) {
+    for (int p = 0; p < tensor_C.extent().h(); ++p) {
+      for (int q = 0; q < tensor_C.extent().w(); ++q) {
+        for (int g = 0; g < tensor_C.extent().c(); ++g) {
+          ElementAccumulator acc = ElementAccumulator();
+          for (int r = 0; r < tensor_B.extent().h(); ++r) {
+            for (int s = 0; s < tensor_B.extent().w(); ++s) {
+
+              // input activation H and W
+              int h = p * conv_stride[0] - padding[0] + r * dilation[0];
+              int w = q * conv_stride[1] - padding[2] + s * dilation[1];
+
+              if (h < tensor_A.extent().h() && h >= 0 && w < tensor_A.extent().w() && w >= 0) {
+                ElementA a = tensor_A.at(mutlass::make_Coord(n, h, w, g));
+
+                ElementB b = (mode == mutlass::conv::Mode::kCrossCorrelation)
+                                   ? tensor_B.at(mutlass::make_Coord(g, r, s, 0))
+                                   : tensor_B.at(mutlass::make_Coord(
+                                         g, tensor_B.extent().h() - r - 1, tensor_B.extent().w() - s - 1, 0));
+
+                acc = inner_product_op(ElementAccumulator(a), ElementAccumulator(b), acc);
+              }
+            }
+          }
+
+          // Apply Epilogue, compute ElementCompute, convert and store ElementC
+          ElementC c_ref = tensor_C.at(mutlass::make_Coord(n, p, q, g));
+          tensor_D.at(mutlass::make_Coord(n, p, q, g)) =
+              convert_op(alpha * ElementCompute(acc) + beta * ElementCompute(c_ref));
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// Dgrad
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// dx = dgrad(dy, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ElementD = ElementC,
+  typename ConvertOp = NumericConverter<ElementD, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+void Conv2dDgrad(
+  mutlass::conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_dx_in,
+  TensorRef<ElementD, LayoutC> tensor_dx_out,
+  ElementCompute alpha,
+  ElementCompute beta) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  // Apply MMA and accumulate ElementAccumulator
+  for (int n = 0; n < problem_size.N; ++n) {
+    for (int h = 0; h < problem_size.H; ++h) {
+      for (int w = 0; w < problem_size.W; ++w) {
+        for (int c = 0; c < problem_size.C; ++c) {
+
+          ElementAccumulator acc = ElementAccumulator();
+
+          for (int r = 0; r < problem_size.R; ++r) {
+            for (int s = 0; s < problem_size.S; ++s) {
+              for (int k = 0; k < problem_size.K; ++k) {
+
+                int filter_r = r;
+                int filter_s = s;
+
+                if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+                  filter_r = problem_size.R - 1 - r;
+                  filter_s = problem_size.S - 1 - s;
+                }
+
+                int p = h + problem_size.pad_h - filter_r * problem_size.dilation_h;
+                int q = w + problem_size.pad_w - filter_s * problem_size.dilation_w;
+
+                if (p >= 0 && (p % problem_size.stride_h) == 0 &&
+                    q >= 0 && (q % problem_size.stride_w) == 0) {
+
+                  p = p / problem_size.stride_h;
+                  q = q / problem_size.stride_w;
+#if 0
+                  std::cout << "row:"
+                  << n * problem_size.H * problem_size.W +
+                    h * problem_size.W +
+                    w << " "
+                  << "n, p, q: ("
+                  << n << ", "
+                  << p << ", "
+                  << q << ") * "
+                  << "r, s: ("
+                  << r << ", "
+                  << s << ") ["
+                  << ((p < problem_size.P && q < problem_size.Q) ? "true":"false") << "]"
+                  << std::endl;
+#endif
+                  if (p < problem_size.P && q < problem_size.Q) {
+
+                    ElementA a = tensor_dy.at(mutlass::make_Coord(n, p, q, k));
+                    ElementB b = tensor_w.at(mutlass::make_Coord(k, r, s, c));
+
+                    acc = inner_product_op(ElementAccumulator(a), ElementAccumulator(b), acc);
+                  }
+                }
+
+              } // for (K)
+            } // for (S)
+          } // for (R)
+
+          // Apply Epilogue, compute ElementCompute, convert and store ElementC
+          ElementC c_ref = ElementC();
+
+          if (beta != ElementCompute()) {
+            c_ref = tensor_dx_in.at(mutlass::make_Coord(n, h, w, c));
+          }
+
+          tensor_dx_out.at(mutlass::make_Coord(n, h, w, c)) =
+              convert_op(alpha * ElementCompute(acc) + beta * ElementCompute(c_ref));
+
+        } // for (C)
+      } // for (W)
+    } // for (H)
+  } // for (N)
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// Wgrad
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// dw = wgrad(dy, x)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ElementD = ElementC,
+  typename ConvertOp = NumericConverter<ElementD, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+void Conv2dWgrad(
+  mutlass::conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_x,
+  TensorRef<ElementC, LayoutC> tensor_dw_in,
+  TensorRef<ElementD, LayoutC> tensor_dw_out,
+  ElementCompute alpha,
+  ElementCompute beta) {
+
+  InnerProductOp inner_product_op;
+  ConvertOp convert_op;
+
+  // Apply MMA and accumulate ElementAccumulator
+  for (int k = 0; k < problem_size.K; ++k) {
+    for (int r = 0; r < problem_size.R; ++r) {
+      for (int s = 0; s < problem_size.S; ++s) {
+        for (int c = 0; c < problem_size.C; ++c) {
+
+          ElementAccumulator acc = ElementAccumulator();
+
+          for (int n = 0; n < problem_size.N; ++n) {
+            for (int p = 0; p < problem_size.P; ++p) {
+              for (int q = 0; q < problem_size.Q; ++q) {
+
+                mutlass::Tensor4DCoord b_coord;
+
+                int filter_r = r;
+                int filter_s = s;
+
+                if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+                  filter_r = problem_size.R - 1 - r;
+                  filter_s = problem_size.S - 1 - s;
+                }
+
+                b_coord = make_Coord(
+                    n,
+                    p * problem_size.stride_h - problem_size.pad_h + filter_r * problem_size.dilation_h,
+                    q * problem_size.stride_w - problem_size.pad_w + filter_s * problem_size.dilation_w,
+                    c);
+
+                if (b_coord.h() < problem_size.H && b_coord.h() >= 0 &&
+                    b_coord.w() < problem_size.W && b_coord.w() >= 0) {
+
+                  ElementAccumulator a = ElementAccumulator(tensor_dy.at(mutlass::make_Coord(n, p, q, k)));
+                  ElementAccumulator b = ElementAccumulator(tensor_x.at(b_coord));
+                  acc = inner_product_op(a, b, acc);
+                }
+              }
+            }
+          }
+
+          // Apply Epilogue, compute ElementCompute, convert and store ElementC
+          ElementC c_ref = ElementC();
+
+          if (beta != ElementCompute()) {
+            c_ref = tensor_dw_in.at(mutlass::make_Coord(k, r, s, c));
+          }
+
+          tensor_dw_out.at(mutlass::make_Coord(k, r, s, c)) =
+              convert_op(alpha * ElementCompute(acc) + beta * ElementCompute(c_ref));
+
+        } // for (C)
+      } // for (S)
+    } // for (R)
+  } // for (K)
+}
+
+/// Generic 2D convolution targeting Conv2dFprop, Conv2dDgrad, and Conv2dWgrad.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ElementD = ElementC,
+  typename ConvertOp = NumericConverter<ElementD, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+void Conv2d(
+  conv::Operator convolutional_operator,
+  conv::Conv2dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_A,
+  TensorRef<ElementB, LayoutB> tensor_B,
+  TensorRef<ElementC, LayoutC> tensor_C,
+  TensorRef<ElementD, LayoutC> tensor_D,
+  ElementCompute alpha,
+  ElementCompute beta) {
+
+  switch (convolutional_operator) {
+  case conv::Operator::kFprop:
+    Conv2dFprop<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ElementD,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta);
+    break;
+
+  case conv::Operator::kDgrad:
+    Conv2dDgrad<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ElementD,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta);
+    break;
+
+  case conv::Operator::kWgrad:
+    Conv2dWgrad<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ElementD,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta);
+    break;
+
+  default:
+    break;
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// 3D convolution
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// y = conv3d(x, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+void Conv3dFprop(
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_x,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_y_in,
+  TensorRef<ElementC, LayoutC> tensor_y_out,
+  ElementCompute alpha,
+  ElementCompute beta) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  // Apply MMA and accumulate ElementAccumulator
+  for (int n = 0; n < problem_size.N; ++n) {
+    for (int z = 0; z < problem_size.Z; ++z) {
+      for (int p = 0; p < problem_size.P; ++p) {
+        for (int q = 0; q < problem_size.Q; ++q) {
+          for (int k = 0; k < problem_size.K; ++k) {
+
+            ElementAccumulator acc = ElementAccumulator();
+
+            for (int t = 0; t < problem_size.T; ++t) {
+              for (int r = 0; r < problem_size.R; ++r) {
+                for (int s = 0; s < problem_size.S; ++s) {
+                  for (int c = 0; c < problem_size.C; ++c) {
+
+                    int filter_t = t;
+                    int filter_r = r;
+                    int filter_s = s;
+
+                    if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+                      filter_t = problem_size.T - 1 - t;
+                      filter_r = problem_size.R - 1 - r;
+                      filter_s = problem_size.S - 1 - s;
+                    }
+
+                    int d = z * problem_size.stride_d - problem_size.pad_d + filter_t * problem_size.dilation_d;
+                    int h = p * problem_size.stride_h - problem_size.pad_h + filter_r * problem_size.dilation_h;
+                    int w = q * problem_size.stride_w - problem_size.pad_w + filter_s * problem_size.dilation_w;
+
+                    if (d >= 0 && d < problem_size.D &&
+                      h >=0 && h < problem_size.H &&
+                      w >= 0 && w < problem_size.W) {
+
+                      ElementA a = tensor_x.at({n, d, h, w, c});
+                      ElementB b = tensor_w.at({k, t, r, s, c});
+
+                      acc = inner_product_op(ElementAccumulator(a), ElementAccumulator(b), acc);
+                    }
+                  }
+                }
+              }
+            }
+
+            // Apply Epilogue, compute ElementCompute, convert and store ElementC
+            ElementC c_ref = ElementC();
+
+            if (beta != ElementCompute()) {
+              c_ref = tensor_y_in.at(mutlass::make_Coord(n, z, p, q, k));
+            }
+
+            tensor_y_out.at(mutlass::make_Coord(n, z, p, q, k)) =
+                convert_op(alpha * ElementCompute(acc) + beta * ElementCompute(c_ref));
+          }
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// Dgrad
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// dx = dgrad(dy, w)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+void Conv3dDgrad(
+  mutlass::conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_w,
+  TensorRef<ElementC, LayoutC> tensor_dx_in,
+  TensorRef<ElementC, LayoutC> tensor_dx_out,
+  ElementCompute alpha,
+  ElementCompute beta) {
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  // Apply MMA and accumulate ElementAccumulator
+  for (int n = 0; n < problem_size.N; ++n) {
+    for (int d = 0; d < problem_size.D; ++d) {
+      for (int h = 0; h < problem_size.H; ++h) {
+        for (int w = 0; w < problem_size.W; ++w) {
+          for (int c = 0; c < problem_size.C; ++c) {
+
+            ElementAccumulator acc = ElementAccumulator();
+
+            for (int t = 0; t < problem_size.T; ++t) {
+              for (int r = 0; r < problem_size.R; ++r) {
+                for (int s = 0; s < problem_size.S; ++s) {
+                  for (int k = 0; k < problem_size.K; ++k) {
+
+                    int filter_t = t;
+                    int filter_r = r;
+                    int filter_s = s;
+
+                    if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+                      filter_t = problem_size.T - 1 - t;
+                      filter_r = problem_size.R - 1 - r;
+                      filter_s = problem_size.S - 1 - s;
+                    }
+
+                    int z = d + problem_size.pad_d - filter_t * problem_size.dilation_d;
+                    int p = h + problem_size.pad_h - filter_r * problem_size.dilation_h;
+                    int q = w + problem_size.pad_w - filter_s * problem_size.dilation_w;
+
+                    if (z >= 0 && (z % problem_size.stride_d) == 0 &&
+                        p >= 0 && (p % problem_size.stride_h) == 0 &&
+                        q >= 0 && (q % problem_size.stride_w) == 0) {
+
+                      z = z / problem_size.stride_d;
+                      p = p / problem_size.stride_h;
+                      q = q / problem_size.stride_w;
+
+                      if (z < problem_size.Z && p < problem_size.P && q < problem_size.Q) {
+
+                        ElementA a = tensor_dy.at(mutlass::make_Coord(n, z, p, q, k));
+                        ElementB b = tensor_w.at(mutlass::make_Coord(k, t, r, s, c));
+
+                        acc = inner_product_op(ElementAccumulator(a), ElementAccumulator(b), acc);
+                      }
+                    }
+
+                  } // for (K)
+                } // for (S)
+              } // for (R)
+            } // for (T)
+
+            // Apply Epilogue, compute ElementCompute, convert and store ElementC
+            ElementC c_ref = ElementC();
+
+            if (beta != ElementCompute()) {
+              c_ref = tensor_dx_in.at(mutlass::make_Coord(n, d, h, w, c));
+            }
+
+            tensor_dx_out.at(mutlass::make_Coord(n, d, h, w, c)) =
+                convert_op(alpha * ElementCompute(acc) + beta * ElementCompute(c_ref));
+
+          } // for (C)
+        } // for (W)
+      } // for (H)
+    } // for (D)
+  } // for (N)
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+/// Wgrad
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// dw = wgrad(dy, x)
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+void Conv3dWgrad(
+  mutlass::conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_dy,
+  TensorRef<ElementB, LayoutB> tensor_x,
+  TensorRef<ElementC, LayoutC> tensor_dw_in,
+  TensorRef<ElementC, LayoutC> tensor_dw_out,
+  ElementCompute alpha,
+  ElementCompute beta) {
+
+  InnerProductOp inner_product_op;
+  ConvertOp convert_op;
+
+  // Apply MMA and accumulate ElementAccumulator
+  for (int k = 0; k < problem_size.K; ++k) {
+    for (int t = 0; t < problem_size.T; ++t) {
+      for (int r = 0; r < problem_size.R; ++r) {
+        for (int s = 0; s < problem_size.S; ++s) {
+          for (int c = 0; c < problem_size.C; ++c) {
+
+            ElementAccumulator acc = ElementAccumulator();
+
+            for (int n = 0; n < problem_size.N; ++n) {
+              for (int z = 0; z < problem_size.Z; ++z) {
+                for (int p = 0; p < problem_size.P; ++p) {
+                  for (int q = 0; q < problem_size.Q; ++q) {
+
+                    int filter_t = t;
+                    int filter_r = r;
+                    int filter_s = s;
+
+                    if (problem_size.mode == mutlass::conv::Mode::kConvolution) {
+                      filter_t = problem_size.T - 1 - t;
+                      filter_r = problem_size.R - 1 - r;
+                      filter_s = problem_size.S - 1 - s;
+                    }
+
+                    Tensor5DCoord b_coord = make_Coord(
+                        n,
+                        z * problem_size.stride_d - problem_size.pad_d + filter_t * problem_size.dilation_d,
+                        p * problem_size.stride_h - problem_size.pad_h + filter_r * problem_size.dilation_h,
+                        q * problem_size.stride_w - problem_size.pad_w + filter_s * problem_size.dilation_w,
+                        c);
+
+                    if (b_coord.d() < problem_size.D && b_coord.d() >= 0 &&
+                        b_coord.h() < problem_size.H && b_coord.h() >= 0 &&
+                        b_coord.w() < problem_size.W && b_coord.w() >= 0) {
+
+                      ElementAccumulator a = ElementAccumulator(tensor_dy.at(mutlass::make_Coord(n, z, p, q, k)));
+                      ElementAccumulator b = ElementAccumulator(tensor_x.at(b_coord));
+
+                      acc = inner_product_op(a, b, acc);
+                    }
+                  }
+                }
+              }
+            }
+
+            // Apply Epilogue, compute ElementCompute, convert and store ElementC
+            ElementC c_ref = ElementC();
+
+            if (beta != ElementCompute()) {
+              c_ref = tensor_dw_in.at(mutlass::make_Coord(k, t, r, s, c));
+            }
+
+            tensor_dw_out.at(mutlass::make_Coord(k, t, r, s, c)) =
+                convert_op(alpha * ElementCompute(acc) + beta * ElementCompute(c_ref));
+
+          } // for (C)
+        } // for (S)
+      } // for (R)
+    } // for (T)
+  } // for (K)
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Generic 3D convolution targeting Conv2dFprop, Conv2dDgrad, and Conv2dWgrad.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ElementCompute,
+  typename ElementAccumulator = ElementCompute,
+  typename ConvertOp = NumericConverter<ElementC, ElementCompute>,
+  typename InnerProductOp = multiply_add<ElementAccumulator>
+>
+void Conv3d(
+  conv::Operator convolutional_operator,
+  conv::Conv3dProblemSize problem_size,
+  TensorRef<ElementA, LayoutA> tensor_A,
+  TensorRef<ElementB, LayoutB> tensor_B,
+  TensorRef<ElementC, LayoutC> tensor_C,
+  TensorRef<ElementC, LayoutC> tensor_D,
+  ElementCompute alpha,
+  ElementCompute beta) {
+
+  switch (convolutional_operator) {
+  case conv::Operator::kFprop:
+    Conv3dFprop<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta);
+    break;
+
+  case conv::Operator::kDgrad:
+    Conv3dDgrad<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta);
+    break;
+
+  case conv::Operator::kWgrad:
+    Conv3dWgrad<
+      ElementA, LayoutA,
+      ElementB, LayoutB,
+      ElementC, LayoutC,
+      ElementCompute,
+      ElementAccumulator,
+      ConvertOp, InnerProductOp
+    >(problem_size, tensor_A, tensor_B, tensor_C, tensor_D, alpha, beta);
+    break;
+
+  default:
+    break;
+  }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace host
+}  // namespace reference
+}  // namespace mutlass
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/error_metrics.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/error_metrics.h
new file mode 100644
index 000000000..d99b05d92
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/error_metrics.h
@@ -0,0 +1,67 @@
+
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <cmath>
+
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+#include "mutlass/util/reference/host/tensor_reduce.h"
+#include "mutlass/core_io.h"
+
+namespace mutlass  {
+namespace reference {
+namespace host {
+
+/// Helper to compute the relative error metric for tensor A_computed  w.r.t. to tensor A_reference
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = double
+>
+ComputeType TensorRelativeErrorMetric(
+  TensorView<Element, Layout> view_A_computed,
+  TensorView<Element, Layout> view_B_reference,
+  ComputeType identity = ComputeType()
+) {
+
+  return mutlass::reference::host::TensorNormDiff(view_A_computed, view_B_reference, identity) /
+   mutlass::reference::host::TensorNorm(view_B_reference, identity);
+}
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm.h
new file mode 100644
index 000000000..f679a21b9
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm.h
@@ -0,0 +1,532 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for GEMM in host-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/arch/mma.h"
+#include "mutlass/util/host_tensor.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+template<typename Out, typename In>
+struct CastIfScalar {
+  static Out cast(In in) {
+    return Out(in);
+  }
+};
+
+template<typename OutScalar, typename In>
+struct CastIfScalar<mutlass::complex<OutScalar>, In> {
+  typedef mutlass::complex<OutScalar> Out;
+  static Out cast(In in) {
+    return Out(static_cast<OutScalar>(in));
+  }
+};
+
+template<typename OutScalar, typename InScalar>
+struct CastIfScalar<mutlass::complex<OutScalar>, mutlass::complex<InScalar>> {
+  typedef mutlass::complex<OutScalar> Out;
+  typedef mutlass::complex<InScalar> In;
+  static Out cast(In in) {
+    return Out(in);
+  }
+};
+
+template<typename Out, typename In>
+Out cast_if_scalar(In in) {
+  return CastIfScalar<Out, In>::cast(in);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_gemm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  for (int row_block = 0; row_block < M; row_block += Mblock) {
+    for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+      ComputeType accum[Mblock][Nblock];
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          accum[i][j] = initial_accum;
+        }
+      }
+
+      for (int k_block = 0; k_block < K; ++k_block) {
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            if (row < M && col < N) {
+              ElementA a = tensor_a.at(MatrixCoord(row, k_block));
+              ElementB b = tensor_b.at(MatrixCoord(k_block, col));
+
+              ComputeType compute_a(cast_if_scalar<ComputeType>(a));
+              ComputeType compute_b(cast_if_scalar<ComputeType>(b));
+
+              accum[i][j] = inner_product_op(compute_a, compute_b, accum[i][j]);
+            }
+          }
+        }
+      }
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          int row = row_block + i;
+          int col = col_block + j;
+
+          MatrixCoord coord = MatrixCoord(row, col);
+
+          if (row < M && col < N) {
+            tensor_d.at(coord) = convert_op(
+              alpha * ScalarType(accum[i][j]) +
+              beta * ScalarType(tensor_c.at(coord)));
+          }
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_gemm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  ComputeType initial_accum) {
+  compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+               ScalarType, ComputeType, InnerProductOp, ConvertOp>(
+      problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_c,
+      initial_accum);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = mutlass::arch::OpMultiplyAdd
+>
+struct Gemm;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            ComputeType, arch::OpMultiplyAdd> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            ComputeType, arch::OpMultiplyAddFastBF16> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add-saturate
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            ComputeType, arch::OpMultiplyAddSaturate> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>,
+                 NumericConverterClamp<ElementC, ScalarType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>,
+                 NumericConverterClamp<ElementC, ScalarType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for XOR-popc
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            ComputeType, arch::OpXorPopc> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, xor_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, xor_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+/// Partial specialization for AND-popc
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            ComputeType, arch::OpAndPopc> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, and_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, and_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA, typename ElementB,
+          typename LayoutB, typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct Gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            ComputeType, arch::OpMultiplyAddFastF32> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_gemm<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Batched GEMM
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a batch of GEMMs over a set of matrices of common dimension.
+//
+// TensorRefCollection* is a type satisfying the TensorRefCollection concept.
+//
+template <
+  typename TensorRefCollectionA,
+  typename TensorRefCollectionB,
+  typename TensorRefCollectionC,
+  typename ScalarType,
+  typename AccumulatorType
+>
+void BatchedGemm(
+  gemm::GemmCoord problem_size,
+  int batch_count,
+  ScalarType alpha,
+  TensorRefCollectionA const& tensor_a,
+  TensorRefCollectionB const& tensor_b,
+  ScalarType beta,
+  TensorRefCollectionC &tensor_c,
+  AccumulatorType initial_accum) {
+
+  typename TensorRefCollectionA::ConstIterator tensor_a_it = tensor_a.begin();
+  typename TensorRefCollectionB::ConstIterator tensor_b_it = tensor_b.begin();
+  typename TensorRefCollectionC::ConstIterator tensor_c_it = tensor_c.begin();
+
+  for (int batch = 0;
+    batch < batch_count;
+    ++batch, ++tensor_a_it, ++tensor_b_it, ++tensor_c_it) {
+
+    Gemm<typename TensorRefCollectionA::Element,
+         typename TensorRefCollectionA::Layout,
+         typename TensorRefCollectionB::Element,
+         typename TensorRefCollectionB::Layout,
+         typename TensorRefCollectionC::Element,
+         typename TensorRefCollectionC::Layout,
+         typename TensorRefCollectionC::Element,
+         typename TensorRefCollectionC::Element>
+        gemm;
+
+    gemm(problem_size, alpha, *tensor_a_it, *tensor_b_it, beta, *tensor_c_it,
+         initial_accum);
+  }
+}
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+//
+// TensorRefCollection* is a type satisfying the TensorRefCollection concept.
+//
+template <
+  typename TensorRefCollectionA,
+  typename TensorRefCollectionB,
+  typename TensorRefCollectionC,
+  typename ScalarType,
+  typename AccumulatorType
+>
+void BatchedGemm(
+  gemm::GemmCoord problem_size,
+  int batch_count,
+  ScalarType alpha,
+  TensorRefCollectionA const& tensor_a,
+  TensorRefCollectionB const& tensor_b,
+  ScalarType beta,
+  TensorRefCollectionC &tensor_c) {
+
+  BatchedGemm(problem_size, batch_count, alpha, tensor_a, tensor_b, beta, tensor_c, ScalarType(0));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm_complex.h
new file mode 100644
index 000000000..cd4c24742
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm_complex.h
@@ -0,0 +1,211 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued GEMM in host-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/complex.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/matrix_coord.h"
+
+#include "mutlass/tensor_view.h"
+
+#include "mutlass/gemm/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ElementD = ElementC,
+  typename ConvertOp = NumericConverter<ElementD, ScalarType>,
+  typename InnerProductOp = multiply_add<ComputeType>
+>
+void GemmComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementD, LayoutC> tensor_d,
+  ComputeType initial_accum,
+  int batch_count = 1,
+  int64_t batch_stride_A = 0,
+  int64_t batch_stride_B = 0,
+  int64_t batch_stride_C = 0,
+  int64_t batch_stride_D = 0) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  for (int batch_idx = 0; batch_idx < batch_count; ++batch_idx) {
+
+    // Compute matrix product using blocks
+    for (int row_block = 0; row_block < M; row_block += Mblock) {
+      for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+        ComputeType accum[Mblock][Nblock];
+
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            accum[i][j] = initial_accum;
+          }
+        }
+
+        for (int k_block = 0; k_block < K; ++k_block) {
+          for (int j = 0; j < Nblock; j++) {
+            for (int i = 0; i < Mblock; i++) {
+              int row = row_block + i;
+              int col = col_block + j;
+
+              if (row < M && col < N) {
+                ElementA a = tensor_a.at(MatrixCoord(row, k_block));
+                ElementB b = tensor_b.at(MatrixCoord(k_block, col));
+
+                ComputeType a_ik = ComputeType(a);
+                ComputeType b_kj = ComputeType(b);
+
+                if (transform_a == ComplexTransform::kConjugate) {
+                  a_ik = conj(a_ik);
+                }
+
+                if (transform_b == ComplexTransform::kConjugate) {
+                  b_kj = conj(b_kj);
+                }
+
+                accum[i][j] = inner_product_op(a_ik, b_kj,  accum[i][j]);
+              }
+            }
+          }
+        }
+
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            MatrixCoord coord = MatrixCoord(row, col);
+
+            if (row < M && col < N) {
+
+              tensor_d.at(coord) = convert_op(
+                alpha * ScalarType(accum[i][j]) +
+                beta * ScalarType(tensor_c.at(coord)));
+            }
+          }
+        }
+
+      } // for (col_block)
+    } // for (row_block)
+
+    tensor_a.add_pointer_offset(batch_stride_A);
+    tensor_b.add_pointer_offset(batch_stride_B);
+    tensor_c.add_pointer_offset(batch_stride_C);
+    tensor_d.add_pointer_offset(batch_stride_D);
+
+  } // for (batch_idx)
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// This assumes the accumulator type is the same type as the scalars.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ElementD = ElementC
+>
+void GemmComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementD, LayoutC> tensor_d) {
+
+  GemmComplex(problem_size, alpha, tensor_a, transform_a, tensor_b, transform_b, beta, tensor_c, tensor_d, ScalarType(0));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm_planar_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm_planar_complex.h
new file mode 100644
index 000000000..c845de75a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gemm_planar_complex.h
@@ -0,0 +1,229 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued GEMM in host-side code.
+*/
+
+#pragma once
+
+#include "mutlass/coord.h"
+#include "mutlass/complex.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_ref_planar_complex.h"
+
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>,
+  typename InnerProductOp = multiply_add<complex<ComputeType>>
+>
+void GemmPlanarComplex(
+  gemm::GemmCoord problem_size,
+  complex<ScalarType> alpha,
+  TensorRefPlanarComplex<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRefPlanarComplex<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  complex<ScalarType> beta,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_c,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_d,
+  complex<ComputeType> initial_accum) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  using ComplexA = typename TensorRefPlanarComplex<ElementA, LayoutA>::ComplexElement;
+  using ComplexB = typename TensorRefPlanarComplex<ElementB, LayoutB>::ComplexElement;
+  using ComplexC = typename TensorRefPlanarComplex<ElementC, LayoutC>::ComplexElement;
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  for (int row_block = 0; row_block < M; row_block += Mblock) {
+    for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+      complex<ComputeType> accum[Mblock][Nblock];
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          accum[i][j] = initial_accum;
+        }
+      }
+
+      for (int k_block = 0; k_block < K; ++k_block) {
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            if (row < M && col < N) {
+
+              ComplexA a_ik = tensor_a.at(MatrixCoord(row, k_block));
+              ComplexB b_kj = tensor_b.at(MatrixCoord(k_block, col));
+
+              complex<ComputeType> a = complex<ComputeType>{
+                ComputeType(a_ik.real()),
+                ComputeType(a_ik.imag())
+              };
+
+              complex<ComputeType> b = complex<ComputeType>{
+                ComputeType(b_kj.real()),
+                ComputeType(b_kj.imag())
+              };
+
+              if (transform_a == ComplexTransform::kConjugate) {
+                a = conj(a);
+              }
+
+              if (transform_b == ComplexTransform::kConjugate) {
+                b = conj(b);
+              }
+
+              accum[i][j] = inner_product_op(a, b,  accum[i][j]);
+            }
+          }
+        }
+      }
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          int row = row_block + i;
+          int col = col_block + j;
+
+          MatrixCoord coord = MatrixCoord(row, col);
+
+          if (row < M && col < N) {
+
+            complex<ScalarType> acc{
+              ScalarType(accum[i][j].real()),
+              ScalarType(accum[i][j].imag())
+            };
+
+            ComplexC d_ij = tensor_c.at(coord);
+
+            complex<ScalarType> src{
+              ScalarType(d_ij.real()),
+              ScalarType(d_ij.imag())
+            };
+
+            complex<ScalarType> result = alpha * acc + beta * src;
+
+            d_ij.real() = convert_op(result.real());
+            d_ij.imag() = convert_op(result.imag());
+
+            tensor_d.at(coord) = d_ij;
+          }
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// This assumes the accumulator type is the same type as the scalars.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType
+>
+void GemmPlanarComplex(
+  gemm::GemmCoord problem_size,
+  complex<ScalarType> alpha,
+  TensorRefPlanarComplex<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRefPlanarComplex<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  complex<ScalarType> beta,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_c,
+  TensorRefPlanarComplex<ElementC, LayoutC> tensor_d) {
+
+  GemmPlanarComplex(
+    problem_size,
+    alpha,
+    tensor_a, transform_a,
+    tensor_b, transform_b,
+    beta,
+    tensor_c,
+    tensor_d,
+    complex<ScalarType>());
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gett.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gett.hpp
new file mode 100644
index 000000000..1b92400ee
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/gett.hpp
@@ -0,0 +1,514 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for GETT in host-side code.
+*/
+
+#pragma once
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+#include "mutlass/complex.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/epilogue/thread/activation.h"
+
+#include "mute/tensor.hpp"
+// #include "mute/pointer_base.hpp"
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass::reference::host {
+
+template<class T, class = void>
+struct ElementTraits {
+  using type = T;
+};
+
+template<class T>
+struct ElementTraits<T, std::enable_if_t<!std::is_same_v<decltype(std::declval<T>().get()), void> > >  {
+  using type = decltype(std::declval<T>().get());
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+  class ElementAccumulator_,
+  class TensorA_,                                                                                         // (M, K, L)
+  class TensorB_                                                                                          // (N, K, L)
+>
+struct GettMainloopParams {
+  using ElementAccumulator = ElementAccumulator_;
+  using TensorA = TensorA_;
+  using TensorB = TensorB_;
+  using EngineA = typename TensorA::engine_type;
+  using LayoutA = typename TensorA::layout_type;
+  using EngineB = typename TensorB::engine_type;
+  using LayoutB = typename TensorB::layout_type;
+
+  TensorA A{};
+  TensorB B{};
+
+  ComplexTransform transform_A = ComplexTransform::kNone;
+  ComplexTransform transform_B = ComplexTransform::kNone;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+  class ElementScalar_,
+  class ElementScalingFactor_,
+  class ElementAccumulator_,
+  class ElementCompute_,
+  class TensorC_,                                                                                          // (M, N, L)
+  class TensorD_,                                                                                          // (M, N, L)
+  class VectorBias_ = TensorD_,                                                                            //    (M, 1)
+  class TensorAux_ = TensorD_,                                                                             // (M, N, L)
+  class VectorAlpha_ = TensorD_,                                                                           //    (M, 1)
+  class VectorBeta_ = VectorAlpha_,                                                                        //    (M, 1)
+  class ActivationFunctor_ = mutlass::epilogue::thread::Identity<ElementCompute_>,
+  class BiasBinaryOp_ = mutlass::plus<ElementCompute_>
+>
+struct GettEpilogueParams {
+  using ElementScalar = ElementScalar_;
+  using ElementScalingFactor = ElementScalingFactor_;
+  using ElementAccumulator = ElementAccumulator_;
+  using ElementCompute = ElementCompute_;
+  using TensorC = TensorC_;
+  using TensorD = TensorD_;
+  using TensorAux = TensorAux_;
+  using VectorBias = VectorBias_;
+  using VectorAlpha = VectorAlpha_;
+  using VectorBeta = VectorBeta_;
+  using ActivationFunctor = ActivationFunctor_;
+  using BiasBinaryOp = BiasBinaryOp_;
+
+  using EngineC = typename TensorC::engine_type;
+  using LayoutC = typename TensorC::layout_type;
+  using EngineD =  typename TensorD::engine_type;
+  using LayoutD = typename TensorD::layout_type;
+
+  ElementScalar alpha = ElementScalar(1);
+  ElementScalar beta = ElementScalar(0);
+
+  TensorC C{};
+  TensorD D{};
+  VectorBias Bias{};
+  TensorAux Aux{};
+  VectorAlpha Valpha{};
+  VectorBeta Vbeta{};
+
+  ElementAccumulator* abs_max_D = nullptr;
+  ElementAccumulator* abs_max_Aux = nullptr;
+
+  ElementScalingFactor scale_a = ElementScalingFactor(1);
+  ElementScalingFactor scale_b = ElementScalingFactor(1);
+  ElementScalingFactor scale_c = ElementScalingFactor(1);
+  ElementScalingFactor scale_d = ElementScalingFactor(1);
+  ElementScalingFactor scale_aux = ElementScalingFactor(1);
+
+  bool beta_per_channel_scaling = false;
+};
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// GETT - General Tensor-Tensor contraction reference kernel
+template <
+  class MainloopParams,
+  class EpilogueParams
+>
+void Gett(
+    MainloopParams const& mainloop_params,
+    EpilogueParams const& epilogue_params)
+{
+
+  static int constexpr kBlockM = 64;
+  static int constexpr kBlockN = 64;
+
+#if defined(_OPENMP)
+  #pragma omp parallel for collapse(3)
+#endif
+  for (int64_t l = 0; l < mute::size<2>(mainloop_params.A.layout()); ++l) {
+    for (int64_t m = 0; m < mute::size<0>(mainloop_params.A.layout()); m += kBlockM) {
+      for (int64_t n = 0; n < mute::size<0>(mainloop_params.B.layout()); n += kBlockN) {
+        typename MainloopParams::ElementAccumulator acc[kBlockM][kBlockN];
+        gett_mainloop(mainloop_params, m, n, l, acc);
+        gett_epilogue(epilogue_params, m, n, l, acc);
+      }
+    }
+  }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// GETT - Mainloop
+template <class MainloopParams, class ElementAccumulator, int kBlockM, int kBlockN>
+void gett_mainloop(
+    MainloopParams const& mainloop_params,
+    int64_t m,
+    int64_t n,
+    int64_t l,
+    ElementAccumulator (&acc)[kBlockM][kBlockN])
+{
+
+  static_assert(mute::rank(typename MainloopParams::LayoutA{}) == 3, "M, K, B");
+  static_assert(mute::rank(typename MainloopParams::LayoutB{}) == 3, "N, K, B");
+
+  using ElementA = typename ElementTraits<typename MainloopParams::EngineA::value_type>::type;
+  using ElementB = typename ElementTraits<typename MainloopParams::EngineB::value_type>::type;
+
+  using RingOp = multiply_add<ElementAccumulator, ElementAccumulator, ElementAccumulator>;
+  RingOp fma_op;
+
+  // Zero out accumulators
+  for (int m_b = 0; m_b < kBlockM; ++m_b) {
+    for (int n_b = 0; n_b < kBlockN; ++n_b) {
+      acc[m_b][n_b] = ElementAccumulator(0); // RingOp::AdditionIdentity
+    }
+  }
+
+  // Compute on this k-block
+  for (int64_t k = 0; k < mute::size<1>(mainloop_params.A.layout()); ++k) {
+    // Load A
+    ElementAccumulator a_frag[kBlockM];
+    for (int m_b = 0; m_b < kBlockM; ++m_b) {
+      if (m + m_b < mute::size<0>(mainloop_params.A.layout())) {
+        // Perform reference GEMM calculations at the accumulator's precision. Cast A value to accumulator type.
+        a_frag[m_b] = static_cast<ElementAccumulator>(ElementA(mainloop_params.A(m + m_b, k, l)));
+        if (mainloop_params.transform_A == ComplexTransform::kConjugate) {
+          a_frag[m_b] = conj(a_frag[m_b]);
+        }
+      } else {
+        a_frag[m_b] = ElementAccumulator(0); // RingOp::AdditionIdentity
+      }
+    }
+
+    // Load B
+    ElementAccumulator b_frag[kBlockN];
+    for (int n_b = 0; n_b < kBlockN; ++n_b) {
+      if (n + n_b < mute::size<0>(mainloop_params.B.layout())) {
+        // Perform reference GEMM calculations at the accumulator's precision. Cast A value to accumulator type.
+        b_frag[n_b] = static_cast<ElementAccumulator>(ElementB(mainloop_params.B(n + n_b, k, l)));
+        if (mainloop_params.transform_B == ComplexTransform::kConjugate) {
+          b_frag[n_b] = conj(b_frag[n_b]);
+        }
+      } else {
+        b_frag[n_b] = ElementAccumulator(0); // RingOp::AdditionIdentity
+      }
+    }
+
+    // do compute
+    for (int m_b = 0; m_b < kBlockM; ++m_b) {
+      for (int n_b = 0; n_b < kBlockN; ++n_b) {
+        acc[m_b][n_b] = fma_op(a_frag[m_b], b_frag[n_b], acc[m_b][n_b]);
+      }
+    }
+  }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// GETT - Epilogue
+template <class EpilogueParams, class ElementAccumulator, int kBlockM, int kBlockN>
+void gett_epilogue(
+    EpilogueParams const& epilogue_params,
+    int64_t m,
+    int64_t n,
+    int64_t l,
+    ElementAccumulator (&acc)[kBlockM][kBlockN])
+{
+  static_assert(mute::rank(typename EpilogueParams::LayoutC{}) == 3, "M, K, B");
+  static_assert(mute::rank(typename EpilogueParams::LayoutD{}) == 3, "N, K, B");
+
+  using ElementCompute = typename EpilogueParams::ElementCompute;
+  using ElementC = typename EpilogueParams::TensorC::value_type;
+  using ElementD = typename EpilogueParams::TensorD::value_type;
+  using ElementAux = typename EpilogueParams::TensorAux::value_type;
+  using ElementBias = typename EpilogueParams::VectorBias::value_type;
+  using ElementScalar = typename EpilogueParams::ElementScalar;
+  using ElementScalingFactor = typename EpilogueParams::ElementScalingFactor;
+  using ActivationFunctor = typename EpilogueParams::ActivationFunctor;
+  using BiasBinaryOp = typename EpilogueParams::BiasBinaryOp;
+
+  // constexpr bool IsScalingAndAmaxOutputNeeded =
+  //     std::is_same_v<ElementD, mutlass::float_e4m3_t> or
+  //     std::is_same_v<ElementD, mutlass::float_e5m2_t>;
+
+  // constexpr bool IsScalingAndAmaxAuxOutputNeeded =
+  //     std::is_same_v<ElementAux, mutlass::float_e4m3_t> or
+  //     std::is_same_v<ElementAux, mutlass::float_e5m2_t>;
+
+  constexpr bool IsScalingAndAmaxOutputNeeded = false;
+  constexpr bool IsScalingAndAmaxAuxOutputNeeded = false;
+
+  constexpr bool IsReLUAuxNeeded =
+      mute::is_same_v<ActivationFunctor, mutlass::epilogue::thread::ReLu<ElementCompute>> and
+      mute::is_same_v<ElementAux, mutlass::uint1b_t>;
+
+  constexpr bool IsBackpropFusion =
+      mute::is_same_v<ActivationFunctor, mutlass::epilogue::thread::dGELU<ElementCompute>> or
+      mute::is_same_v<ActivationFunctor, mutlass::epilogue::thread::dReLU<ElementCompute>>;
+
+  // Input related converter
+  NumericConverter<ElementCompute, ElementAccumulator> accumulator_converter;
+  NumericConverter<ElementCompute, ElementC> source_converter;
+  NumericConverter<ElementCompute, ElementBias> bias_converter;
+  NumericConverter<ElementCompute, ElementAux> aux_source_converter;
+
+  // Scale related converter
+  NumericConverter<ElementCompute, ElementScalar> scale_converter;
+  NumericConverter<ElementCompute, ElementScalingFactor> scaling_factor_converter;
+
+  // Abs max converter
+  [[maybe_unused]] NumericConverter<ElementAccumulator, ElementCompute> abs_max_output_converter;
+
+  // Output related converter
+  NumericConverter<ElementD, ElementCompute> destination_converter;
+  NumericConverter<ElementAux, ElementCompute> aux_destination_converter;
+  NumericConverter<ElementBias, ElementCompute> dBias_converter;
+
+  // Epilogue operations
+  multiply_add<ElementCompute, ElementCompute, ElementCompute> epilogue_fma;
+  multiplies<ElementCompute> mul;
+  plus<ElementCompute> add;
+
+  // Activation operation
+  ActivationFunctor activation;
+
+  // Bias binary operation
+  BiasBinaryOp bias_op;
+
+  // Do conversion
+  ElementCompute converted_alpha = scale_converter(epilogue_params.alpha);
+  ElementCompute converted_beta = scale_converter(epilogue_params.beta);
+  ElementCompute converted_scale_a = scaling_factor_converter(epilogue_params.scale_a);
+  ElementCompute converted_scale_b = scaling_factor_converter(epilogue_params.scale_b);
+  ElementCompute converted_scale_c = scaling_factor_converter(epilogue_params.scale_c);
+  ElementCompute converted_scale_d = scaling_factor_converter(epilogue_params.scale_d);
+  ElementCompute converted_scale_aux = scaling_factor_converter(epilogue_params.scale_aux);
+
+  // Init local var
+  [[maybe_unused]] ElementCompute local_abs_max_output = ElementCompute(0);
+  [[maybe_unused]] ElementCompute local_abs_max_aux_output = ElementCompute(0);
+
+  converted_alpha = mul(converted_alpha, mul(converted_scale_a, converted_scale_b));
+  converted_beta = mul(converted_beta, converted_scale_c);
+
+  for (int m_b = 0; m_b < kBlockM; ++m_b) {
+    ElementCompute local_dBias = ElementCompute(0);
+
+    for (int n_b = 0; n_b < kBlockN; ++n_b) {
+      if (m + m_b < mute::size<0>(epilogue_params.D.layout()) && n + n_b < mute::size<1>(epilogue_params.D.layout())) {
+        // Convert every type to ElementCompute first, do compute, convert to output type, write it out
+        ElementCompute converted_acc = accumulator_converter(acc[m_b][n_b]);
+        // per-row alpha
+        if (mute::raw_pointer_cast(epilogue_params.Valpha.data())) {
+          converted_alpha = scale_converter(epilogue_params.Valpha(m + m_b, n + n_b, l));
+          converted_alpha = mul(converted_alpha, mul(converted_scale_a, converted_scale_b));
+        }
+        ElementCompute output = mul(converted_alpha, converted_acc);
+
+        if (mute::raw_pointer_cast(epilogue_params.Bias.data()) && not IsBackpropFusion) {
+          ElementCompute converted_bias = bias_converter(epilogue_params.Bias(m + m_b));
+          output = bias_op(output, converted_bias);
+        }
+
+        if (mute::raw_pointer_cast(epilogue_params.C.data())) {
+          ElementCompute converted_src = source_converter(epilogue_params.C(m + m_b, n + n_b, l));
+          // per-row beta
+          if (epilogue_params.Vbeta.data()) {
+            converted_beta = scale_converter(epilogue_params.Vbeta(m + m_b, n + n_b, l));
+            converted_beta = mul(converted_beta, converted_scale_c);
+          }
+          output = epilogue_fma(converted_beta, converted_src, output);
+        }
+
+        if constexpr (IsBackpropFusion) {
+          ElementAux aux_input = ElementAux(0);
+          if (mute::raw_pointer_cast(epilogue_params.Aux.data())) {
+            aux_input = epilogue_params.Aux(m + m_b, n + n_b, l);
+          }
+
+          output = activation(output, aux_source_converter(aux_input));
+          local_dBias = add(local_dBias, output);
+        }
+        else {
+          if (mute::raw_pointer_cast(epilogue_params.Aux.data())) {
+            auto aux_output = output;
+            if constexpr (IsScalingAndAmaxAuxOutputNeeded) {
+              maximum_absolute_value_reduction<ElementCompute, true> amax_op;
+              local_abs_max_aux_output = amax_op(local_abs_max_aux_output, aux_output);
+              aux_output = epilogue_fma(converted_scale_aux, aux_output, ElementCompute(0));
+            }
+
+            if constexpr (IsReLUAuxNeeded) {
+              epilogue_params.Aux(m + m_b, n + n_b, l) = not (aux_output < 0) ? uint1b_t(1) : uint1b_t(0);
+            } else {
+              epilogue_params.Aux(m + m_b, n + n_b, l) = aux_destination_converter(aux_output);
+            }
+          }
+
+          output = activation(output);
+        }
+
+        if constexpr (IsScalingAndAmaxOutputNeeded) {
+          maximum_absolute_value_reduction<ElementCompute, true> amax_op;
+          local_abs_max_output = amax_op(local_abs_max_output, output);
+          output = epilogue_fma(converted_scale_d, output, ElementCompute(0));
+        }
+
+        epilogue_params.D(m + m_b, n + n_b, l) = destination_converter(output);
+      }
+    } // n_b
+
+    if (m + m_b < mute::size<0>(epilogue_params.D.layout()) && n < mute::size<1>(epilogue_params.D.layout())) {
+      if (mute::raw_pointer_cast(epilogue_params.Bias.data()) && IsBackpropFusion) {
+        ElementCompute converted_dBias = bias_converter(epilogue_params.Bias(m + m_b));
+        local_dBias = add(local_dBias, converted_dBias);
+        epilogue_params.Bias(m + m_b) = dBias_converter(local_dBias);
+      }
+    }
+  } // m_b
+#if defined(_OPENMP)
+  #pragma omp critical(Abs_Max_Data_Update)
+#endif
+  {
+    if constexpr (IsScalingAndAmaxOutputNeeded) {
+      if (epilogue_params.abs_max_D) {
+        *epilogue_params.abs_max_D = maximum_with_nan_propogation<ElementAccumulator>{}(
+          *epilogue_params.abs_max_D, abs_max_output_converter(local_abs_max_output));
+      }
+    }
+
+    if constexpr (IsScalingAndAmaxAuxOutputNeeded) {
+      if (epilogue_params.abs_max_Aux) {
+        *epilogue_params.abs_max_Aux = maximum_with_nan_propogation<ElementAccumulator>{}(
+            *epilogue_params.abs_max_Aux, abs_max_output_converter(local_abs_max_aux_output));
+      }
+    }
+  }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <class TensorType>
+auto make_layout_rank3(const TensorType& tensor) {
+  // append a batch mode of size 1 if we do not have tensors that are rank 3
+  return make_layout(
+      make_shape(mute::get<0>(tensor.shape()), mute::get<1>(tensor.shape()), mute::Int<1>{}),
+      make_stride(mute::get<0>(tensor.stride()), mute::get<1>(tensor.stride()), int64_t(cosize(tensor.layout()))));
+}
+
+/// GEMM - General Matrix-Matrix contraction without conjugation options
+template <
+  class MainloopParams,
+  class EpilogueParams
+>
+void Gemm3x(
+    MainloopParams const& mainloop_params,
+    EpilogueParams const& epilogue_params)
+{
+  using namespace mute;
+
+  static_assert(mute::rank(typename MainloopParams::LayoutA{}) == mute::rank(typename MainloopParams::LayoutB{}));
+  static_assert(mute::rank(typename EpilogueParams::LayoutC{}) == mute::rank(typename EpilogueParams::LayoutD{}));
+  static_assert(mute::rank(typename MainloopParams::LayoutA{}) == mute::rank(typename EpilogueParams::LayoutC{}));
+
+  if constexpr (mute::rank(typename MainloopParams::LayoutA{}) == 2) {
+    Layout layout_A = make_layout_rank3(mainloop_params.A);
+    Layout layout_B = make_layout_rank3(mainloop_params.B);
+    Layout layout_C = make_layout_rank3(epilogue_params.C);
+    Layout layout_D = make_layout_rank3(epilogue_params.D);
+    Layout layout_Aux = make_layout_rank3(epilogue_params.Aux);
+    Layout layout_Bias = make_layout_rank3(epilogue_params.Bias);
+    Layout layout_Valpha = make_layout_rank3(epilogue_params.Valpha);
+    Layout layout_Vbeta = make_layout_rank3(epilogue_params.Vbeta);
+
+    auto TensorA = make_tensor(mainloop_params.A.data(), layout_A);
+    auto TensorB = make_tensor(mainloop_params.B.data(), layout_B);
+    auto TensorC = make_tensor(epilogue_params.C.data(), layout_C);
+    auto TensorD = make_tensor(epilogue_params.D.data(), layout_D);
+    auto TensorAux = make_tensor(epilogue_params.Aux.data(), layout_Aux);
+    auto VectorBias = make_tensor(epilogue_params.Bias.data(), layout_Bias);
+    auto VectorAlpha = make_tensor(epilogue_params.Valpha.data(), layout_Valpha);
+    auto VectorBeta = make_tensor(epilogue_params.Vbeta.data(), layout_Vbeta);
+
+    // Reconstruct mainloop params
+    GettMainloopParams<typename MainloopParams::ElementAccumulator,
+                       decltype(TensorA),
+                       decltype(TensorB)>
+        mainloop_params_converted{TensorA,
+                                  TensorB,
+                                  mainloop_params.transform_A,
+                                  mainloop_params.transform_B};
+
+    // Reconstruct epilogue params
+    GettEpilogueParams<typename EpilogueParams::ElementScalar,
+                       typename EpilogueParams::ElementScalingFactor,
+                       typename EpilogueParams::ElementAccumulator,
+                       typename EpilogueParams::ElementCompute,
+                       decltype(TensorC),
+                       decltype(TensorD),
+                       decltype(VectorBias),
+                       decltype(TensorAux),
+                       decltype(VectorAlpha),
+                       decltype(VectorBeta)
+                      >
+        epilogue_params_converted{epilogue_params.alpha,
+                                  epilogue_params.beta,
+                                  TensorC,
+                                  TensorD,
+                                  VectorBias,
+                                  TensorAux,
+                                  VectorAlpha,
+                                  VectorBeta,
+                                  epilogue_params.abs_amax_D,
+                                  epilogue_params.abs_amax_Aux,
+                                  epilogue_params.scale_a,
+                                  epilogue_params.scale_b,
+                                  epilogue_params.scale_c,
+                                  epilogue_params.scale_d,
+                                  epilogue_params.scale_aux
+                                  };
+
+    Gett(mainloop_params_converted, epilogue_params_converted);
+  }
+  else {
+    // if we already have a batch mode, just pass it through
+    Gett(mainloop_params, epilogue_params);
+  }
+}
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // mutlass::reference::host
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_2k.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_2k.h
new file mode 100644
index 000000000..31432dd60
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_2k.h
@@ -0,0 +1,262 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for Rank 2k update in host-side code.
+
+
+
+*/
+
+#pragma once
+
+#include "mutlass/blas3.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/arch/mma.h"
+#include "mutlass/util/host_tensor.h"
+#include "mutlass/util/reference/host/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  FillMode FillModeC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_rank2k(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2,
+    "Tensors must be of rank 2");
+
+  static_assert(
+    FillModeC == FillMode::kLower ||
+    FillModeC == FillMode::kUpper,
+    "Fill Mode can either be Lower or Upper.");
+
+  using CompareOp = typename platform::conditional<(FillModeC == FillMode::kLower),
+                                                    std::greater_equal<int>,
+                                                    std::less_equal<int>>::type;
+
+  // Note: batch is ignored.
+  // Note: M is same as N for Rank 2k update
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  // Blocking necessary to speedup reference implementation
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+  CompareOp compare_op;
+
+  for (int row_block = 0; row_block < N; row_block += Nblock) {
+    for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+      ComputeType accum[Nblock][Nblock];
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Nblock; i++) {
+          accum[i][j] = initial_accum;
+        }
+      }
+
+      for (int k_block = 0; k_block < K; ++k_block) {
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Nblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            if (row < N && col < N && compare_op(row, col))
+            {
+
+              // A x B^T
+              ElementA a = tensor_a.at(MatrixCoord(row, k_block));
+              ElementB b_t = tensor_b.at(MatrixCoord(col, k_block));
+
+              ComputeType compute_a(cast_if_scalar<ComputeType>(a));
+              ComputeType compute_b_t(cast_if_scalar<ComputeType>(b_t));
+
+              accum[i][j] = inner_product_op(compute_a, compute_b_t, accum[i][j]);
+
+              // B x A^T
+              ElementB b = tensor_b.at(MatrixCoord(row, k_block));
+              ElementA a_t = tensor_a.at(MatrixCoord(col, k_block));
+
+              ComputeType compute_b(cast_if_scalar<ComputeType>(b));
+              ComputeType compute_a_t(cast_if_scalar<ComputeType>(a_t));
+
+              accum[i][j] = inner_product_op(compute_b, compute_a_t, accum[i][j]);
+            }
+          }
+        }
+      }
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Nblock; i++) {
+          int row = row_block + i;
+          int col = col_block + j;
+
+          MatrixCoord coord = MatrixCoord(row, col);
+
+          if (row < N && col < N &&
+              ( (FillModeC == FillMode::kLower && row >= col) ||
+                (FillModeC == FillMode::kUpper && row <= col) )
+          ) {
+            tensor_d.at(coord) = convert_op(
+              alpha * ScalarType(accum[i][j]) +
+              beta * ScalarType(tensor_c.at(coord)));
+          }
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general Rank 2k update (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  FillMode FillModeC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_rank2k(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  ComputeType initial_accum) {
+  compute_rank2k<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, FillModeC,
+               ScalarType, ComputeType, InnerProductOp, ConvertOp>(
+      problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_c,
+      initial_accum);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  FillMode FillModeC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = mutlass::arch::OpMultiplyAdd
+>
+struct Rank2K;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA,
+          typename ElementB, typename LayoutB,
+          typename ElementC, typename LayoutC, FillMode FillModeC,
+          typename ScalarType, typename ComputeType>
+struct Rank2K<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, FillModeC, ScalarType,
+            ComputeType, arch::OpMultiplyAdd> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_rank2k<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, FillModeC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_rank2k<ElementA, LayoutA, ElementB, LayoutB, ElementC, LayoutC, FillModeC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_2k_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_2k_complex.h
new file mode 100644
index 000000000..5b1a8b70a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_2k_complex.h
@@ -0,0 +1,319 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued Rank 2K update in host-side code.
+
+
+*/
+
+#pragma once
+
+#include "mutlass/blas3.h"
+#include "mutlass/complex.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+#include <assert.h>
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>,
+  typename InnerProductOp = multiply_add<ComputeType>
+>
+void Rank2KComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum,
+  FillMode fill_mode_c,
+  BlasMode blas_mode,
+  int batch_count = 1,
+  int64_t batch_stride_A = 0,
+  int64_t batch_stride_B = 0,
+  int64_t batch_stride_C = 0,
+  int64_t batch_stride_D = 0) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  // Rank2K update operates on A=NxK, B=NxK, and C=NxN
+  assert(M==N);
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  for (int batch_idx = 0; batch_idx < batch_count; ++batch_idx) {
+
+    // Compute matrix product using blocks
+    for (int row_block = 0; row_block < M; row_block += Mblock) {
+      for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+        ComputeType accum[Mblock][Nblock];
+
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            accum[i][j] = initial_accum;
+          }
+        }
+
+        for (int k_block = 0; k_block < K; ++k_block) {
+          for (int j = 0; j < Nblock; j++) {
+            for (int i = 0; i < Mblock; i++) {
+              int row = row_block + i;
+              int col = col_block + j;
+
+              if (row < M && col < N &&
+                 ( (fill_mode_c == FillMode::kLower && row >= col) ||
+                  (fill_mode_c == FillMode::kUpper && row <= col) )
+                ) {
+
+                // A x B^T (Symmetric) or A x B^H (Hermitian)
+                // complex conjugation on operandB (b_t) is function of blas3 computation
+                ElementA a = tensor_a.at(MatrixCoord(row, k_block));
+                ElementB b_t = (blas_mode == BlasMode::kHermitian) ?
+                              conj(tensor_b.at(MatrixCoord(col, k_block))) :
+                              tensor_b.at(MatrixCoord(col, k_block));
+
+                ComputeType a_ik = ComputeType(a);
+                ComputeType b_jk = ComputeType(b_t);
+
+                // complex conjugation is a function of operand layouts
+                if (transform_a == ComplexTransform::kConjugate) {
+                  a_ik = conj(a_ik);
+                }
+                // complex conjugation is a function of operand layouts
+                if (transform_b == ComplexTransform::kConjugate) {
+                  b_jk = conj(b_jk);
+                }
+
+                accum[i][j] = inner_product_op(a_ik, b_jk,  accum[i][j]);
+              }
+            }
+          }
+        }
+
+        /* HER2K need two epilogues to handle complex alpha value */
+        if ( blas_mode == BlasMode::kHermitian ) {
+          for (int j = 0; j < Nblock; j++) {
+            for (int i = 0; i < Mblock; i++) {
+              int row = row_block + i;
+              int col = col_block + j;
+
+              MatrixCoord coord = MatrixCoord(row, col);
+
+              if (row < M && col < N &&
+                  ((fill_mode_c == FillMode::kLower && row >= col) ||
+                  (fill_mode_c == FillMode::kUpper && row <= col))
+                ) {
+
+                ScalarType c = tensor_c.at(coord);
+                // The imaginary parts of the diagonal elements of
+                // a complex data type are assumed and set to zero
+                if (blas_mode == BlasMode::kHermitian) {
+                  c = (row == col) ? real(c) : c;
+                }
+
+                tensor_d.at(coord) = convert_op(alpha *
+                  ScalarType(accum[i][j]) +
+                  beta * c);
+              }
+            }
+          }
+
+          /* Zeoring out accum for second HERK */
+          for (int j = 0; j < Nblock; j++) {
+            for (int i = 0; i < Mblock; i++) {
+              accum[i][j] = initial_accum;
+            }
+          }
+        }
+
+        for (int k_block = 0; k_block < K; ++k_block) {
+          for (int j = 0; j < Nblock; j++) {
+            for (int i = 0; i < Mblock; i++) {
+              int row = row_block + i;
+              int col = col_block + j;
+
+              if (row < M && col < N &&
+                 ( (fill_mode_c == FillMode::kLower && row >= col) ||
+                  (fill_mode_c == FillMode::kUpper && row <= col) )
+                ) {
+
+                // B x A^T (Symmetric) or B x A^H (Hermitian)
+                // complex conjugation on operandB (a_t) is function of blas3 computation
+                ElementB b = tensor_b.at(MatrixCoord(row, k_block));
+                ElementA a_t = (blas_mode == BlasMode::kHermitian) ?
+                                conj(tensor_a.at(MatrixCoord(col, k_block))):
+                                tensor_a.at(MatrixCoord(col, k_block));
+
+                ComputeType b_ik = ComputeType(b);
+                ComputeType a_jk = ComputeType(a_t);
+
+                // complex conjugation here is a function of operand layouts
+                if (transform_b == ComplexTransform::kConjugate) {
+                  b_ik = conj(b_ik);
+                }
+                // complex conjugation here is a function of operand layouts
+                if (transform_a == ComplexTransform::kConjugate) {
+                  a_jk = conj(a_jk);
+                }
+
+                accum[i][j] = inner_product_op(b_ik, a_jk, accum[i][j]);
+              }
+            }
+          }
+        }
+
+        ScalarType alpha_hermitian = (blas_mode == BlasMode::kHermitian) ?
+                                      conj(alpha) : alpha;
+        ScalarType beta_hermitian = (blas_mode == BlasMode::kHermitian) ?
+                                      1 : beta;
+
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            MatrixCoord coord = MatrixCoord(row, col);
+
+            if (row < M && col < N &&
+                ((fill_mode_c == FillMode::kLower && row >= col) ||
+                 (fill_mode_c == FillMode::kUpper && row <= col))
+              ) {
+
+              ScalarType d = (blas_mode == BlasMode::kHermitian) ?
+                             tensor_d.at(coord) : tensor_c.at(coord);
+
+              ScalarType tmp_d = convert_op(
+                alpha_hermitian * ScalarType(accum[i][j]) +
+                beta_hermitian * d);
+
+              if (blas_mode == BlasMode::kHermitian && row == col ) {
+                tensor_d.at(coord) = real(tmp_d);
+              } else {
+                tensor_d.at(coord) = tmp_d;
+              }
+            }
+          }
+        }
+
+      } // for (col_block)
+    } // for (row_block)
+
+    tensor_a.add_pointer_offset(batch_stride_A);
+    tensor_b.add_pointer_offset(batch_stride_B);
+    tensor_c.add_pointer_offset(batch_stride_C);
+    tensor_d.add_pointer_offset(batch_stride_D);
+
+  } // for (batch_idx)
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// This assumes the accumulator type is the same type as the scalars.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType
+>
+void Rank2KComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ComplexTransform transform_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  FillMode fill_mode_c,
+  BlasMode blas_mode) {
+
+  Rank2KComplex(
+    problem_size, alpha,
+    tensor_a, transform_a,
+    tensor_b, transform_b,
+    beta, tensor_c, tensor_d,
+    ScalarType(0),
+    fill_mode_c,
+    blas_mode);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_k_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_k_complex.h
new file mode 100644
index 000000000..ac9b4bccb
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/rank_k_complex.h
@@ -0,0 +1,235 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued Rank 2K update in host-side code.
+
+
+*/
+
+#pragma once
+
+#include "mutlass/blas3.h"
+#include "mutlass/complex.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+#include <assert.h>
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>,
+  typename InnerProductOp = multiply_add<ComputeType>
+>
+void Rank2KComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum,
+  FillMode fill_mode_c,
+  BlasMode blas_mode,
+  int batch_count = 1,
+  int64_t batch_stride_A = 0,
+  int64_t batch_stride_C = 0,
+  int64_t batch_stride_D = 0) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  int const K = problem_size.k();
+
+  // Rank2K update operates on A=NxK, B=NxK, and C=NxN
+  assert(M==N);
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+
+  for (int batch_idx = 0; batch_idx < batch_count; ++batch_idx) {
+
+    // Compute matrix product using blocks
+    for (int row_block = 0; row_block < M; row_block += Mblock) {
+      for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+        ComputeType accum[Mblock][Nblock];
+
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            accum[i][j] = initial_accum;
+          }
+        }
+
+        for (int k_block = 0; k_block < K; ++k_block) {
+          for (int j = 0; j < Nblock; j++) {
+            for (int i = 0; i < Mblock; i++) {
+              int row = row_block + i;
+              int col = col_block + j;
+
+              if (row < M && col < N &&
+                 ( (fill_mode_c == FillMode::kLower && row >= col) ||
+                  (fill_mode_c == FillMode::kUpper && row <= col) )
+                ) {
+
+                // A x A^T (Symmetric) or A x A^H (Hermitian)
+                // complex conjugation on operandB (a_t) (function of blas3 computation)
+                ElementA a = tensor_a.at(MatrixCoord(row, k_block));
+                ElementA a_t = (blas_mode == BlasMode::kHermitian) ?
+                              conj(tensor_a.at(MatrixCoord(col, k_block))) :
+                              tensor_a.at(MatrixCoord(col, k_block));
+
+                ComputeType a_ik = ComputeType(a);
+                ComputeType b_jk = ComputeType(a_t);
+
+                // complex conjugation (function of input layouts)
+                if (transform_a == ComplexTransform::kConjugate) {
+                  a_ik = conj(a_ik);
+                }
+                // complex conjugation (function of input layouts)
+                if (transform_a == ComplexTransform::kConjugate) {
+                  b_jk = conj(b_jk);
+                }
+
+                accum[i][j] = inner_product_op(a_ik, b_jk,  accum[i][j]);
+
+              }
+            }
+          }
+        }
+
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            MatrixCoord coord = MatrixCoord(row, col);
+
+            if (row < M && col < N &&
+                ((fill_mode_c == FillMode::kLower && row >= col) ||
+                 (fill_mode_c == FillMode::kUpper && row <= col))
+              ) {
+
+              ScalarType c = tensor_c.at(coord);
+              // The imaginary parts of the diagonal elements of
+              // a complex data type are assumed and set to zero
+              if (blas_mode == BlasMode::kHermitian) {
+                c = (row == col) ? real(c) : c;
+              }
+
+              ScalarType tmp_d = convert_op(
+                alpha * ScalarType(accum[i][j]) +
+                beta * c);
+
+              if (blas_mode == BlasMode::kHermitian && row == col ) {
+                tensor_d.at(coord) = real(tmp_d);
+              } else {
+                tensor_d.at(coord) = tmp_d;
+              }
+            }
+          }
+        }
+
+      } // for (col_block)
+    } // for (row_block)
+
+    tensor_a.add_pointer_offset(batch_stride_A);
+    tensor_c.add_pointer_offset(batch_stride_C);
+    tensor_d.add_pointer_offset(batch_stride_D);
+
+  } // for (batch_idx)
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// This assumes the accumulator type is the same type as the scalars.
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType
+>
+void RankKComplex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  ComplexTransform transform_a,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  FillMode fill_mode_c,
+  BlasMode blas_mode) {
+
+  Rank2KComplex(
+    problem_size, alpha,
+    tensor_a, transform_a,
+    beta, tensor_c, tensor_d,
+    ScalarType(0),
+    fill_mode_c,
+    blas_mode);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/symm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/symm.h
new file mode 100644
index 000000000..f440be8c8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/symm.h
@@ -0,0 +1,286 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for SYMM update in host-side code.
+
+
+
+*/
+
+#pragma once
+
+#include "mutlass/blas3.h"
+#include "mutlass/numeric_conversion.h"
+
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/arch/mma.h"
+#include "mutlass/util/host_tensor.h"
+#include "mutlass/util/reference/host/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename ElementA,
+  typename LayoutA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_symm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2,
+    "Tensors must be of rank 2");
+
+  static_assert(SideModeA != SideMode::kInvalid
+                , "Side Mode can either be Left or Right.");
+
+  static_assert(
+    FillModeA == FillMode::kLower ||
+    FillModeA == FillMode::kUpper,
+    "Fill Mode can either be Lower or Upper.");
+
+  using CompareOp_w_diag =  typename TrMatrixCompareOp<FillModeA, DiagType::kNonUnit>::Type;
+  using CompareOp_wo_diag = typename TrMatrixCompareOp<FillModeA, DiagType::kZero>::Type;
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  // Assuming correct k-dimension value is passed
+  int const K = problem_size.k();
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+  CompareOp_w_diag compare_op_1;
+  CompareOp_wo_diag compare_op_2;
+
+  for (int row_block = 0; row_block < M; row_block += Mblock) {
+    for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+      ComputeType accum[Mblock][Nblock];
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          accum[i][j] = initial_accum;
+        }
+      }
+
+      for (int k_block = 0; k_block < K; ++k_block) {
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            if (row < M && col < N) {
+              ElementA a_1 = ElementA();
+              ElementB b_1 = ElementB();
+              ElementA a_2 = ElementA();
+              ElementB b_2 = ElementB();
+
+              // A x B or B x A (with diagonal)
+              if (SideModeA == SideMode::kLeft) {
+                a_1 = (compare_op_1(row, k_block)) ?
+                      (tensor_a.at(MatrixCoord(row, k_block))) : ElementA();
+                b_1 = tensor_b.at(MatrixCoord(k_block, col));
+              } else if (SideModeA == SideMode::kRight) {
+                a_1 = tensor_b.at(MatrixCoord(row, k_block));
+                b_1 = (compare_op_1(k_block, col)) ?
+                      tensor_a.at(MatrixCoord(k_block, col)) : ElementA();
+              }
+
+              ComputeType compute_a_1(cast_if_scalar<ComputeType>(a_1));
+              ComputeType compute_b_1(cast_if_scalar<ComputeType>(b_1));
+
+              accum[i][j] = inner_product_op(compute_a_1, compute_b_1, accum[i][j]);
+
+              // A^T x B or B x A^T (without diagonal)
+              if (SideModeA == SideMode::kLeft) {
+                a_2 = (compare_op_2(k_block, row)) ?
+                      (tensor_a.at(MatrixCoord(k_block, row))) : ElementA();
+                b_2 = tensor_b.at(MatrixCoord(k_block, col));
+              } else if (SideModeA == SideMode::kRight) {
+                a_2 = tensor_b.at(MatrixCoord(row, k_block));
+                b_2 = (compare_op_2(col, k_block)) ?
+                      tensor_a.at(MatrixCoord(col, k_block)) : ElementA();
+              }
+
+              ComputeType compute_a_2(cast_if_scalar<ComputeType>(a_2));
+              ComputeType compute_b_2(cast_if_scalar<ComputeType>(b_2));
+
+              accum[i][j] = inner_product_op(compute_a_2, compute_b_2, accum[i][j]);
+            }
+          }
+        }
+      }
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          int row = row_block + i;
+          int col = col_block + j;
+
+          MatrixCoord coord = MatrixCoord(row, col);
+
+          if (row < M && col < N) {
+            tensor_d.at(coord) = convert_op(
+              alpha * ScalarType(accum[i][j]) +
+              beta * ScalarType(tensor_c.at(coord)));
+          }
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general Symm update (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename ElementA,
+  typename LayoutA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_symm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  ComputeType initial_accum) {
+  compute_symm<ElementA, LayoutA, SideModeA, FillModeA, ElementB, LayoutB, ElementC, LayoutC,
+               ScalarType, ComputeType, InnerProductOp, ConvertOp>(
+      problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_c,
+      initial_accum);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = mutlass::arch::OpMultiplyAdd
+>
+struct Symm;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA,
+          SideMode SideModeA, FillMode FillModeA,
+          typename ElementB, typename LayoutB,
+          typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct Symm<ElementA, LayoutA, SideModeA, FillModeA, ElementB, LayoutB, ElementC, LayoutC, ScalarType,
+            ComputeType, arch::OpMultiplyAdd> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_symm<ElementA, LayoutA, SideModeA, FillModeA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, initial_accum);
+  }
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutB::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_symm<ElementA, LayoutA, SideModeA, FillModeA, ElementB, LayoutB, ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+        problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/symm_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/symm_complex.h
new file mode 100644
index 000000000..3138b2909
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/symm_complex.h
@@ -0,0 +1,320 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued SYMM update in host-side code.
+
+
+*/
+
+#pragma once
+
+#include "mutlass/blas3.h"
+#include "mutlass/complex.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+#include <assert.h>
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Computes a general matrix product among matrices (tensors of rank=2) pointed to by TensorRef
+/// objects.
+///
+/// Explicitly naming types needed by this template can be cumbersome, particularly for the
+/// accumulator type, so a function argument 'initial_accum' is exposed. Passing
+/// AccumulatorType(0) as the last function argument can be easier than naming all template
+/// arguments explicitly.
+template <
+  typename ElementA,
+  typename LayoutA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  BlasMode BlasMode_ = BlasMode::kSymmetric,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_symm_complex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  ScalarType beta,
+  TensorRef<ElementC, LayoutC> tensor_c,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum,
+  int batch_count = 1,
+  int64_t batch_stride_A = 0,
+  int64_t batch_stride_B = 0,
+  int64_t batch_stride_C = 0,
+  int64_t batch_stride_D = 0) {
+
+  static SideMode const kSideModeA = SideModeA;
+  static FillMode const kFillModeA = FillModeA;
+  static BlasMode const kBlasMode  = BlasMode_;
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutB::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  static_assert(kSideModeA != SideMode::kInvalid
+                , "Side Mode can either be Left or Right.");
+
+  static_assert(
+    kFillModeA == FillMode::kLower ||
+    kFillModeA == FillMode::kUpper,
+    "Fill Mode can either be Lower or Upper.");
+
+  using CompareOp_w_diag =  typename TrMatrixCompareOp<kFillModeA, DiagType::kNonUnit>::Type;
+  using CompareOp_wo_diag = typename TrMatrixCompareOp<kFillModeA, DiagType::kZero>::Type;
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  // Assuming correct k-dimension value is passed
+  int const K = problem_size.k();
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+  CompareOp_w_diag compare_op_1;
+  CompareOp_wo_diag compare_op_2;
+
+  for (int batch_idx = 0; batch_idx < batch_count; ++batch_idx) {
+
+    // Compute matrix product using blocks
+    for (int row_block = 0; row_block < M; row_block += Mblock) {
+      for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+        ComputeType accum[Mblock][Nblock];
+
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            accum[i][j] = initial_accum;
+          }
+        }
+
+        for (int k_block = 0; k_block < K; ++k_block) {
+          for (int j = 0; j < Nblock; j++) {
+            for (int i = 0; i < Mblock; i++) {
+              int row = row_block + i;
+              int col = col_block + j;
+
+              if (row < M && col < N)
+              {
+                ElementA a_1 = ElementA();
+                ElementB b_1 = ElementB();
+                ElementA a_2 = ElementA();
+                ElementB b_2 = ElementB();
+
+                // A x B or B x A (with diagonal)
+                if (kSideModeA == SideMode::kLeft) {
+                  a_1 = (compare_op_1(row, k_block)) ?
+                        (tensor_a.at(MatrixCoord(row, k_block))) : ElementA();
+                  b_1 = tensor_b.at(MatrixCoord(k_block, col));
+                } else if (kSideModeA == SideMode::kRight) {
+                  a_1 = tensor_b.at(MatrixCoord(row, k_block));
+                  b_1 = (compare_op_1(k_block, col)) ?
+                        tensor_a.at(MatrixCoord(k_block, col)) : ElementA();
+                }
+                ComputeType compute_a_1 = ComputeType(a_1);
+                ComputeType compute_b_1 = ComputeType(b_1);
+
+                // The imaginary parts of the diagonal elements of
+                // a complex data type are assumed and set to zero
+                if (kBlasMode == BlasMode::kHermitian && kSideModeA == SideMode::kLeft && row == k_block) {
+                  compute_a_1 = real(compute_a_1);
+                } else if (kBlasMode == BlasMode::kHermitian && kSideModeA == SideMode::kRight && k_block == col) {
+                  compute_b_1 = real(compute_b_1);
+                }
+
+                accum[i][j] = inner_product_op(compute_a_1, compute_b_1,  accum[i][j]);
+
+                // A^T x B or B x A^T (without diagonal)
+                if (kSideModeA == SideMode::kLeft) {
+                  a_2 = (compare_op_2(k_block, row)) ?
+                        (tensor_a.at(MatrixCoord(k_block, row))) : ElementA();
+                  b_2 = tensor_b.at(MatrixCoord(k_block, col));
+                  if (kBlasMode == BlasMode::kHermitian)
+                    a_2 = conj(a_2);
+                } else if (kSideModeA == SideMode::kRight) {
+                  a_2 = tensor_b.at(MatrixCoord(row, k_block));
+                  b_2 = (compare_op_2(col, k_block)) ?
+                        tensor_a.at(MatrixCoord(col, k_block)) : ElementA();
+                  if (kBlasMode == BlasMode::kHermitian)
+                    b_2 = conj(b_2);
+                }
+
+                ComputeType compute_a_2 = ComputeType(a_2);
+                ComputeType compute_b_2 = ComputeType(b_2);
+
+                accum[i][j] = inner_product_op(compute_a_2, compute_b_2, accum[i][j]);
+              }
+            }
+          }
+        }
+
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            MatrixCoord coord = MatrixCoord(row, col);
+
+            if (row < M && col < N) {
+
+              ScalarType c = tensor_c.at(coord);
+
+              tensor_d.at(coord) = convert_op(
+                alpha * ScalarType(accum[i][j]) +
+                beta * c);
+            }
+          }
+        }
+
+      } // for (col_block)
+    } // for (row_block)
+
+    tensor_a.add_pointer_offset(batch_stride_A);
+    tensor_b.add_pointer_offset(batch_stride_B);
+    tensor_c.add_pointer_offset(batch_stride_C);
+    tensor_d.add_pointer_offset(batch_stride_D);
+
+  } // for (batch_idx)
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  BlasMode BlasMode_ = mutlass::BlasMode::kSymmetric,
+  typename InnerProductOp = mutlass::arch::OpMultiplyAddComplex
+>
+struct SymmComplex;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA,
+          SideMode SideModeA, FillMode FillModeA,
+          typename ElementB, typename LayoutB,
+          typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType,
+          BlasMode BlasMode_>
+struct SymmComplex<ElementA, LayoutA,
+                   SideModeA, FillModeA,
+                   ElementB, LayoutB,
+                   ElementC, LayoutC, ScalarType,
+                   ComputeType, BlasMode_,
+                   arch::OpMultiplyAddComplex> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_symm_complex<ElementA, LayoutA,
+                 SideModeA, FillModeA,
+                 ElementB, LayoutB,
+                 ElementC, LayoutC,
+                 ScalarType, ComputeType, BlasMode_, multiply_add<ComputeType>>(
+                 problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for gaussian multiply-add
+template <typename ElementA, typename LayoutA,
+          SideMode SideModeA, FillMode FillModeA,
+          typename ElementB, typename LayoutB,
+          typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType,
+          BlasMode BlasMode_>
+struct SymmComplex<ElementA, LayoutA,
+                   SideModeA, FillModeA,
+                   ElementB, LayoutB,
+                   ElementC, LayoutC, ScalarType,
+                   ComputeType, BlasMode_,
+                   arch::OpMultiplyAddGaussianComplex> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b, ScalarType beta,
+                  TensorRef<ElementC, LayoutC> tensor_c,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_symm_complex<ElementA, LayoutA,
+                 SideModeA, FillModeA,
+                 ElementB, LayoutB,
+                 ElementC, LayoutC,
+                 ScalarType, ComputeType, BlasMode_, multiply_add<ComputeType>>(
+                 problem_size, alpha, tensor_a, tensor_b, beta, tensor_c, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_compare.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_compare.h
new file mode 100644
index 000000000..8d4c90d13
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_compare.h
@@ -0,0 +1,424 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 host-side elementwise operations on TensorView.
+*/
+
+#pragma once
+
+// Standard Library includes
+#include <utility>
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/relatively_equal.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/tensor_view_planar_complex.h"
+
+#include "mutlass/util/distribution.h"
+#include "tensor_foreach.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorEqualsFunc {
+
+  //
+  // Data members
+  //
+
+  TensorView<Element, Layout> lhs;
+  TensorView<Element, Layout> rhs;
+  bool result;
+
+  /// Ctor
+  TensorEqualsFunc(): result(true) { }
+
+  /// Ctor
+  TensorEqualsFunc(
+    TensorView<Element, Layout> const &lhs_,
+    TensorView<Element, Layout> const &rhs_
+  ) :
+    lhs(lhs_), rhs(rhs_), result(true) { }
+
+  /// Visits a coordinate
+  void operator()(Coord<Layout::kRank> const &coord) {
+
+    Element lhs_ = lhs.at(coord);
+    Element rhs_ = rhs.at(coord);
+
+    if (lhs_ != rhs_) {
+      result = false;
+    }
+  }
+
+  /// Returns true if equal
+  operator bool() const {
+    return result;
+  }
+};
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorRelativelyEqualsFunc {
+
+  //
+  // Data members
+  //
+
+  TensorView<Element, Layout> lhs;
+  TensorView<Element, Layout> rhs;
+  Element epsilon;
+  Element nonzero_floor;
+  bool result;
+
+  /// Ctor
+  TensorRelativelyEqualsFunc(
+    TensorView<Element, Layout> const &lhs_,
+    TensorView<Element, Layout> const &rhs_,
+    Element epsilon_,
+    Element nonzero_floor_
+  ) :
+    lhs(lhs_),
+    rhs(rhs_),
+    epsilon(epsilon_),
+    nonzero_floor(nonzero_floor_),
+    result(true) { }
+
+  /// Visits a coordinate
+  void operator()(Coord<Layout::kRank> const &coord) {
+
+    Element lhs_ = lhs.at(coord);
+    Element rhs_ = rhs.at(coord);
+
+    if (!relatively_equal(lhs_, rhs_, epsilon, nonzero_floor)) {
+      result = false;
+    }
+  }
+
+  /// Returns true if equal
+  operator bool() const {
+    return result;
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns true if two tensor views are equal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+bool TensorEquals(
+  TensorView<Element, Layout> const &lhs,
+  TensorView<Element, Layout> const &rhs) {
+
+  // Extents must be identical
+  if (lhs.extent() != rhs.extent()) {
+    return false;
+  }
+
+  detail::TensorEqualsFunc<Element, Layout> func(lhs, rhs);
+  TensorForEach(
+    lhs.extent(),
+    func
+  );
+
+  return bool(func);
+}
+
+/// Returns true if two tensor views are equal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+bool TensorEquals(
+  TensorViewPlanarComplex<Element, Layout> const &lhs,
+  TensorViewPlanarComplex<Element, Layout> const &rhs) {
+
+  // Extents must be identical
+  if (lhs.extent() != rhs.extent()) {
+    return false;
+  }
+
+  detail::TensorEqualsFunc<Element, Layout> real_func(
+    {lhs.data(), lhs.layout(), lhs.extent()},
+    {rhs.data(), rhs.layout(), rhs.extent()}
+  );
+
+  TensorForEach(
+    lhs.extent(),
+    real_func
+  );
+
+  if (!bool(real_func)) {
+    return false;
+  }
+
+  detail::TensorEqualsFunc<Element, Layout> imag_func(
+    {lhs.data() + lhs.imaginary_stride(), lhs.layout(), lhs.extent()},
+    {rhs.data() + rhs.imaginary_stride(), rhs.layout(), rhs.extent()}
+    );
+
+  TensorForEach(
+    lhs.extent(),
+    imag_func
+  );
+
+  return bool(imag_func);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns true if two tensor views are relatively equal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+bool TensorRelativelyEquals(
+  TensorView<Element, Layout> const &lhs,
+  TensorView<Element, Layout> const &rhs,
+  Element epsilon,
+  Element nonzero_floor) {
+
+  // Extents must be identical
+  if (lhs.extent() != rhs.extent()) {
+    return false;
+  }
+
+  detail::TensorRelativelyEqualsFunc<Element, Layout> func(lhs, rhs, epsilon, nonzero_floor);
+  TensorForEach(
+    lhs.extent(),
+    func
+  );
+
+  return bool(func);
+}
+
+/// Returns true if two tensor views are relatively equal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+bool TensorRelativelyEquals(
+  TensorViewPlanarComplex<Element, Layout> const &lhs,
+  TensorViewPlanarComplex<Element, Layout> const &rhs,
+  Element epsilon,
+  Element nonzero_floor) {
+
+  // Extents must be identical
+  if (lhs.extent() != rhs.extent()) {
+    return false;
+  }
+
+  detail::TensorRelativelyEqualsFunc<Element, Layout> real_func(
+    {lhs.data(), lhs.layout(), lhs.extent()},
+    {rhs.data(), rhs.layout(), rhs.extent()},
+    epsilon,
+    nonzero_floor
+  );
+
+  TensorForEach(
+    lhs.extent(),
+    real_func
+  );
+
+  if (!bool(real_func)) {
+    return false;
+  }
+
+  detail::TensorEqualsFunc<Element, Layout> imag_func(
+    {lhs.data() + lhs.imaginary_stride(), lhs.layout(), lhs.extent()},
+    {rhs.data() + rhs.imaginary_stride(), rhs.layout(), rhs.extent()},
+    epsilon,
+    nonzero_floor
+  );
+
+  TensorForEach(
+    lhs.extent(),
+    imag_func
+  );
+
+  return bool(imag_func);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns true if two tensor views are NOT equal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+bool TensorNotEquals(
+  TensorView<Element, Layout> const &lhs,
+  TensorView<Element, Layout> const &rhs) {
+
+  // Extents must be identical
+  if (lhs.extent() != rhs.extent()) {
+    return true;
+  }
+
+  detail::TensorEqualsFunc<Element, Layout> func(lhs, rhs);
+  TensorForEach(
+    lhs.extent(),
+    func
+  );
+
+  return !bool(func);
+}
+
+/// Returns true if two tensor views are equal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+bool TensorNotEquals(
+  TensorViewPlanarComplex<Element, Layout> const &lhs,
+  TensorViewPlanarComplex<Element, Layout> const &rhs) {
+
+  return !TensorEquals(lhs, rhs);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorContainsFunc {
+
+  //
+  // Data members
+  //
+
+  TensorView<Element, Layout> view;
+  Element value;
+  bool contains;
+  Coord<Layout::kRank> location;
+
+  //
+  // Methods
+  //
+
+  /// Ctor
+  TensorContainsFunc(): contains(false) { }
+
+  /// Ctor
+  TensorContainsFunc(
+    TensorView<Element, Layout> const &view_,
+    Element value_
+  ) :
+    view(view_), value(value_), contains(false) { }
+
+  /// Visits a coordinate
+  void operator()(Coord<Layout::kRank> const &coord) {
+
+    if (view.at(coord) == value) {
+      if (!contains) {
+        location = coord;
+      }
+      contains = true;
+    }
+  }
+
+  /// Returns true if equal
+  operator bool() const {
+    return contains;
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns true if a value is present in a tensor
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+bool TensorContains(
+  TensorView<Element, Layout> const & view,
+  Element value) {
+
+  detail::TensorContainsFunc<Element, Layout> func(
+    view,
+    value
+  );
+
+  TensorForEach(
+    view.extent(),
+    func
+  );
+
+  return bool(func);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns a pair containing a boolean of whether a value exists in a tensor and the location of
+/// of the first occurrence. If the value is not contained in the tensor, the second element of the
+/// pair is undefined.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+std::pair<bool, Coord<Layout::kRank> > TensorFind(
+  TensorView<Element, Layout> const & view,
+  Element value) {
+
+  detail::TensorContainsFunc<Element, Layout> func(
+    view,
+    value
+  );
+
+  TensorForEach(
+    view.extent(),
+    func
+  );
+
+  return std::make_pair(bool(func), func.location);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_compare.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_compare.hpp
new file mode 100644
index 000000000..76eac628e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_compare.hpp
@@ -0,0 +1,102 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Provides several functions for filling tensors with data.
+*/
+
+#pragma once
+
+// Standard Library includes
+#include <utility>
+#include <cstdlib>
+#include <cmath>
+
+// Mute includes
+#include "mute/tensor.hpp"
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+#include "mutlass/quaternion.h"
+#include "mutlass/array.h"
+#include "mutlass/numeric_types.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Returns true if two tensor views are equal.
+template <
+  typename TensorL,
+  typename TensorR
+>
+bool TensorEquals(
+  TensorL lhs,
+  TensorR rhs) {
+
+  // Extents must be identical
+  if (mute::size(lhs) != mute::size(rhs)) {
+    return false;
+  }
+
+  for (int64_t idx = 0; idx < mute::size(lhs); ++idx) {
+    if (lhs(idx) != rhs(idx)) {
+      return false;
+    }
+  }
+
+  return true;
+}
+
+/// Returns true if two tensor views are NOT equal.
+template <
+  typename TensorL,
+  typename TensorR
+>
+bool TensorNotEquals(
+  TensorL lhs,
+  TensorR rhs) {
+
+  return TensorEquals(lhs, rhs);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_copy.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_copy.h
new file mode 100644
index 000000000..565c32d9a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_copy.h
@@ -0,0 +1,257 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 host-side elementwise operations on TensorView.
+*/
+
+#pragma once
+
+// Standard Library includes
+#include <utility>
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "tensor_foreach.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Helper to convert between types
+template <
+  typename DstElement,
+  typename SrcElement
+>
+struct TrivialConvert {
+
+  TrivialConvert() { }
+
+  DstElement operator()(SrcElement src) const {
+    return DstElement(src);
+  }
+};
+
+/// Helper to conditionally copy between tensor views.
+template <
+  typename DstElement,
+  typename DstLayout,
+  typename SrcElement,
+  typename SrcLayout,
+  typename F
+>
+struct TensorCopyIf {
+
+  using DstTensorView = TensorView<DstElement, DstLayout>;
+  using SrcTensorView = TensorView<SrcElement, SrcLayout>;
+
+  //
+  // Data members
+  //
+
+  DstTensorView dst;
+  SrcTensorView src;
+  F convert;
+
+  //
+  // Methods
+  //
+
+  TensorCopyIf() { }
+
+  TensorCopyIf(
+    DstTensorView const &dst_,
+    SrcTensorView const &src_,
+    F const &convert_): dst(dst_), src(src_), convert(convert_) {}
+
+  /// Copies based on destination and source bounds
+  void operator()(Coord<DstLayout::kRank> const &coord) {
+    if (dst.contains(coord) && src.contains(coord)) {
+      dst.at(coord) = convert(src.at(coord));
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies elements from one tensor view into another, satisfying bounds of each tensor.
+template <
+  typename DstElement,          /// Destination tensor's element type
+  typename DstLayout,           /// Destination tensor's layout
+  typename SrcElement,          /// Source tensor's element type
+  typename SrcLayout,           /// Source tensor's layout
+  typename F                    /// Transformation functor
+>
+void TensorCopy(
+  TensorView<DstElement, DstLayout> dst,
+  TensorView<SrcElement, SrcLayout> src,
+  F const &transform) {
+
+  using CopyIf = detail::TensorCopyIf<
+    DstElement,
+    DstLayout,
+    SrcElement,
+    SrcLayout,
+    F>;
+
+  CopyIf copy_if(dst, src, transform);
+
+  TensorForEach(dst.extent(), copy_if);
+}
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies elements from a TensorRef into a TensorView. Assumes source tensor has sufficient extent
+/// to avoid out of bounds accesses.
+template <
+  typename DstElement,          /// Destination tensor's element type
+  typename DstLayout,           /// Destination tensor's layout
+  typename SrcElement,          /// Source tensor's element type
+  typename SrcLayout,           /// Source tensor's layout
+  typename F                    /// Transformation functor
+>
+void TensorCopy(
+  TensorView<DstElement, DstLayout> dst,
+  TensorRef<SrcElement, SrcLayout> src,
+  F const &transform) {
+
+  using CopyIf = detail::TensorCopyIf<
+    DstElement,
+    DstLayout,
+    SrcElement,
+    SrcLayout,
+    F>;
+
+  TensorView<SrcElement, SrcLayout> src_view(src, dst.extent());
+
+  CopyIf copy_if(dst, src_view, transform);
+
+  TensorForEach(dst.extent(), copy_if);
+}
+
+/// Copies elements from a TensorRef into a TensorView. Assumes source tensor has sufficient extent
+/// to avoid out of bounds accesses.
+template <
+  typename DstElement,          /// Destination tensor's element type
+  typename DstLayout,           /// Destination tensor's layout
+  typename SrcElement,          /// Source tensor's element type
+  typename SrcLayout,           /// Source tensor's layout
+  typename F                    /// Transformation functor
+>
+void TensorCopy(
+  TensorRef<DstElement, DstLayout> dst,
+  TensorView<SrcElement, SrcLayout> src,
+  F const &transform) {
+
+  using CopyIf = detail::TensorCopyIf<
+    DstElement,
+    DstLayout,
+    SrcElement,
+    SrcLayout,
+    F>;
+
+  TensorView<DstElement, DstLayout> dst_view(dst, src.extent());
+
+  CopyIf copy_if(dst_view, src, transform);
+
+  TensorForEach(src.extent(), copy_if);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies elements from one tensor view into another, satisfying bounds of each tensor. Succeeds
+/// if SrcElement can be converted to DstElement.
+template <
+  typename DstElement,          /// Destination tensor's element type
+  typename DstLayout,           /// Destination tensor's layout
+  typename SrcElement,          /// Source tensor's element type
+  typename SrcLayout            /// Source tensor's layout
+>
+void TensorCopy(
+  TensorView<DstElement, DstLayout> dst,
+  TensorView<SrcElement, SrcLayout> src) {
+
+  detail::TrivialConvert<DstElement, SrcElement> convert;
+
+  TensorCopy(dst, src, convert);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies elements from one tensor view into another, satisfying bounds of each tensor. Succeeds
+/// if SrcElement can be converted to DstElement.
+template <
+  typename DstElement,          /// Destination tensor's element type
+  typename DstLayout,           /// Destination tensor's layout
+  typename SrcElement,          /// Source tensor's element type
+  typename SrcLayout,           /// Source tensor's layout
+  typename F                    /// Transformation functor
+>
+void TensorCopy(
+  TensorView<DstElement, DstLayout> dst,
+  TensorRef<SrcElement, SrcLayout> src) {
+
+  detail::TrivialConvert<DstElement, SrcElement> convert;
+
+  TensorCopy(dst, src, convert);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies elements from one tensor view into another, satisfying bounds of each tensor. Succeeds
+/// if SrcElement can be converted to DstElement.
+template <
+  typename DstElement,          /// Destination tensor's element type
+  typename DstLayout,           /// Destination tensor's layout
+  typename SrcElement,          /// Source tensor's element type
+  typename SrcLayout            /// Source tensor's layout
+>
+void TensorCopy(
+  TensorRef<DstElement, DstLayout> dst,
+  TensorView<SrcElement, SrcLayout> src) {
+
+  detail::TrivialConvert<DstElement, SrcElement> convert;
+
+  TensorCopy(dst, src, convert);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_elementwise.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_elementwise.h
new file mode 100644
index 000000000..05fda4e16
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_elementwise.h
@@ -0,0 +1,342 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 host-side elementwise operations on TensorView.
+*/
+
+#pragma once
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/functional.h"
+
+#include "tensor_foreach.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper to apply a binary operator in place
+template <
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementD,
+  typename LayoutD,
+  typename BinaryFunc>
+struct TensorFuncBinaryOp {
+
+  //
+  // Data members
+  //
+
+  /// View of left-hand-side tensor
+  TensorView<ElementD, LayoutD> view_d;
+  TensorRef<ElementA, LayoutA> view_a;
+  TensorRef<ElementB, LayoutB> view_b;
+  BinaryFunc func;
+
+  //
+  // Methods
+  //
+
+  /// Constructor
+  TensorFuncBinaryOp() { }
+
+  /// Constructor
+  TensorFuncBinaryOp(
+    TensorView<ElementD, LayoutD> const & view_d_,
+    TensorRef<ElementA, LayoutA> const & view_a_,
+    TensorRef<ElementB, LayoutB> const & view_b_,
+    BinaryFunc func = BinaryFunc()
+  ):
+    view_d(view_d_), view_a(view_a_), view_b(view_b_), func(func) { }
+
+  /// Equality check
+  void operator()(Coord<LayoutD::kRank> const &coord) const {
+    view_d.at(coord) = func(
+      ElementD(view_a.at(coord)),
+      ElementD(view_b.at(coord))
+    );
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Adds two tensors and stores in the destination tensor: d = a + b
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB
+>
+void TensorAdd(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a,       ///< A tensor reference
+  TensorRef<ElementB, LayoutB> b        ///< B tensor reference
+) {
+
+  detail::TensorFuncBinaryOp<
+    ElementD,
+    LayoutD,
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    mutlass::plus<ElementD>
+  > func(d, a, b);
+
+  TensorForEach(
+    d.extent(),
+    func);
+}
+
+/// Adds a tensor in place: d = d .+ a
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA
+>
+void TensorAdd(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a        ///< A tensor reference
+) {
+  TensorAdd(d, d, a);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Subtracts two tensors and stores in the destination tensor: d = a - b
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB
+>
+void TensorSub(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a,       ///< A tensor reference
+  TensorRef<ElementB, LayoutB> b        ///< B tensor reference
+  ) {
+
+  detail::TensorFuncBinaryOp<
+    ElementD,
+    LayoutD,
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    mutlass::minus<ElementD>
+  > func(d, a, b);
+
+  TensorForEach(
+    d.extent(),
+    func);
+}
+
+/// Subtracts two tensors in place: d = d .- a
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB
+>
+void TensorSub(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a        ///< A tensor reference
+  ) {
+
+  TensorSub(d, d, a);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Multiplies two tensors and stores in the destination tensor: d = a .* b
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB
+>
+void TensorMul(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a,       ///< A tensor reference
+  TensorRef<ElementB, LayoutB> b        ///< B tensor reference
+) {
+
+  detail::TensorFuncBinaryOp<
+    ElementD,
+    LayoutD,
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    mutlass::multiplies<ElementD>
+  > func(d, a, b);
+
+  TensorForEach(
+    d.extent(),
+    func);
+}
+
+/// Multiplies tensors in place: d = d .* a
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA
+>
+void TensorMul(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a        ///< A tensor reference
+) {
+  TensorMul(d, d, a);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Divides two tensors and stores in the destination tensor: d = a ./ b
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB
+>
+void TensorDiv(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a,       ///< A tensor reference
+  TensorRef<ElementB, LayoutB> b        ///< B tensor reference
+) {
+
+  detail::TensorFuncBinaryOp<
+    ElementD,
+    LayoutD,
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    mutlass::divides<ElementD>
+  > func(d, a, b);
+
+  TensorForEach(
+    d.extent(),
+    func);
+}
+
+/// Divides tensors in place: d = d ./ a
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA
+>
+void TensorDiv(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a        ///< A tensor reference
+) {
+  TensorDiv(d, d, a);
+}
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Divides two tensors and stores in the destination tensor: d = a ./ b
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA,
+  typename ElementB,
+  typename LayoutB
+>
+void TensorModulus(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a,       ///< A tensor reference
+  TensorRef<ElementB, LayoutB> b        ///< B tensor reference
+) {
+
+  detail::TensorFuncBinaryOp<
+    ElementD,
+    LayoutD,
+    ElementA,
+    LayoutA,
+    ElementB,
+    LayoutB,
+    mutlass::divides<ElementD>
+  > func(d, a, b);
+
+  TensorForEach(
+    d.extent(),
+    func);
+}
+
+/// Divides tensors in place: d = d ./ a
+template <
+  typename ElementD,
+  typename LayoutD,
+  typename ElementA,
+  typename LayoutA
+>
+void TensorModulus(
+  TensorView<ElementD, LayoutD> d,      ///< destination tensor view
+  TensorRef<ElementA, LayoutA> a        ///< A tensor reference
+) {
+  TensorDiv(d, d, a);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_fill.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_fill.h
new file mode 100644
index 000000000..f4fda5279
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_fill.h
@@ -0,0 +1,1719 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Provides several functions for filling tensors with data.
+*/
+
+#pragma once
+
+// Standard Library includes
+#include <utility>
+#include <cstdlib>
+#include <cmath>
+#include <random>
+#include <stdexcept>
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+#include "mutlass/quaternion.h"
+#include "mutlass/array.h"
+#include "mutlass/numeric_types.h"
+#include "mutlass/subbyte_reference.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/tensor_view_planar_complex.h"
+#include "mutlass/blas3.h"
+
+#include "mutlass/util/distribution.h"
+#include "tensor_foreach.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  Element value;
+
+  //
+  // Methods
+  //
+
+  TensorFillFunc(
+    TensorView const &view_ = TensorView(),
+    Element value_ = Element(0)
+  ): view(view_), value(value_) { }
+
+  void operator()(Coord<Layout::kRank> const & coord) const {
+    view.at(coord) = value;
+  }
+};
+
+/// Returns a pair of values of the Gaussian distribution generated by the Box Muller method
+struct BoxMullerFunc {
+
+  BoxMullerFunc() {}
+
+  void operator()(
+    double* rnd,                     ///< Size-2 vector to be filled with random values
+    double  mean = 0,                ///< Mean of the Gaussian distribution
+    double  stddev = 1,              ///< Standard deviation of the Gaussian distribution
+    double  pi = std::acos(-1)) const {
+
+    double u1 = double(std::rand()) / double(RAND_MAX);
+    double u2 = double(std::rand()) / double(RAND_MAX);
+    rnd[0] = std::sqrt(-2 * std::log(u1)) * std::cos(2 * pi * u2);
+    rnd[1] = std::sqrt(-2 * std::log(u1)) * std::sin(2 * pi * u2);
+    rnd[0] = mean + stddev * rnd[0];
+    rnd[1] = mean + stddev * rnd[1];
+  }
+};
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with a uniform value
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFill(
+  TensorView<Element, Layout> dst,    ///< destination tensor
+  Element val = Element(0)) {               ///< value to uniformly fill it with
+
+  detail::TensorFillFunc<Element, Layout> func(dst, val);
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+/// Fills a tensor with a uniform value
+template <
+  typename Element,                                                   ///< Element type
+  typename Layout>                                                    ///< Layout function
+void TensorFill(
+  TensorViewPlanarComplex<Element, Layout> dst,                       ///< destination tensor
+  mutlass::complex<Element> val = mutlass::complex<Element>(0)) {     ///< value to uniformly fill it with
+
+  TensorFill(dst.view_real(), val.real());
+  TensorFill(dst.view_imag(), val.imag());
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <typename Element>
+struct RandomGaussianFunc {
+
+  uint64_t seed;
+  double mean;
+  double stddev;
+  int int_scale;
+  double pi;
+  double pnz;
+  bool exclude_zero;
+
+  //
+  // Methods
+  //
+  RandomGaussianFunc(
+    uint64_t seed_ = 0,
+    double mean_ = 0,
+    double stddev_ = 1,
+    int int_scale_ = -1,
+    double pnz_ = 1.0,
+    bool exclude_zero_ = false
+  ):
+    seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)), pnz(pnz_), exclude_zero(exclude_zero_) {
+      std::srand((unsigned)seed);
+  }
+
+  /// Compute random value and update RNG state
+  Element operator()() const {
+
+    // Box-Muller transform to generate random numbers with Normal distribution
+    double u1 = double(std::rand()) / double(RAND_MAX);
+    double u2 = double(std::rand()) / double(RAND_MAX);
+
+    // Compute Gaussian random value
+    double rnd = std::sqrt(-2 * std::log(u1)) * std::cos(2 * pi * u2);
+    rnd = mean + stddev * rnd;
+
+    // Scale and convert final result
+    Element result;
+
+    // Sample from the Bernoulli distribution, and use the result to sample from the Gaussian
+    std::random_device rnd_device;
+    std::mt19937 bernoulli_rnd(rnd_device());
+    std::bernoulli_distribution bernoulli_dist(pnz);
+    bool bernoulli_result = bernoulli_dist(bernoulli_rnd);
+
+    // Sample from the Gaussian distribution for a nonzero element
+    if (bernoulli_result) {
+      if (int_scale >= 0) {
+        rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
+        result = static_cast<Element>(rnd);
+      }
+      else {
+        result = static_cast<Element>(rnd);
+      }
+    }
+    else {
+      result = static_cast<Element>(0);
+    }
+
+    // Note that exclude_zero = true will disable the bernoulli_result above by unsetting zeros
+    if (exclude_zero && result == Element(0)) {
+      if (rnd > 0) {
+        rnd += 1;
+      } else {
+        rnd -= 1;
+      }
+      result = Element(rnd);
+    }
+
+    return result;
+  }
+};
+
+/// Partial specialization for initializing a complex value.
+template <typename Element>
+struct RandomGaussianFunc<complex<Element> > {
+
+  uint64_t seed;
+  double mean;
+  double stddev;
+  int int_scale;
+  double pi;
+  double pnz;
+  bool exclude_zero;
+
+  //
+  // Methods
+  //
+  RandomGaussianFunc(
+    uint64_t seed_ = 0,
+    double mean_ = 0,
+    double stddev_ = 1,
+    int int_scale_ = -1,
+    double pnz_ = 1.0,
+    bool exclude_zero_ = false
+  ):
+    seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)), pnz(pnz_), exclude_zero(exclude_zero_) {
+      std::srand((unsigned)seed);
+  }
+
+  /// Compute random value and update RNG state
+  complex<Element> operator()() const {
+
+    Element reals[2];
+
+    double rnd[2];
+    detail::BoxMullerFunc func;
+    func(rnd, mean, stddev, pi);
+
+    // Sample from the Bernoulli distribution, and use the result to sample from the Gaussian
+    std::random_device rnd_device;
+    std::mt19937 bernoulli_rnd(rnd_device());
+    std::bernoulli_distribution bernoulli_dist(pnz);
+    bool bernoulli_result = bernoulli_dist(bernoulli_rnd);
+
+    // Sample from the Gaussian distribution for a nonzero element
+    if (bernoulli_result) {
+      if (int_scale >= 0) {
+        rnd[0] = double(int(rnd[0] * double(1 << int_scale)));
+        rnd[1] = double(int(rnd[1] * double(1 << int_scale)));
+        reals[0] = from_real<Element>(rnd[0] / double(1 << int_scale));
+        reals[1] = from_real<Element>(rnd[1] / double(1 << int_scale));
+      }
+      else {
+        reals[0] = from_real<Element>(rnd[0]);
+        reals[1] = from_real<Element>(rnd[1]);
+      }
+    }
+    else {
+      reals[0] = from_real<Element>(0);
+      reals[1] = from_real<Element>(0);
+    }
+
+    // Note that this will invalidate the above else statement because it unsets zero elements
+    if (exclude_zero &&
+        reals[0] == from_real<Element>(0.0) &&
+        reals[1] == from_real<Element>(0.0)) {
+
+      if (rnd[0] > 0.0) {
+        rnd[0] += 1.0;
+      } else {
+        rnd[0] -= 1.0;
+      }
+      reals[0] = from_real<Element>(rnd[0]);
+    }
+
+    return complex<Element>(reals[0], reals[1]);
+  }
+};
+
+/// Partial specialization for initializing a complex value.
+template <typename Element>
+struct RandomGaussianFunc<Quaternion<Element> > {
+
+  uint64_t seed;
+  double mean;
+  double stddev;
+  int int_scale;
+  double pi;
+  double pnz;
+  bool exclude_zero;
+
+  //
+  // Methods
+  //
+  RandomGaussianFunc(
+    uint64_t seed_ = 0,
+    double mean_ = 0,
+    double stddev_ = 1,
+    int int_scale_ = -1,
+    double pnz_ = 1.0,
+    bool exclude_zero_ = false
+  ):
+    seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)), pnz(pnz_), exclude_zero(exclude_zero_) {
+      std::srand((unsigned)seed);
+  }
+
+  /// Compute random value and update RNG state
+  Quaternion<Element> operator()() const {
+
+    Element reals[4];
+
+    double rnd1[2];
+    double rnd2[2];
+    detail::BoxMullerFunc func;
+    func(rnd1, mean, stddev, pi);
+    func(rnd2, mean, stddev, pi);
+
+    // Sample from the Bernoulli distribution, and use the result to sample from the Gaussian
+    std::random_device rnd_device;
+    std::mt19937 bernoulli_rnd(rnd_device());
+    std::bernoulli_distribution bernoulli_dist(pnz);
+    bool bernoulli_result = bernoulli_dist(bernoulli_rnd);
+
+    // Sample from the Gaussian distribution for a nonzero element
+    if (bernoulli_result) {
+      if (int_scale >= 0) {
+        rnd1[0] = double(int(rnd1[0] * double(1 << int_scale)));
+        rnd1[1] = double(int(rnd1[1] * double(1 << int_scale)));
+        rnd2[0] = double(int(rnd2[0] * double(1 << int_scale)));
+        rnd2[1] = double(int(rnd2[1] * double(1 << int_scale)));
+
+        reals[0] = from_real<Element>(rnd1[0] / double(1 << int_scale));
+        reals[1] = from_real<Element>(rnd1[1] / double(1 << int_scale));
+        reals[2] = from_real<Element>(rnd2[0] / double(1 << int_scale));
+        reals[3] = from_real<Element>(rnd2[1] / double(1 << int_scale));
+      }
+      else {
+        reals[0] = from_real<Element>(rnd1[0]);
+        reals[1] = from_real<Element>(rnd1[1]);
+        reals[2] = from_real<Element>(rnd2[0]);
+        reals[3] = from_real<Element>(rnd2[1]);
+      }
+    }
+    else {
+      reals[0] = from_real<Element>(0);
+      reals[1] = from_real<Element>(0);
+      reals[2] = from_real<Element>(0);
+      reals[3] = from_real<Element>(0);
+    }
+
+    // Note that this will invalidate the above else statement because it unsets zero elements
+    if (exclude_zero &&
+        reals[0] == from_real<Element>(0) &&
+        reals[1] == from_real<Element>(0) &&
+        reals[2] == from_real<Element>(0) &&
+        reals[3] == from_real<Element>(0)) {
+
+      if (rnd1[0] > 0.0) {
+        rnd1[0] += 1.0;
+      } else {
+        rnd1[0] -= 1.0;
+      }
+      reals[0] = from_real<Element>(rnd1[0]);
+    }
+
+    return Quaternion<Element>(reals[0], reals[1], reals[2], reals[3]);
+  }
+};
+
+/// Computes a random Gaussian distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillGaussianFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  RandomGaussianFunc<Element> func;
+
+  //
+  // Methods
+  //
+
+  /// Construction of Gaussian RNG functor.
+  TensorFillGaussianFunc(
+    TensorView view_ = TensorView(),
+    RandomGaussianFunc<Element> func_ = RandomGaussianFunc<Element>()
+  ):
+    view(view_), func(func_) {
+
+  }
+
+  /// Compute random value and update RNG state
+  void operator()(Coord<Layout::kRank> const &coord) const {
+    view.at(coord) = func();
+  }
+};
+
+/// Computes a random Gaussian distribution for a rank-2 tensor
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillSymmetricGaussianFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  RandomGaussianFunc<Element> func;
+  mutlass::FillMode fill_mode;
+
+  //
+  // Methods
+  //
+
+  /// Construction of Gaussian RNG functor.
+  TensorFillSymmetricGaussianFunc(
+    TensorView view_ = TensorView(),
+    RandomGaussianFunc<Element> func_ = RandomGaussianFunc<Element>(),
+    mutlass::FillMode fill_mode_ = mutlass::FillMode::kInvalid
+  ):
+    view(view_), func(func_), fill_mode(fill_mode_) {
+
+  }
+
+  /// Compute random value and update RNG state
+  void operator()(Coord<Layout::kRank> const &coord) const {
+    // Fill half of matrix based on FillMode
+    if (Layout::kRank == 2 &&
+        fill_mode == mutlass::FillMode::kLower &&
+        coord[0] >= coord[1]) {
+      view.at(coord) = func();
+    } else if (Layout::kRank == 2 &&
+        fill_mode == mutlass::FillMode::kUpper &&
+        coord[0] <= coord[1]) {
+      view.at(coord) = func();
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a Gaussian distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandomGaussian(
+  TensorView<Element, Layout> dst,        ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  double mean = 0,                        ///< Gaussian distribution's mean
+  double stddev = 1,                      ///< Gaussian distribution's standard deviation
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+  double pnz = 1.0,                     ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+  bool exclude_zero = false) {            ///< Exclude zeros from tensor init.
+
+  detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits, pnz, exclude_zero);
+
+  detail::TensorFillGaussianFunc<Element, Layout> func(
+    dst,
+    random_func
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+/// Fills a tensor with random values with a Gaussian distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandomGaussian(
+  TensorViewPlanarComplex<Element, Layout> dst,         ///< destination tensor
+  uint64_t seed,                                        ///< seed for RNG
+  double mean = 0,                                      ///< Gaussian distribution's mean
+  double stddev = 1,                                    ///< Gaussian distribution's standard deviation
+  int bits = -1,                                        ///< If non-negative, specifies number of fractional bits that
+  double pnz = 1.0,                                   ///  are not truncated to zero. Permits reducing precision of
+                                                        ///  data.
+  bool exclude_zero = false) {                          ///< Exclude zeros from tensor init.
+
+  TensorFillRandomGaussian(dst.view_real(), seed, mean, stddev, bits, pnz);
+  TensorFillRandomGaussian(dst.view_imag(), ~seed, mean, stddev, bits, pnz);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+/// Fills the upper or lower part of a symmetric rank-2 tensor with random values of a Gaussian distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillSymmetricRandomGaussian(
+  TensorView<Element, Layout> dst,        ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  mutlass::FillMode fill_mode,            ///< FillMode for symmetric matrices
+  double mean = 0,                        ///< Gaussian distribution's mean
+  double stddev = 1,                      ///< Gaussian distribution's standard deviation
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+  double pnz = 1.0) {                   ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+
+  detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits, pnz);
+
+  detail::TensorFillSymmetricGaussianFunc<Element, Layout> func(
+    dst,
+    random_func,
+    fill_mode
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values of a Gaussian distribution.
+template <
+  typename Element                        ///< Element type
+>
+void BlockFillRandomGaussian(
+  Element *ptr,                           ///< destination buffer
+  size_t capacity,                        ///< number of elements
+  uint64_t seed,                          ///< seed for RNG
+  double mean = 0,                        ///< Gaussian distribution's mean
+  double stddev = 1,                      ///< Gaussian distribution's standard deviation
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+  double pnz = 1.0) {                   ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+
+
+  detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits, pnz);
+
+  for (size_t i = 0; i < capacity; ++i) {
+    ReferenceFactory<Element>::get(ptr, i) = random_func();
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <typename Element>
+struct RandomUniformFunc {
+
+  using Real = typename RealType<Element>::Type;
+
+  uint64_t seed;
+  double range;
+  double min;
+  int int_scale;
+
+  double pnan;
+private:
+  using engine_type = std::mt19937;
+public:
+  engine_type bernoulli_rnd;
+  std::bernoulli_distribution bernoulli_dist;
+
+  bool exclude_zero;
+
+  RandomUniformFunc(
+    uint64_t seed_ = 0,
+    double max = 1,
+    double min_ = 0,
+    int int_scale_ = -1,
+    double pnan_ = 0,
+    bool exclude_zero_ = false
+  ):
+    seed(seed_), range(max - min_), min(min_), int_scale(int_scale_), pnan(pnan_)
+    , bernoulli_rnd{static_cast<engine_type::result_type>(seed_)}
+    , bernoulli_dist(pnan_)
+    , exclude_zero(exclude_zero_)
+    {
+      std::srand((unsigned)seed);
+
+      // Handle cases where min = 0 or max = 0 for excluding zeros
+      if (exclude_zero) {
+        min = (min == 0.0) ? min + 1: min;
+        range = (max == 0.0) ? range - 1: range;
+      }
+  }
+
+
+  /// Compute random value and update RNG state
+  Element operator()() {
+
+    // Sample from NaN distribution.
+    if constexpr (std::numeric_limits<Element>::has_quiet_NaN) {
+      if (pnan > 0 && bernoulli_dist(bernoulli_rnd)) {
+        return Element(NAN);
+      }
+    }
+
+    double rnd = double(std::rand()) / double(RAND_MAX);
+
+    rnd = min + range * rnd;
+
+    // Random values are cast to integer after scaling by a power of two to facilitate error
+    // testing
+    Element result;
+    if (int_scale >= 0) {
+      rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
+      result = static_cast<Element>(Real(rnd));
+    }
+    else {
+      result = static_cast<Element>(Real(rnd));
+    }
+
+    if (exclude_zero && result == Element(0)) {
+      if (rnd > 0.0) {
+        rnd = std::min(min + range, rnd + 1.0);
+      } else {
+        rnd = std::max(min, rnd - 1.0);
+      }
+      result = static_cast<Element>(Real(rnd));
+    }
+
+    return result;
+  }
+};
+
+/// Partial specialization for initializing a complex value.
+template <typename Element>
+struct RandomUniformFunc<complex<Element> > {
+
+  using Real = typename RealType<Element>::Type;
+
+  uint64_t seed;
+  double range;
+  double min;
+  int int_scale;
+
+  double pnan;
+private:
+  using engine_type = std::mt19937;
+public:
+  engine_type bernoulli_rnd;
+  std::bernoulli_distribution bernoulli_dist;
+
+  bool exclude_zero;
+
+  //
+  // Methods
+  //
+
+  RandomUniformFunc(
+    uint64_t seed_ = 0,
+    double max = 1,
+    double min_ = 0,
+    int int_scale_ = -1,
+    double pnan_ = 0,
+    bool exclude_zero_ = false
+  ):
+    seed(seed_), range(max - min_), min(min_), int_scale(int_scale_), pnan(pnan_)
+    , bernoulli_rnd{static_cast<engine_type::result_type>(seed_)}
+    , bernoulli_dist(pnan_)
+    , exclude_zero(exclude_zero_) {
+      std::srand((unsigned)seed);
+
+      // Handle cases where min = 0 or max = 0 for excluding zeros
+      if (exclude_zero) {
+        min = (min == 0.0) ? min + 1: min;
+        range = (max == 0.0) ? range - 1: range;
+      }
+  }
+
+
+  /// Compute random value and update RNG state
+  complex<Element> operator()() {
+
+    // Sample from NaN distribution.
+    if constexpr (std::numeric_limits<Element>::has_quiet_NaN) {
+      if (pnan > 0 && bernoulli_dist(bernoulli_rnd)) {
+        return Element(NAN);
+      }
+    }
+
+    Element reals[2];
+
+    for (int i = 0; i < 2; ++i) {
+      double rnd = double(std::rand()) / double(RAND_MAX);
+
+      rnd = min + range * rnd;
+
+      // Random values are cast to integer after scaling by a power of two to facilitate error
+      // testing
+
+      if (int_scale >= 0) {
+        rnd = double(int(rnd * double(1 << int_scale)));
+        reals[i] = from_real<Element>(Real(rnd / double(1 << int_scale)));
+      }
+      else {
+        reals[i] = from_real<Element>(Real(rnd));
+      }
+
+      if (exclude_zero &&
+          i == 0 &&
+          reals[0] == from_real<Element>(0.0)) {
+
+        if (rnd > 0.0) {
+          rnd = std::min(min + range, rnd + 1.0);
+        } else {
+          rnd = std::max(min, rnd - 1.0);
+        }
+        reals[0] = from_real<Element>(Real(rnd));
+      }
+
+    }
+
+    return complex<Element>(reals[0], reals[1]);
+  }
+};
+
+/// Partial specialization for initializing a Quaternion value.
+template <typename Element>
+struct RandomUniformFunc<Quaternion<Element> > {
+
+  using Real = typename RealType<Element>::Type;
+
+  uint64_t seed;
+  double range;
+  double min;
+  int int_scale;
+
+  double pnan;
+private:
+  using engine_type = std::mt19937;
+public:
+  engine_type bernoulli_rnd;
+  std::bernoulli_distribution bernoulli_dist;
+
+  //
+  // Methods
+  //
+
+  RandomUniformFunc(
+    uint64_t seed_ = 0,
+    double max = 1,
+    double min_ = 0,
+    int int_scale_ = -1,
+    double pnan_ = 0
+  ):
+    seed(seed_), range(max - min_), min(min_), int_scale(int_scale_), pnan(pnan_),
+    bernoulli_rnd{static_cast<engine_type::result_type>(seed_)},
+    bernoulli_dist(pnan_)
+  {
+    std::srand((unsigned)seed);
+  }
+
+
+  /// Compute random value and update RNG state
+  Quaternion<Element> operator()() {
+
+    // Sample from NaN distribution.
+    if constexpr (std::numeric_limits<Element>::has_quiet_NaN) {
+      if (pnan > 0 && bernoulli_dist(bernoulli_rnd)) {
+        return Element(NAN);
+      }
+    }
+
+    Element reals[4];
+
+    for (int i = 0; i < 4; ++i) {
+      double rnd = double(std::rand()) / double(RAND_MAX);
+
+      rnd = min + range * rnd;
+
+      // Random values are cast to integer after scaling by a power of two to facilitate error
+      // testing
+
+      if (int_scale >= 0) {
+        rnd = double(int(rnd * double(1 << int_scale)));
+        reals[i] = from_real<Element>(Real(rnd / double(1 << int_scale)));
+      }
+      else {
+        reals[i] = from_real<Element>(Real(rnd));
+      }
+    }
+
+    return make_Quaternion(reals[0], reals[1], reals[2], reals[3]);
+  }
+};
+
+/// Computes a random uniform distribution
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillRandomUniformFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  RandomUniformFunc<Element> func;
+
+  //
+  // Methods
+  //
+
+  /// Construction of uniform RNG functor.
+  TensorFillRandomUniformFunc(
+    TensorView view_ = TensorView(),
+    RandomUniformFunc<Element> func_ = RandomUniformFunc<Element>()
+  ):
+    view(view_), func(func_) {
+
+  }
+
+  /// Compute random value and update RNG state
+  void operator()(Coord<Layout::kRank> const &coord) {
+
+    view.at(coord) = func();
+  }
+};
+
+/// Fills the upper or lower part of a symmetric rank-2 tensor with random values of a uniform distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillSymmetricRandomUniformFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  RandomUniformFunc<Element> func;
+  mutlass::FillMode fill_mode;
+
+  //
+  // Methods
+  //
+
+  /// Construction of uniform RNG functor.
+  TensorFillSymmetricRandomUniformFunc(
+    TensorView view_ = TensorView(),
+    RandomUniformFunc<Element> func_ = RandomUniformFunc<Element>(),
+    mutlass::FillMode fill_mode_ = mutlass::FillMode::kInvalid
+  ):
+    view(view_), func(func_), fill_mode(fill_mode_) {
+
+  }
+
+  /// Compute random value and update RNG state
+  void operator()(Coord<Layout::kRank> const &coord) {
+    // Fill half of matrix based on FillMode
+    if (Layout::kRank == 2 &&
+        fill_mode == mutlass::FillMode::kLower &&
+        coord[0] >= coord[1]) {
+      view.at(coord) = func();
+    } else if (Layout::kRank == 2 &&
+        fill_mode == mutlass::FillMode::kUpper &&
+        coord[0] <= coord[1]) {
+      view.at(coord) = func();
+    }
+  }
+};
+
+/// Computes a random Uniform distribution and pads diagonal with zeros
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillPadDiagonalRandomUniformFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  RandomUniformFunc<Element> func;
+  mutlass::FillMode fill_mode;
+  int alignment;
+
+  //
+  // Methods
+  //
+
+  /// Construction of uniform RNG functor.
+  TensorFillPadDiagonalRandomUniformFunc(
+    TensorView view_ = TensorView(),
+    RandomUniformFunc<Element> func_ = RandomUniformFunc<Element>(),
+    mutlass::FillMode fill_mode_ = mutlass::FillMode::kInvalid,
+    int alignment_ = 1
+  ):
+    view(view_), func(func_), fill_mode(fill_mode_), alignment(alignment_) {
+
+  }
+
+  /// Compute random value and update RNG state
+  void operator()(Coord<Layout::kRank> const &coord) {
+    // Fill half of matrix based on FillMode
+    if (Layout::kRank == 2 &&
+        (fill_mode == mutlass::FillMode::kLower) &&
+        (coord[0] >= coord[1]) ||
+        ((coord[1] - coord[0]) >= alignment)) {
+      view.at(coord) = func();
+    } else if (Layout::kRank == 2 &&
+        fill_mode == mutlass::FillMode::kUpper &&
+        (coord[0] <= coord[1]) ||
+        ((coord[0] - coord[1]) >= alignment)) {
+      view.at(coord) = func();
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values of a uniform random distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandomUniform(
+  TensorView<Element, Layout> dst,        ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  double max = 1,                         ///< upper bound of distribution
+  double min = 0,                         ///< lower bound for distribution
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+  double pnan = 0,                        ///< Percentage of NaN elements.
+  bool exclude_zero = false) {            ///< Exclude zero from tensor init
+  detail::RandomUniformFunc<Element> random_func(seed, max, min, bits, pnan, exclude_zero);
+
+  detail::TensorFillRandomUniformFunc<Element, Layout> func(
+    dst,
+    random_func
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+/// Fills a tensor with random values of a uniform random distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandomUniform(
+  TensorViewPlanarComplex<Element, Layout> dst,        ///< destination tensor
+  uint64_t seed,                                       ///< seed for RNG
+  double max = 1,                                      ///< upper bound of distribution
+  double min = 0,                                      ///< lower bound for distribution
+  int bits = -1,                                       ///< If non-negative, specifies number of fractional bits that
+                                                       ///  are not truncated to zero. Permits reducing precision of
+                                                       ///  data.
+  double pnan = 0,                                     ///< Percentage of NaN elements.
+  bool exclude_zero = false) {                         ///< Exclude zero from tensor init
+
+  TensorFillRandomUniform(dst.view_real(), seed, max, min, bits, pnan, exclude_zero);
+  TensorFillRandomUniform(dst.view_imag(), ~seed, max, min, bits, pnan, exclude_zero);
+}
+
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandomUniform(
+  TensorView<Quaternion<Element>, Layout> dst,        ///< destination tensor
+  uint64_t seed,                                      ///< seed for RNG
+  double max = 1,                                     ///< upper bound of distribution
+  double min = 0,                                     ///< lower bound for distribution
+  int bits = -1) {                                    ///< If non-negative, specifies number of fractional bits that
+                                                      ///  are not truncated to zero. Permits reducing precision of
+                                                      ///  data.
+  detail::RandomUniformFunc<Quaternion<Element>> random_func(seed, max, min, bits);
+
+  detail::TensorFillRandomUniformFunc<Quaternion<Element>, Layout> func(
+    dst,
+    random_func
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillSymmetricRandomUniform(
+  TensorView<Element, Layout> dst,        ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  mutlass::FillMode fill_mode,            ///< FillMode for symmetric matrices
+  double max = 1,                         ///< upper bound of distribution
+  double min = 0,                         ///< lower bound for distribution
+  int bits = -1) {                        ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+
+  detail::RandomUniformFunc<Element> random_func(seed, max, min, bits);
+
+  detail::TensorFillSymmetricRandomUniformFunc<Element, Layout> func(
+    dst,
+    random_func,
+    fill_mode
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+/// Fills a tensor with random values with a uniform random distribution pads zeros along diagonal
+template <
+  typename Element,                       ///< Element type
+  typename Layout>                        ///< Layout function
+void TensorFillPadDiagonalRandomUniform(
+  TensorView<Element, Layout> dst,        ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  mutlass::FillMode fill_mode,            ///< FillMode for symmetric matrices
+  double max = 1,                         ///< upper bound of distribution
+  double min = 0,                         ///< lower bound for distribution
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+  int alignment = 1
+) {
+
+  detail::RandomUniformFunc<Element> random_func(seed, max, min, bits);
+
+  detail::TensorFillPadDiagonalRandomUniformFunc<Element, Layout> func(
+    dst,
+    random_func,
+    fill_mode,
+    alignment
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with a uniform value
+template <
+  typename Element                        ///< Element type
+>
+void BlockFill(
+  Element *ptr,
+  size_t capacity,
+  Element val
+  ) {
+  for (size_t i = 0; i < capacity; ++i) {
+    ReferenceFactory<Element>::get(ptr, i) = val;
+  }
+}
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <
+  typename Element                        ///< Element type
+>
+void BlockFillRandomUniform(
+  Element *ptr,
+  size_t capacity,
+  uint64_t seed,                          ///< seed for RNG
+  double max = 1,                         ///< upper bound of distribution
+  double min = 0,                         ///< lower bound for distribution
+  int bits = -1,                          ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+  double pnan = 0) {                      ///< Percentage of NaN elements.
+  detail::RandomUniformFunc<Element> random_func(seed, max, min, bits, pnan);
+
+  for (size_t i = 0; i < capacity; ++i) {
+    ReferenceFactory<Element>::get(ptr, i) = random_func();
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillDiagonalFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  Element diag;
+  Element other;
+
+  //
+  // Methods
+  //
+
+  TensorFillDiagonalFunc(
+    TensorView const &view_ = TensorView(),
+    Element diag_ = Element(1),
+    Element other_ = Element(0)
+  ):
+    view(view_), diag(diag_), other(other_) { }
+
+  void operator()(Coord<Layout::kRank> const & coord) const {
+    bool is_diag = true;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < Layout::kRank; ++i) {
+      if (coord[i] != coord[i - 1]) {
+        is_diag = false;
+        break;
+      }
+    }
+
+    view.at(coord) = (is_diag ? diag : other);
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor everywhere with a unique value for its diagonal.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillDiagonal(
+  TensorView<Element, Layout> dst,        ///< destination tensor
+  Element diag = Element(1),              ///< value to write in the diagonal
+  Element other = Element(0)) {           ///< value to write off the diagonal
+
+  detail::TensorFillDiagonalFunc<Element, Layout> func(
+    dst,
+    diag,
+    other
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Helper to fill a tensor's diagonal with 1 and 0 everywhere else.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillIdentity(
+  TensorView<Element, Layout> dst) {               ///< destination tensor
+
+  TensorFillDiagonal(dst, Element(1), Element(0));
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Writes a uniform value to the diagonal of a tensor without modifying off-diagonal elements.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorUpdateDiagonal(
+  TensorView<Element, Layout> dst,                 ///< destination tensor
+  Element val = Element(1)) {
+
+  typename Layout::Index extent = dst.extent().min();
+
+  for (typename Layout::Index i = 0; i < extent; ++i) {
+    Coord<Layout::kRank> coord(i);
+    dst.at(coord) = val;
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorUpdateOffDiagonalFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  Element other;
+
+  //
+  // Methods
+  //
+
+  TensorUpdateOffDiagonalFunc(
+    TensorView const &view_ = TensorView(),
+    Element other_ = Element(0)
+  ):
+    view(view_), other(other_) { }
+
+  void operator()(Coord<Layout::kRank> const & coord) const {
+    bool is_diag = true;
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 1; i < Layout::kRank; ++i) {
+      if (coord[i] != coord[i - 1]) {
+        is_diag = false;
+        break;
+      }
+    }
+
+    if (!is_diag) {
+      view.at(coord) = other;
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Writes a uniform value to all elements in the tensor without modifying diagonal elements.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorUpdateOffDiagonal(
+  TensorView<Element, Layout> dst,      ///< destination tensor
+  Element other = Element(1)) {
+
+  detail::TensorUpdateOffDiagonalFunc<Element, Layout> func(
+    dst,
+    other
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillLinearFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  Array<Element, Layout::kRank> v;
+  Element s;
+
+  //
+  // Methods
+  //
+
+  TensorFillLinearFunc() { }
+
+  /// Constructs functor
+  TensorFillLinearFunc(
+    TensorView const &view_,
+    Array<Element, Layout::kRank> const & v_,
+    Element s_ = Element(0)
+  ):
+    view(view_), v(v_), s(s_) { }
+
+  /// Updates the tensor
+  void operator()(Coord<Layout::kRank> const & coord) const {
+
+    Element sum(s);
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < Layout::kRank; ++i) {
+      sum += Element(coord[i]) * v[i];
+    }
+
+    view.at(coord) = sum;
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills tensor with a linear combination of its coordinate and another vector
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillLinear(
+  TensorView<Element, Layout> dst,      ///< destination tensor
+  Array<Element, Layout::kRank> const & v,
+  Element s = Element(0)) {
+
+  detail::TensorFillLinearFunc<Element, Layout> func(
+    dst,
+    v,
+    s
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills tensor with a linear combination of its coordinate and another vector
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillSequential(
+  TensorView<Element, Layout> dst,     ///< destination tensor
+  Element s = Element(0)) {
+
+  Array<Element, Layout::kRank> stride;
+
+  stride[0] = Element(1);
+
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 1; i < Layout::kRank; ++i) {
+    stride[i] = stride[i - 1] * Element(dst.extent()[i - 1]);
+  }
+
+  TensorFillLinear(dst, stride, s);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values from a distribution.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorFillRandom(
+  TensorView<Element, Layout> view,       ///< destination tensor
+  uint64_t seed,
+  Distribution dist,
+  bool exclude_zero = false               ///< If true, excludes 0.
+                                          ///  Note that setting this flag will result in more 1's,
+                                          ///  as we use a simple mechanism to replace 0's by adding/subtracting 1's.
+) {
+
+  using Real = typename RealType<Element>::Type;
+
+  if (dist.kind == Distribution::Gaussian) {
+    TensorFillRandomGaussian(
+      view,
+      seed,
+      dist.gaussian.mean,
+      dist.gaussian.stddev,
+      dist.int_scale,
+      dist.gaussian.pnz,
+      exclude_zero);
+  } else if (dist.kind == Distribution::Uniform) {
+    TensorFillRandomUniform(
+      view,
+      seed,
+      dist.uniform.max,
+      dist.uniform.min,
+      dist.int_scale,
+      dist.uniform.pnan,
+      exclude_zero);
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a block of data with sequential elements
+template <
+  typename Element
+>
+void BlockFillSequential(
+  Element *ptr,
+  int64_t capacity,
+  Element v = Element(1),
+  Element s = Element(0)) {
+  int i = 0;
+
+  while (i < capacity) {
+    mutlass::ReferenceFactory<Element, (mutlass::sizeof_bits<Element>::value <
+                                        8)>::get(ptr, i) = s;
+
+    s = Element(s + v);
+    ++i;
+  }
+}
+
+/// Fills a block of data with sequential elements
+template <
+  typename Element
+>
+void BlockFillSequentialModN(
+  Element *ptr,
+  int64_t capacity,
+  int64_t mod,
+  int64_t v = int64_t(1),
+  int64_t s = int64_t(0)) {
+  int i = 0;
+
+  while (i < capacity) {
+    mutlass::ReferenceFactory<Element, (mutlass::sizeof_bits<Element>::value <
+                                        8)>::get(ptr, i) = Element(s);
+
+    s = int64_t(s + v) % mod;
+    ++i;
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a block of data with sequential elements
+template <
+  typename Element
+>
+void BlockFillRandom(
+  Element *ptr,
+  size_t capacity,
+  uint64_t seed,
+  Distribution dist) {
+
+  if (dist.kind == Distribution::Gaussian) {
+    BlockFillRandomGaussian<Element>(
+      ptr,
+      capacity,
+      seed,
+      dist.gaussian.mean,
+      dist.gaussian.stddev,
+      dist.int_scale,
+      dist.gaussian.pnz);
+  }
+  else if (dist.kind == Distribution::Uniform) {
+    BlockFillRandomUniform<Element>(
+      ptr,
+      capacity,
+      seed,
+      dist.uniform.max,
+      dist.uniform.min,
+      dist.int_scale,
+      dist.uniform.pnan);
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <typename Element>
+struct RandomSparseMetaFunc {
+
+  uint64_t seed;
+  int range;
+  int MetaSizeInBits;
+
+  //
+  // Methods
+  //
+
+  RandomSparseMetaFunc(
+    uint64_t seed_ = 0,
+    int MetaSizeInBits_ = 2
+  ):
+    seed(seed_), MetaSizeInBits(MetaSizeInBits_) {
+      std::srand((unsigned)seed);
+      if (MetaSizeInBits_ == 2) {
+        range = 6;
+      }
+      else if (MetaSizeInBits_ == 4) {
+        range = 2;
+      }
+      else {
+        throw std::invalid_argument("Invalid MetaSizeInBits");
+      }
+    }
+
+  /// Compute random value and update RNG state
+  Element operator()() const {
+    Element FourToTwoMeta[6] = {0x4, 0x8, 0x9, 0xc, 0xd, 0xe};
+    Element TwoToOneMeta[2] = {0x4, 0xe};
+
+    Element * MetaArray = (MetaSizeInBits == 2) ? FourToTwoMeta : TwoToOneMeta;
+
+    Element result = 0x0;
+
+    for (int i = 0; i < mutlass::sizeof_bits<Element>::value / 4; ++i) {
+      int rnd = std::rand() % range;
+      Element meta = MetaArray[rnd];
+
+      result = (Element)(result | ((Element)(meta << (i * 4))));
+    }
+
+    return result;
+  }
+};
+
+/// Computes a random sparse meta
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+struct TensorFillRandomSparseMetaFunc {
+
+  using TensorView = TensorView<Element, Layout>;
+
+  //
+  // Data members
+  //
+
+  TensorView view;
+  RandomSparseMetaFunc<Element> func;
+
+  //
+  // Methods
+  //
+
+  /// Construction of Gaussian RNG functor.
+  TensorFillRandomSparseMetaFunc(
+    TensorView view_ = TensorView(),
+    RandomSparseMetaFunc<Element> func_ = RandomSparseMetaFunc<Element>()
+  ):
+    view(view_), func(func_) {
+
+  }
+
+  /// Compute random value and update RNG state
+  void operator()(Coord<Layout::kRank> const &coord) const {
+
+    view.at(coord) = func();
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <
+  typename Element,                 ///< Element type
+  typename Layout>                  ///< Layout function
+void TensorFillRandomSparseMeta(
+  TensorView<Element, Layout> dst,  ///< destination tensor
+  uint64_t seed,                    ///< seed for RNG
+  int MetaSizeInBits) {             ///< 2 bit or 4 bit
+
+  detail::RandomSparseMetaFunc<Element> random_func(seed, MetaSizeInBits);
+
+  detail::TensorFillRandomSparseMetaFunc<Element, Layout> func(
+    dst,
+    random_func
+  );
+
+  TensorForEach(
+    dst.extent(),
+    func
+  );
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <
+  typename Element                        ///< Element type
+>
+void BlockFillRandomSparseMeta(
+  Element *ptr,
+  size_t capacity,
+  uint64_t seed,                          ///< seed for RNG
+  int MetaSizeInBits) {                   ///< 2 bit or 4bit
+
+  detail::RandomSparseMetaFunc<Element> random_func(seed, MetaSizeInBits);
+
+  for (size_t i = 0; i < capacity; ++i) {
+    ptr[i] = random_func();
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a ell block index matrix with random values with a uniform random distribution.
+template <
+  typename Element,                                ///< Element type
+  typename Layout>                                 ///< Layout function
+void TensorFillRandomEllIdx(
+  TensorView<Element, Layout> dst,                 ///< destination tensor
+  uint64_t seed,                                   ///< seed for RNG
+  int rows, int ell_cols, int cols) {              ///< dimension of the matrix
+
+  std::srand((unsigned)seed);
+
+  for (int i = 0; i < rows; ++i) {
+    int col_idx = std::rand() % cols;
+
+    for (int j = 0; j < ell_cols; ++j) {
+      dst.at({i, j}) = col_idx;
+
+      if (col_idx != -1) {
+        if (col_idx == (cols - 1)) {
+          col_idx = -1;
+        } else {
+          col_idx = std::rand() % (cols - col_idx - 1) + col_idx + 1;
+        }
+      }
+    }
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies a diagonal in from host memory without modifying off-diagonal elements.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorCopyDiagonalIn(
+  TensorView<Element, Layout> dst,          ///< destination tensor
+  Element const *ptr) {                     ///< dense buffer of elements
+
+  typename Layout::Index extent = dst.extent().min();
+
+  for (typename Layout::Index i = 0; i < extent; ++i) {
+    Coord<Layout::kRank> coord(i);
+    dst.at(coord) = ReferenceFactory<Element>::get(ptr, i);
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Copies the diagonal of a tensor into a dense buffer in host memory.
+template <
+  typename Element,               ///< Element type
+  typename Layout>                ///< Layout function
+void TensorCopyDiagonalOut(
+  Element *ptr,                               ///< dense buffer of elements
+  TensorView<Element, Layout> src) {          ///< source tensor
+
+  typename Layout::Index extent = src.extent().min();
+
+  for (typename Layout::Index i = 0; i < extent; ++i) {
+    Coord<Layout::kRank> coord(i);
+    ReferenceFactory<Element>::get(ptr, i) = src.at(coord);
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_fill.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_fill.hpp
new file mode 100644
index 000000000..1375ba696
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_fill.hpp
@@ -0,0 +1,433 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Provides several functions for filling tensors with data.
+*/
+
+#pragma once
+
+// Standard Library includes
+#include <utility>
+#include <cstdlib>
+#include <cmath>
+
+// Mute includes
+#include "mute/tensor.hpp"
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+#include "mutlass/quaternion.h"
+#include "mutlass/array.h"
+#include "mutlass/numeric_types.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Uniform and procedural tensor fills
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with a scalar element
+template <typename Tensor>
+void TensorFill(Tensor dst, typename Tensor::value_type element) {
+
+  for (int64_t idx = 0; idx < mute::size(dst); ++idx) {
+    dst(idx) = element;
+  }
+}
+
+/// Fills a tensor with the contents of its layout
+template <typename Tensor>
+void TensorFillSequential(Tensor dst) {
+
+  auto layout = dst.layout();
+
+  for (int64_t idx = 0; idx < mute::size(dst); ++idx) {
+    dst(idx) = layout(idx);
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Random uniform values
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <typename Element>
+struct RandomUniformFunc {
+
+  using Real = typename RealType<Element>::Type;
+
+  uint64_t seed;
+  double range;
+  double min;
+  int int_scale;
+
+  //
+  // Methods
+  //
+
+  RandomUniformFunc(
+    uint64_t seed_ = 0,
+    double max = 1,
+    double min_ = 0,
+    int int_scale_ = -1
+  ):
+    seed(seed_), range(max - min_), min(min_), int_scale(int_scale_) {
+      std::srand((unsigned)seed);
+    }
+
+
+  /// Compute random value and update RNG state
+  Element operator()() const {
+
+    double rnd = double(std::rand()) / double(RAND_MAX);
+
+    rnd = min + range * rnd;
+
+    // Random values are cast to integer after scaling by a power of two to facilitate error
+    // testing
+    Element result;
+
+    if (int_scale >= 0) {
+      rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
+      result = static_cast<Element>(Real(rnd));
+    }
+    else {
+      result = static_cast<Element>(Real(rnd));
+    }
+
+    return result;
+  }
+};
+
+/// Partial specialization for initializing a complex value.
+template <typename Element>
+struct RandomUniformFunc<complex<Element> > {
+
+  using Real = typename RealType<Element>::Type;
+
+  uint64_t seed;
+  double range;
+  double min;
+  int int_scale;
+
+  //
+  // Methods
+  //
+
+  RandomUniformFunc(
+    uint64_t seed_ = 0,
+    double max = 1,
+    double min_ = 0,
+    int int_scale_ = -1
+  ):
+    seed(seed_), range(max - min_), min(min_), int_scale(int_scale_) {
+      std::srand((unsigned)seed);
+    }
+
+
+  /// Compute random value and update RNG state
+  complex<Element> operator()() const {
+
+    Element reals[2];
+
+    for (int i = 0; i < 2; ++i) {
+      double rnd = double(std::rand()) / double(RAND_MAX);
+
+      rnd = min + range * rnd;
+
+      // Random values are cast to integer after scaling by a power of two to facilitate error
+      // testing
+
+      if (int_scale >= 0) {
+        rnd = double(int(rnd * double(1 << int_scale)));
+        reals[i] = from_real<Element>(Real(rnd / double(1 << int_scale)));
+      }
+      else {
+        reals[i] = from_real<Element>(Real(rnd));
+      }
+    }
+
+    return complex<Element>(reals[0], reals[1]);
+  }
+};
+
+/// Partial specialization for initializing a Quaternion value.
+template <typename Element>
+struct RandomUniformFunc<Quaternion<Element> > {
+
+  using Real = typename RealType<Element>::Type;
+
+  uint64_t seed;
+  double range;
+  double min;
+  int int_scale;
+
+  //
+  // Methods
+  //
+
+  RandomUniformFunc(
+    uint64_t seed_ = 0,
+    double max = 1,
+    double min_ = 0,
+    int int_scale_ = -1
+  ):
+    seed(seed_), range(max - min_), min(min_), int_scale(int_scale_) {
+      std::srand((unsigned)seed);
+    }
+
+
+  /// Compute random value and update RNG state
+  Quaternion<Element> operator()() const {
+
+    Element reals[4];
+
+    for (int i = 0; i < 4; ++i) {
+      double rnd = double(std::rand()) / double(RAND_MAX);
+
+      rnd = min + range * rnd;
+
+      // Random values are cast to integer after scaling by a power of two to facilitate error
+      // testing
+
+      if (int_scale >= 0) {
+        rnd = double(int(rnd * double(1 << int_scale)));
+        reals[i] = from_real<Element>(Real(rnd / double(1 << int_scale)));
+      }
+      else {
+        reals[i] = from_real<Element>(Real(rnd));
+      }
+    }
+
+    return make_Quaternion(reals[0], reals[1], reals[2], reals[3]);
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a uniform random distribution.
+template <typename Tensor>                ///< Tensor object
+void TensorFillRandomUniform(
+  Tensor dst,                             ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  double max = 1,                         ///< upper bound of distribution
+  double min = 0,                         ///< lower bound for distribution
+  int bits = -1) {                        ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+
+  detail::RandomUniformFunc<typename Tensor::value_type> random_func(seed, max, min, bits);
+
+  for (int64_t idx = 0; idx < mute::size(dst); ++idx) {
+    dst(idx) = random_func();
+  }
+}
+
+/// Fills a block with random values with a uniform random distribution.
+template <
+  typename Element                        ///< Element type
+>
+void BlockFillRandomUniform(
+  Element *ptr,
+  size_t capacity,
+  uint64_t seed,                          ///< seed for RNG
+  double max = 1,                         ///< upper bound of distribution
+  double min = 0,                         ///< lower bound for distribution
+  int bits = -1) {                        ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+  detail::RandomUniformFunc<Element> random_func(seed, max, min, bits);
+
+  for (size_t i = 0; i < capacity; ++i) {
+    ptr[i] = random_func();
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Random Gaussian
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+template <typename Element>
+struct RandomGaussianFunc {
+
+  uint64_t seed;
+  double mean;
+  double stddev;
+  int int_scale;
+  double pi;
+
+  //
+  // Methods
+  //
+  RandomGaussianFunc(
+    uint64_t seed_ = 0,
+    double mean_ = 0,
+    double stddev_ = 1,
+    int int_scale_ = -1
+  ):
+    seed(seed_), mean(mean_), stddev(stddev_), int_scale(int_scale_), pi(std::acos(-1)) {
+      std::srand((unsigned)seed);
+  }
+
+  /// Compute random value and update RNG state
+  Element operator()() const {
+
+    // Box-Muller transform to generate random numbers with Normal distribution
+    double u1 = double(std::rand()) / double(RAND_MAX);
+    double u2 = double(std::rand()) / double(RAND_MAX);
+
+    // Compute Gaussian random value
+    double rnd = std::sqrt(-2 * std::log(u1)) * std::cos(2 * pi * u2);
+    rnd = mean + stddev * rnd;
+
+    // Scale and convert final result
+    Element result;
+
+    if (int_scale >= 0) {
+      rnd = double(int64_t(rnd * double(1 << int_scale))) / double(1 << int_scale);
+      result = static_cast<Element>(rnd);
+    }
+    else {
+      result = static_cast<Element>(rnd);
+    }
+
+    return result;
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a tensor with random values with a Gaussian distribution.
+template <
+  typename Tensor
+>
+void TensorFillRandomGaussian(
+  Tensor  dst,                            ///< destination tensor
+  uint64_t seed,                          ///< seed for RNG
+  double mean = 0,                        ///< Gaussian distribution's mean
+  double stddev = 1,                      ///< Gaussian distribution's standard deviation
+  int bits = -1) {                        ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+
+  detail::RandomGaussianFunc<typename Tensor::value_type> random_func(seed, mean, stddev, bits);
+
+  for (int64_t idx = 0; idx < mute::size(dst); ++idx) {
+    dst(idx) = random_func();
+  }
+}
+
+/// Fills a block with random values with a Gaussian distribution.
+template <
+  typename Element                        ///< Element type
+>
+void BlockFillRandomGaussian(
+  Element *ptr,                           ///< destination buffer
+  size_t capacity,                        ///< number of elements
+  uint64_t seed,                          ///< seed for RNG
+  double mean = 0,                        ///< Gaussian distribution's mean
+  double stddev = 1,                      ///< Gaussian distribution's standard deviation
+  int bits = -1) {                        ///< If non-negative, specifies number of fractional bits that
+                                          ///  are not truncated to zero. Permits reducing precision of
+                                          ///  data.
+
+  detail::RandomGaussianFunc<Element> random_func(seed, mean, stddev, bits);
+
+  for (size_t i = 0; i < capacity; ++i) {
+    ptr[i] = random_func();
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Fills a block of data with sequential elements
+template <
+  typename Element
+>
+void BlockFillSequential(
+  Element *ptr,
+  int64_t capacity,
+  Element v = Element(1),
+  Element s = Element(0)) {
+  int i = 0;
+
+  while (i < capacity) {
+
+    ptr[i] = Element(s + v);
+    ++i;
+  }
+}
+
+/// Fills a block of data with sequential elements
+template <
+  typename Element
+>
+void BlockFillSequentialModN(
+  Element *ptr,
+  int64_t capacity,
+  int64_t mod,
+  int64_t v = int64_t(1),
+  int64_t s = int64_t(0)) {
+  int i = 0;
+
+  while (i < capacity) {
+
+    ptr[i] = static_cast<Element>(int32_t(int64_t(s + v) % mod));
+    ++i;
+  }
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_foreach.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_foreach.h
new file mode 100644
index 000000000..4072009b7
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_foreach.h
@@ -0,0 +1,135 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <stdexcept>
+#include "mutlass/mutlass.h"
+
+namespace mutlass  {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Defines several helpers
+namespace detail {
+
+/// Helper to perform for-each operation
+template <typename Func, int Rank, int RankRemaining>
+struct TensorForEachHelper {
+
+  /// Index of the active rank
+  static int const kActiveRank = Rank - RankRemaining - 1;
+
+  /// Constructor for general rank
+  TensorForEachHelper(
+    Func &func,
+    Coord<Rank> const &extent,
+    Coord<Rank> &coord) {
+
+    for (int i = 0; i < extent.at(kActiveRank); ++i) {
+      coord[kActiveRank] = i;
+      TensorForEachHelper<Func, Rank, RankRemaining - 1>(func, extent, coord);
+    }
+  }
+};
+
+/// Helper to perform for-each operation
+template <typename Func, int Rank>
+struct TensorForEachHelper<Func, Rank, 0> {
+
+  /// Index of the active rank
+  static int const kActiveRank = Rank - 1;
+
+  /// Constructor for fastest changing rank
+  TensorForEachHelper(
+    Func &func,
+    Coord<Rank> const &extent,
+    Coord<Rank> &coord) {
+
+    for (int i = 0; i < extent.at(kActiveRank); ++i) {
+      coord[kActiveRank] = i;
+      func(coord);
+    }
+  }
+};
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Iterates over the index space of a tensor
+template <
+  typename Func,          ///< function applied to each point in a tensor's index space
+  int Rank>               ///< rank of index space
+void TensorForEach(Coord<Rank> extent, Func & func) {
+  Coord<Rank> coord;
+  detail::TensorForEachHelper<Func, Rank, Rank - 1>(func, extent, coord);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Iterates over the index space of a tensor and calls a C++ lambda
+template <
+  typename Func,          ///< function applied to each point in a tensor's index space
+  int Rank>               ///< rank of index space
+void TensorForEachLambda(Coord<Rank> extent, Func func) {
+  Coord<Rank> coord;
+  detail::TensorForEachHelper<Func, Rank, Rank - 1>(func, extent, coord);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Element, typename Func>
+struct BlockForEach {
+
+  /// Constructor performs the operation.
+  BlockForEach(
+    Element *ptr,
+    size_t capacity,
+    typename Func::Params params = typename Func::Params()) {
+
+    Func func(params);
+
+    for (size_t index = 0; index < capacity; ++index) {
+      ptr[index] = func();
+    }
+  }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_norm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_norm.h
new file mode 100644
index 000000000..1a6013721
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_norm.h
@@ -0,0 +1,41 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+
+#include "mutlass/mutlass.h"
+
+// The contents of this file have been moved  to 'tensor_reduce' to cover other types of reductions.
+
+#include "mutlass/util/reference/host/tensor_reduce.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_reduce.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_reduce.h
new file mode 100644
index 000000000..4cfd17066
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_reduce.h
@@ -0,0 +1,204 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include <cmath>
+
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+#include "mutlass/tensor_ref.h"
+
+#include "mutlass/util/reference/detail/linear_to_coordinate.h"
+#include "mutlass/core_io.h"
+
+namespace mutlass  {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Transform-reduce operation over the elements of a tensor. This helper allocates the device-side
+/// workspace
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp
+>
+ComputeType TensorTransformReduce(
+  TensorView<Element, Layout> view,
+  ComputeType identity,
+  ReduceOp reduce,
+  TransformOp transform
+) {
+
+  for (int64_t idx = 0; idx < view.size(); ++idx) {
+    typename Layout::TensorCoord coord;
+    mutlass::reference::detail::LinearToCoordinate<Layout::kRank>()(coord, idx, view.extent());
+
+    if (view.contains(coord)) {
+      Element x = view.at(coord);
+      identity = reduce(identity, transform(x));
+    }
+  }
+
+  return identity;
+}
+
+/// Transform-reduce operation over the elements of a tensor. This helper allocates the device-side
+/// workspace
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp
+>
+ComputeType TensorTransformReduce(
+  TensorView<Element, Layout> view_A,
+  TensorView<Element, Layout> view_B,
+  ComputeType identity,
+  ReduceOp reduce,
+  TransformOp transform) {
+
+  if (view_A.extent() != view_B.extent()) {
+    throw std::runtime_error("Tensor extents must match.");
+  }
+
+  for (int64_t idx = 0; idx < view_A.size(); ++idx) {
+
+    typename Layout::TensorCoord coord;
+    mutlass::reference::detail::LinearToCoordinate<Layout::kRank>()(coord, idx, view_A.extent());
+
+    if (view_A.contains(coord)) {
+      Element a = view_A.at(coord);
+      Element b = view_B.at(coord);
+      identity = reduce(identity, transform(a, b));
+    }
+  }
+
+  return identity;
+}
+
+/// Helper to compute the sum of the elements of a tensor
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = Element
+>
+ComputeType TensorSum(
+  TensorView<Element, Layout> view,
+  ComputeType identity = ComputeType()
+) {
+
+  plus<ComputeType> reduce;
+  NumericConverter<ComputeType, Element> transform;
+
+  return TensorTransformReduce(
+    view, identity, reduce, transform);
+}
+
+/// Helper to compute the sum of the squares of the elements of a tensor
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = Element
+>
+ComputeType TensorSumSq(
+  TensorView<Element, Layout> view,
+  ComputeType identity = ComputeType()
+) {
+
+  plus<ComputeType> reduce;
+  magnitude_squared<Element, ComputeType> transform;
+
+  return TensorTransformReduce(
+    view, identity, reduce, transform);
+}
+
+/// Helper to compute the norm of the elements of a tensor.
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = double
+>
+ComputeType TensorNorm(
+  TensorView<Element, Layout> view,
+  ComputeType identity = ComputeType()
+) {
+
+  return std::sqrt(TensorSumSq(view, identity));
+}
+
+/// Helper to compute the sum of the squares of the differences of two tensors
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = double
+>
+ComputeType TensorSumSqDiff(
+  TensorView<Element, Layout> view_A,
+  TensorView<Element, Layout> view_B,
+  ComputeType identity = ComputeType()
+) {
+
+  plus<ComputeType> reduce;
+  magnitude_squared_difference<Element, ComputeType> transform;
+
+  return TensorTransformReduce(
+    view_A, view_B, identity, reduce, transform);
+}
+
+
+/// Helper to compute the norm of the tensor computed as the difference of two tensors in memory
+template <
+  typename Element,
+  typename Layout,
+  typename ComputeType = double
+>
+ComputeType TensorNormDiff(
+  TensorView<Element, Layout> view_A,
+  TensorView<Element, Layout> view_B,
+  ComputeType identity = ComputeType()
+) {
+
+  return std::sqrt(TensorSumSqDiff(view_A, view_B, identity));
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_reduce.hpp b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_reduce.hpp
new file mode 100644
index 000000000..28177bffe
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/tensor_reduce.hpp
@@ -0,0 +1,204 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Provides several functions for filling tensors with data.
+*/
+
+#pragma once
+
+// Standard Library includes
+#include <utility>
+#include <cstdlib>
+#include <cmath>
+
+// Mute includes
+#include "mute/tensor.hpp"
+
+// Mutlass includes
+#include "mutlass/mutlass.h"
+#include "mutlass/complex.h"
+#include "mutlass/functional.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/quaternion.h"
+#include "mutlass/array.h"
+#include "mutlass/numeric_types.h"
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Tensor reductions
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Transform-reduce operation over the elements of a tensor. This helper allocates the device-side
+/// workspace
+template <
+  typename Tensor,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp
+>
+ComputeType TensorTransformReduce(
+  Tensor view,
+  ComputeType identity,
+  ReduceOp reduce,
+  TransformOp transform
+) {
+
+  for (int64_t idx = 0; idx < mute::size(view); ++idx) {
+    identity = reduce(identity, transform(view(idx)));
+  }
+
+  return identity;
+}
+
+/// Transform-reduce operation over the elements of a tensor. This helper allocates the device-side
+/// workspace
+template <
+  typename TensorA,
+  typename TensorB,
+  typename ComputeType,
+  typename ReduceOp,
+  typename TransformOp
+>
+ComputeType TensorTransformReduce(
+  TensorA view_A,
+  TensorB view_B,
+  ComputeType identity,
+  ReduceOp reduce,
+  TransformOp transform) {
+
+  if (mute::size(view_A) != mute::size(view_B)) {
+    throw std::runtime_error("Tensor sizes must match.");
+  }
+
+  for (int64_t idx = 0; idx < mute::size(view_A); ++idx) {
+    identity = reduce(identity, transform(view_A(idx), view_B(idx)));
+  }
+
+  return identity;
+}
+
+/// Helper to compute the sum of the elements of a tensor
+template <
+  typename Tensor,
+  typename ComputeType = typename Tensor::value_type
+>
+ComputeType TensorSum(
+  Tensor view,
+  ComputeType identity = ComputeType()
+) {
+
+  plus<ComputeType> reduce;
+  NumericConverter<ComputeType, typename Tensor::value_type> transform;
+
+  return TensorTransformReduce(
+    view, identity, reduce, transform);
+}
+
+/// Helper to compute the sum of the squares of the elements of a tensor
+template <
+  typename Tensor,
+  typename ComputeType = typename Tensor::value_type
+>
+ComputeType TensorSumSq(
+  Tensor view,
+  ComputeType identity = ComputeType()
+) {
+
+  plus<ComputeType> reduce;
+  magnitude_squared<typename Tensor::value_type, ComputeType> transform;
+
+  return TensorTransformReduce(
+    view, identity, reduce, transform);
+}
+
+/// Helper to compute the norm of the elements of a tensor.
+template <
+  typename Tensor,
+  typename ComputeType = double
+>
+ComputeType TensorNorm(
+  Tensor view,
+  ComputeType identity = ComputeType()
+) {
+
+  return std::sqrt(TensorSumSq(view, identity));
+}
+
+/// Helper to compute the sum of the squares of the differences of two tensors
+template <
+  typename TensorA,
+  typename TensorB,
+  typename ComputeType = double
+>
+ComputeType TensorSumSqDiff(
+  TensorA view_A,
+  TensorB view_B,
+  ComputeType identity = ComputeType()
+) {
+
+  plus<ComputeType> reduce;
+  magnitude_squared_difference<typename TensorA::value_type, ComputeType> transform;
+
+  return TensorTransformReduce(
+    view_A, view_B, identity, reduce, transform);
+}
+
+
+/// Helper to compute the norm of the tensor computed as the difference of two tensors in memory
+template <
+  typename TensorA,
+  typename TensorB,
+  typename ComputeType = double
+>
+ComputeType TensorNormDiff(
+  TensorA view_A,
+  TensorB view_B,
+  ComputeType identity = ComputeType()
+) {
+
+  return std::sqrt(TensorSumSqDiff(view_A, view_B, identity));
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/trmm.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/trmm.h
new file mode 100644
index 000000000..d9e4dbea6
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/trmm.h
@@ -0,0 +1,216 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for TRMM in host-side code.
+
+
+*/
+
+#pragma once
+
+#include "mutlass/blas3.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+#include "mutlass/arch/mma.h"
+#include "mutlass/util/host_tensor.h"
+
+#include "mutlass/util/reference/host/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+/// Computes a Triangular Matrix Multiplication (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename ElementA,
+  typename LayoutA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  DiagType DiagTypeA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_trmm(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  static_assert(SideModeA != SideMode::kInvalid
+                , "Side Mode can either be Left or Right.");
+
+  static_assert(FillModeA == FillMode::kLower || FillModeA == FillMode::kUpper
+                , "Fill Mode can either be Lower or Upper.");
+
+  using CompareOp = typename TrMatrixCompareOp<FillModeA, DiagTypeA>::Type;
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  // Assuming correct k-dimension value is passed
+  int const K = problem_size.k();
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+  CompareOp compare_op;
+
+  for (int row_block = 0; row_block < M; row_block += Mblock) {
+    for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+      ComputeType accum[Mblock][Nblock];
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          accum[i][j] = initial_accum;
+        }
+      }
+
+      for (int k_block = 0; k_block < K; ++k_block) {
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            if (row < M && col < N) {
+              ElementA a = ElementA();
+              ElementB b = ElementB();
+
+              if (SideModeA == SideMode::kLeft) {
+                a = (compare_op(row, k_block)) ?
+                            (tensor_a.at(MatrixCoord(row, k_block))) : ElementA(0);
+                if (row == k_block && DiagTypeA == DiagType::kUnit) {
+                  a = ElementA(1);
+                }
+                b = tensor_b.at(MatrixCoord(k_block, col));
+              } else if (SideModeA == SideMode::kRight) {
+                a = tensor_b.at(MatrixCoord(row, k_block));
+                b = (compare_op(k_block, col)) ?
+                      tensor_a.at(MatrixCoord(k_block, col)) : ElementA(0);
+                if (k_block == col && DiagTypeA == DiagType::kUnit) {
+                  b = ElementA(1);
+                }
+              }
+
+              ComputeType compute_a(cast_if_scalar<ComputeType>(a));
+              ComputeType compute_b(cast_if_scalar<ComputeType>(b));
+
+              accum[i][j] = inner_product_op(compute_a, compute_b, accum[i][j]);
+            }
+          }
+        }
+      }
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          int row = row_block + i;
+          int col = col_block + j;
+
+          MatrixCoord coord = MatrixCoord(row, col);
+
+          if (row < M && col < N) {
+            tensor_d.at(coord) = convert_op(
+              alpha * ScalarType(accum[i][j]));
+          }
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  DiagType DiagTypeA,
+  typename ElementB,
+  typename LayoutB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = mutlass::arch::OpMultiplyAdd
+>
+struct Trmm;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA, SideMode SideModeA,
+           FillMode FillModeA, DiagType DiagTypeA,
+           typename ElementB, typename LayoutB,
+           typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct Trmm<ElementA, LayoutA, SideModeA, FillModeA, DiagTypeA, ElementB, LayoutB,
+            ElementC, LayoutC, ScalarType,
+            ComputeType, arch::OpMultiplyAdd> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_trmm<ElementA, LayoutA, SideModeA, FillModeA, DiagTypeA, ElementB, LayoutB,
+                 ElementC, LayoutC, ScalarType, ComputeType, multiply_add<ComputeType>>(
+                 problem_size, alpha, tensor_a, tensor_b, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/trmm_complex.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/trmm_complex.h
new file mode 100644
index 000000000..9cd6ecfa5
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/reference/host/trmm_complex.h
@@ -0,0 +1,263 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Reference implementation for complex-valued TRMM in host-side code.
+
+
+*/
+
+#pragma once
+
+#include "mutlass/blas3.h"
+#include "mutlass/complex.h"
+#include "mutlass/numeric_conversion.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/gemm/gemm.h"
+
+#include "mutlass/util/reference/host/gemm.h"
+
+namespace mutlass {
+namespace reference {
+namespace host {
+
+/// Computes a Triangular Matrix Multiplication (tensors of rank=2) pointed to by TensorRef
+/// objects.
+template <
+  typename ElementA,
+  typename LayoutA,
+  ComplexTransform TransformA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  DiagType DiagTypeA,
+  typename ElementB,
+  typename LayoutB,
+  ComplexTransform TransformB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = multiply_add<ComputeType>,
+  typename ConvertOp = NumericConverter<ElementC, ScalarType>
+>
+void compute_trmm_complex(
+  gemm::GemmCoord problem_size,
+  ScalarType alpha,
+  TensorRef<ElementA, LayoutA> tensor_a,
+  TensorRef<ElementB, LayoutB> tensor_b,
+  TensorRef<ElementC, LayoutC> tensor_d,
+  ComputeType initial_accum) {
+
+  static_assert(
+    LayoutA::kRank == 2 &&
+    LayoutC::kRank == 2, "Tensors must be of rank 2");
+
+  static_assert(SideModeA != SideMode::kInvalid
+                , "Side Mode can either be Left or Right.");
+
+  static_assert(FillModeA == FillMode::kLower || FillModeA == FillMode::kUpper
+                , "Fill Mode can either be Lower or Upper.");
+
+  using CompareOp = typename TrMatrixCompareOp<FillModeA, DiagTypeA>::Type;
+
+  // Note: batch is ignored.
+  int const M = problem_size.m();
+  int const N = problem_size.n();
+  // Assuming correct k-dimension value is passed
+  int const K = problem_size.k();
+
+  // Blocking necessary to speedup reference implementation
+  int const Mblock = 16;
+  int const Nblock = 16;
+
+  ConvertOp convert_op;
+  InnerProductOp inner_product_op;
+  CompareOp compare_op;
+
+  for (int row_block = 0; row_block < M; row_block += Mblock) {
+    for (int col_block = 0; col_block < N; col_block += Nblock) {
+
+      ComputeType accum[Mblock][Nblock];
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          accum[i][j] = initial_accum;
+        }
+      }
+
+      for (int k_block = 0; k_block < K; ++k_block) {
+        for (int j = 0; j < Nblock; j++) {
+          for (int i = 0; i < Mblock; i++) {
+            int row = row_block + i;
+            int col = col_block + j;
+
+            if (row < M && col < N) {
+              ElementA a = ElementA();
+              ElementB b = ElementB();
+
+              if (SideModeA == SideMode::kLeft) {
+                a = (compare_op(row, k_block)) ?
+                              (tensor_a.at(MatrixCoord(row, k_block))) : ElementA(0);
+                if (row == k_block && DiagTypeA == DiagType::kUnit) {
+                  a = ElementA(1);
+                }
+                b = tensor_b.at(MatrixCoord(k_block, col));
+              } else if (SideModeA == SideMode::kRight) {
+                a = tensor_b.at(MatrixCoord(row, k_block));
+                b = (compare_op(k_block, col)) ?
+                      tensor_a.at(MatrixCoord(k_block, col)) : ElementA(0);
+                if (k_block == col && DiagTypeA == DiagType::kUnit) {
+                  b = ElementA(1);
+                }
+              }
+
+              ComputeType a_ik = ComputeType(a);
+              ComputeType b_kj = ComputeType(b);
+
+              // Conjugate, and hence hermitian, is only allowed for the triangular matrix
+              if (SideModeA == SideMode::kLeft && TransformA == ComplexTransform::kConjugate) {
+                a_ik = conj(a_ik);
+              } else if (SideModeA == SideMode::kRight && TransformA == ComplexTransform::kConjugate) {
+                b_kj = conj(b_kj);
+              }
+
+              accum[i][j] = inner_product_op(a_ik, b_kj,  accum[i][j]);
+            }
+          }
+        }
+      }
+
+      for (int j = 0; j < Nblock; j++) {
+        for (int i = 0; i < Mblock; i++) {
+          int row = row_block + i;
+          int col = col_block + j;
+
+          MatrixCoord coord = MatrixCoord(row, col);
+
+          if (row < M && col < N) {
+            tensor_d.at(coord) = convert_op(
+              alpha * ScalarType(accum[i][j]));
+          }
+        }
+      }
+    }
+  }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <
+  typename ElementA,
+  typename LayoutA,
+  ComplexTransform TransformA,
+  SideMode SideModeA,
+  FillMode FillModeA,
+  DiagType DiagTypeA,
+  typename ElementB,
+  typename LayoutB,
+  ComplexTransform TransformB,
+  typename ElementC,
+  typename LayoutC,
+  typename ScalarType,
+  typename ComputeType,
+  typename InnerProductOp = mutlass::arch::OpMultiplyAddComplex
+>
+struct TrmmComplex;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for multiply-add
+template <typename ElementA, typename LayoutA, ComplexTransform TransformA,
+          SideMode SideModeA, FillMode FillModeA, DiagType DiagTypeA,
+          typename ElementB, typename LayoutB, ComplexTransform TransformB,
+          typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct TrmmComplex<ElementA, LayoutA, TransformA,
+                   SideModeA, FillModeA, DiagTypeA,
+                   ElementB, LayoutB, TransformB,
+                   ElementC, LayoutC, ScalarType,
+                   ComputeType, arch::OpMultiplyAddComplex> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_trmm_complex<ElementA, LayoutA, TransformA,
+                 SideModeA, FillModeA, DiagTypeA,
+                 ElementB, LayoutB, TransformB,
+                 ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+                 problem_size, alpha, tensor_a, tensor_b, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Partial specialization for gaussian multiply-add
+template <typename ElementA, typename LayoutA, ComplexTransform TransformA,
+          SideMode SideModeA, FillMode FillModeA, DiagType DiagTypeA,
+          typename ElementB, typename LayoutB, ComplexTransform TransformB,
+          typename ElementC, typename LayoutC,
+          typename ScalarType, typename ComputeType>
+struct TrmmComplex<ElementA, LayoutA, TransformA,
+                   SideModeA, FillModeA, DiagTypeA,
+                   ElementB, LayoutB, TransformB,
+                   ElementC, LayoutC, ScalarType,
+                   ComputeType, arch::OpMultiplyAddGaussianComplex> {
+
+  void operator()(gemm::GemmCoord problem_size, ScalarType alpha,
+                  TensorRef<ElementA, LayoutA> tensor_a,
+                  TensorRef<ElementB, LayoutB> tensor_b,
+                  TensorRef<ElementC, LayoutC> tensor_d,
+                  ComputeType initial_accum = ComputeType(0)) {
+    static_assert(
+        LayoutA::kRank == 2 && LayoutC::kRank == 2,
+        "Tensors must be of rank 2");
+
+    compute_trmm_complex<ElementA, LayoutA, TransformA,
+                 SideModeA, FillModeA, DiagTypeA,
+                 ElementB, LayoutB, TransformB,
+                 ElementC, LayoutC,
+                 ScalarType, ComputeType, multiply_add<ComputeType>>(
+                 problem_size, alpha, tensor_a, tensor_b, tensor_d, initial_accum);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace host
+} // namespace reference
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/tensor_view_io.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/tensor_view_io.h
new file mode 100644
index 000000000..ff2dc55d8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/tensor_view_io.h
@@ -0,0 +1,271 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+#pragma once
+
+#include "mutlass/core_io.h"
+#include "mutlass/tensor_view.h"
+#include "mutlass/tensor_view_planar_complex.h"
+#include "mutlass/complex.h"
+
+namespace mutlass {
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace detail {
+
+/// Helper to write the least significant rank of a TensorView
+template <
+  typename Element,
+  typename Layout
+>
+inline std::ostream & TensorView_WriteLeastSignificantRank(
+  std::ostream& out,
+  TensorView<Element, Layout> const& view,
+  Coord<Layout::kRank> const &start_coord,
+  int rank,
+  std::streamsize width) {
+
+  for (int idx = 0; idx < view.extent(rank); ++idx) {
+
+    Coord<Layout::kRank> coord(start_coord);
+    coord[rank] = idx;
+
+    if (idx) {
+      out.width(0);
+      out << ", ";
+    }
+    if (idx || coord) {
+      out.width(width);
+    }
+    out << ScalarIO<Element>(view.at(coord));
+  }
+
+  return out;
+}
+
+/// Helper to write a rank of a TensorView
+template <
+  typename Element,
+  typename Layout
+>
+inline std::ostream & TensorView_WriteRank(
+  std::ostream& out,
+  TensorView<Element, Layout> const& view,
+  Coord<Layout::kRank> const &start_coord,
+  int rank,
+  std::streamsize width) {
+
+  // If called on the least significant rank, write the result as a row
+  if (rank + 1 == Layout::kRank) {
+    return TensorView_WriteLeastSignificantRank(out, view, start_coord, rank, width);
+  }
+
+  // Otherwise, write a sequence of rows and newlines
+  for (int idx = 0; idx < view.extent(rank); ++idx) {
+
+    Coord<Layout::kRank> coord(start_coord);
+    coord[rank] = idx;
+
+    if (rank + 2 == Layout::kRank) {
+      // Write least significant ranks asa matrix with rows delimited by "\n"
+      if (idx) {
+        out << ",\n";
+      }
+      TensorView_WriteLeastSignificantRank(out, view, coord, rank + 1, width);
+    }
+    else {
+      // Higher ranks are separated by newlines
+      if (idx) {
+        out << ",\n\n";
+      }
+      TensorView_WriteRank(out, view, coord, rank + 1, width);
+    }
+  }
+
+  return out;
+}
+
+/// Helper to write the least significant rank of a TensorView
+template <
+  typename Element,
+  typename Layout
+>
+inline std::ostream & TensorViewPlanarComplex_WriteLeastSignificantRank(
+  std::ostream& out,
+  TensorViewPlanarComplex<Element, Layout> const& view,
+  Coord<Layout::kRank> const &start_coord,
+  int rank,
+  std::streamsize width) {
+
+  for (int idx = 0; idx < view.extent(rank); ++idx) {
+
+    Coord<Layout::kRank> coord(start_coord);
+    coord[rank] = idx;
+
+    if (idx) {
+      out.width(0);
+      out << ", ";
+    }
+    if (idx || coord) {
+      out.width(width);
+    }
+
+    complex<Element> x = view.at(coord);
+    out << x;
+  }
+
+  return out;
+}
+
+/// Helper to write a rank of a TensorView
+template <
+  typename Element,
+  typename Layout
+>
+inline std::ostream & TensorViewPlanarComplex_WriteRank(
+  std::ostream& out,
+  TensorViewPlanarComplex<Element, Layout> const& view,
+  Coord<Layout::kRank> const &start_coord,
+  int rank,
+  std::streamsize width) {
+
+  // If called on the least significant rank, write the result as a row
+  if (rank + 1 == Layout::kRank) {
+    return TensorViewPlanarComplex_WriteLeastSignificantRank(out, view, start_coord, rank, width);
+  }
+
+  // Otherwise, write a sequence of rows and newlines
+  for (int idx = 0; idx < view.extent(rank); ++idx) {
+
+    Coord<Layout::kRank> coord(start_coord);
+    coord[rank] = idx;
+
+    if (rank + 2 == Layout::kRank) {
+      // Write least significant ranks asa matrix with rows delimited by ";\n"
+      if (idx) {
+        out << ";\n";
+      }
+      TensorViewPlanarComplex_WriteLeastSignificantRank(out, view, coord, rank + 1, width);
+    }
+    else {
+      // Higher ranks are separated by newlines
+      if (idx) {
+        out << "\n";
+      }
+      TensorViewPlanarComplex_WriteRank(out, view, coord, rank + 1, width);
+    }
+  }
+
+  return out;
+}
+
+} // namespace detail
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Prints human-readable representation of a TensorView to an ostream
+template <
+  typename Element,
+  typename Layout
+>
+inline std::ostream& TensorViewWrite(
+  std::ostream& out,
+  TensorView<Element, Layout> const& view) {
+
+  // Prints a TensorView according to the following conventions:
+  //   - least significant rank is printed as rows separated by ";\n"
+  //   - all greater ranks are delimited with newlines
+  //
+  // The result is effectively a whitespace-delimited series of 2D matrices.
+
+  return detail::TensorView_WriteRank(out, view, Coord<Layout::kRank>(), 0, out.width());
+}
+
+/// Prints human-readable representation of a TensorView to an ostream
+template <
+  typename Element,
+  typename Layout
+>
+inline std::ostream& operator<<(
+  std::ostream& out,
+  TensorView<Element, Layout> const& view) {
+
+  // Prints a TensorView according to the following conventions:
+  //   - least significant rank is printed as rows separated by ";\n"
+  //   - all greater ranks are delimited with newlines
+  //
+  // The result is effectively a whitespace-delimited series of 2D matrices.
+
+  return TensorViewWrite(out, view);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+/// Prints human-readable representation of a TensorView to an ostream
+template <
+  typename Element,
+  typename Layout
+>
+inline std::ostream& TensorViewWrite(
+  std::ostream& out,
+  TensorViewPlanarComplex<Element, Layout> const& view) {
+
+  // Prints a TensorView according to the following conventions:
+  //   - least significant rank is printed as rows separated by ";\n"
+  //   - all greater ranks are delimited with newlines
+  //
+  // The result is effectively a whitespace-delimited series of 2D matrices.
+
+  return detail::TensorViewPlanarComplex_WriteRank(out, view, Coord<Layout::kRank>(), 0, out.width());
+}
+
+/// Prints human-readable representation of a TensorView to an ostream
+template <
+  typename Element,
+  typename Layout
+>
+inline std::ostream& operator<<(
+  std::ostream& out,
+  TensorViewPlanarComplex<Element, Layout> const& view) {
+
+  // Prints a TensorView according to the following conventions:
+  //   - least significant rank is printed as rows separated by ";\n"
+  //   - all greater ranks are delimited with newlines
+  //
+  // The result is effectively a whitespace-delimited series of 2D matrices.
+
+  return TensorViewWrite(out, view);
+}
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/util/type_traits.h b/vllm_musa/csrc/flash_attn/include/mutlass/util/type_traits.h
new file mode 100644
index 000000000..ad3c252bf
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/util/type_traits.h
@@ -0,0 +1,239 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Type traits for common MUSA types
+*/
+
+#pragma once
+
+#include <mublas_v2.h>
+#include <musa_fp16.h>
+#include <stdint.h>
+
+#include "mutlass/numeric_types.h"
+#include "mutlass/complex.h"
+
+namespace mutlass {
+struct half_t;
+
+template <typename T>
+struct TypeTraits {
+  typedef T host_type;
+  typedef T device_type;
+  static inline T remove_negative_zero(T x) { return x; }
+  static inline T to_print(T x) { return x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+template <>
+struct TypeTraits<int8_t> {
+  static musaDataType_t const mublas_type = MUSA_R_8I;
+  typedef int8_t host_type;
+  typedef int8_t device_type;
+  typedef int8_t integer_type;
+  typedef uint8_t unsigned_type;
+  static inline int8_t remove_negative_zero(int8_t x) { return x; }
+  static inline int to_print(int8_t x) { return (int)x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+template <>
+struct TypeTraits<uint8_t> {
+  static musaDataType_t const mublas_type = MUSA_R_8I;
+  typedef uint8_t host_type;
+  typedef uint8_t device_type;
+  typedef uint8_t integer_type;
+  typedef uint8_t unsigned_type;
+  static inline uint8_t remove_negative_zero(uint8_t x) { return x; }
+  static inline uint32_t to_print(uint8_t x) { return (uint32_t)x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+template <>
+struct TypeTraits<int> {
+  static musaDataType_t const mublas_type = MUSA_R_32I;
+  typedef int host_type;
+  typedef int device_type;
+  typedef int32_t integer_type;
+  typedef uint32_t unsigned_type;
+  static inline int32_t remove_negative_zero(int32_t x) { return x; }
+  static inline int to_print(int x) { return x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+template <>
+struct TypeTraits<unsigned> {
+  static musaDataType_t const mublas_type = MUSA_R_32I;
+  typedef unsigned host_type;
+  typedef unsigned device_type;
+  typedef uint32_t integer_type;
+  typedef uint32_t unsigned_type;
+  static inline uint32_t remove_negative_zero(uint32_t x) { return x; }
+  static inline uint32_t to_print(uint32_t x) { return x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+template <>
+struct TypeTraits<int64_t> {
+  static musaDataType_t const mublas_type = MUSA_R_8I;
+  typedef int64_t host_type;
+  typedef int64_t device_type;
+  typedef int64_t integer_type;
+  typedef uint64_t unsigned_type;
+  static inline int64_t remove_negative_zero(int64_t x) { return x; }
+  static inline int64_t to_print(int64_t x) { return x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+template <>
+struct TypeTraits<uint64_t> {
+  static musaDataType_t const mublas_type = MUSA_R_8I;
+  typedef uint64_t host_type;
+  typedef uint64_t device_type;
+  typedef uint64_t integer_type;
+  typedef uint64_t unsigned_type;
+  static inline uint64_t remove_negative_zero(uint64_t x) { return x; }
+  static inline uint64_t to_print(uint64_t x) { return x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+template <>
+struct TypeTraits<half_t> {
+  static musaDataType_t const mublas_type = MUSA_R_16F;
+  typedef half_t host_type;
+  typedef half_t device_type;
+  typedef int16_t integer_type;
+  typedef uint16_t unsigned_type;
+  static inline half_t remove_negative_zero(half_t x) {
+    return (x.raw() == 0x8000 ? half_t::bitcast(0) : x);
+  }
+  static inline half_t to_print(half_t x) { return x; }
+  static inline device_type to_device(half_t x) { return reinterpret_cast<device_type const &>(x); }
+};
+
+template <>
+struct TypeTraits<float> {
+  static musaDataType_t const mublas_type = MUSA_R_32F;
+  typedef float host_type;
+  typedef float device_type;
+  typedef int32_t integer_type;
+  typedef uint32_t unsigned_type;
+  static inline float remove_negative_zero(float x) { return x == -0.f ? 0.f : x; }
+  static inline float to_print(float x) { return x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+template <>
+struct TypeTraits<double> {
+  static musaDataType_t const mublas_type = MUSA_R_64F;
+  typedef double host_type;
+  typedef double device_type;
+  typedef int64_t integer_type;
+  typedef uint64_t unsigned_type;
+  static inline double remove_negative_zero(double x) { return x == -0.0 ? 0.0 : x; }
+  static inline double to_print(double x) { return x; }
+  static inline device_type to_device(host_type x) { return x; }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// Complex types
+//
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <>
+struct TypeTraits<complex<half> > {
+  static musaDataType_t const mublas_type = MUSA_C_16F;
+  typedef complex<half_t> host_type;
+  typedef complex<half> device_type;
+  typedef int16_t integer_type;
+  typedef uint16_t unsigned_type;
+  static inline device_type to_device(complex<half> x) { return reinterpret_cast<device_type const &>(x); }
+};
+
+template <>
+struct TypeTraits<complex<half_t> > {
+  static musaDataType_t const mublas_type = MUSA_C_16F;
+  typedef complex<half_t> host_type;
+  typedef complex<half> device_type;
+  typedef int16_t integer_type;
+  typedef uint16_t unsigned_type;
+  static inline complex<half_t> remove_negative_zero(complex<half_t> x) {
+    return complex<half_t>(
+      real(x) == -0_hf ? 0_hf : real(x),
+      imag(x) == -0_hf ? 0_hf : imag(x)
+    );
+  }
+  static inline complex<half_t> to_print(complex<half_t> x) { return x; }
+  static inline device_type to_device(complex<half_t> x) { return reinterpret_cast<device_type const &>(x); }
+};
+
+template <>
+struct TypeTraits<complex<float> > {
+
+  static musaDataType_t const mublas_type = MUSA_C_32F;
+  typedef complex<float> host_type;
+  typedef complex<float> device_type;
+  typedef int64_t integer_type;
+  typedef uint64_t unsigned_type;
+
+  static inline complex<float> remove_negative_zero(complex<float> x) {
+    return complex<float>(
+      real(x) == -0.f ? 0.f : real(x),
+      imag(x) == -0.f ? 0.f : imag(x)
+    );
+  }
+
+  static inline complex<float> to_print(complex<float> x) { return x; }
+  static inline device_type to_device(complex<float> x) { return reinterpret_cast<device_type const &>(x); }
+};
+
+template <>
+struct TypeTraits<complex<double> > {
+  static musaDataType_t const mublas_type = MUSA_C_64F;
+  typedef complex<double> host_type;
+  typedef complex<double> device_type;
+  struct integer_type { int64_t real, imag; };
+  struct unsigned_type { uint64_t real, imag; };
+  static inline complex<double> remove_negative_zero(complex<double> x) {
+    return complex<double>(
+      real(x) == -0.0 ? 0.0 : real(x),
+      imag(x) == -0.0 ? 0.0 : imag(x)
+    );
+  }
+  static inline complex<double> to_print(complex<double> x) { return x; }
+  static inline device_type to_device(complex<double> x) { return reinterpret_cast<device_type const &>(x); }
+};
+
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/version.h b/vllm_musa/csrc/flash_attn/include/mutlass/version.h
new file mode 100644
index 000000000..b21444e29
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/version.h
@@ -0,0 +1,81 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2017 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+#pragma once
+
+#include <cstdint>
+#include <string>
+
+#define MUTLASS_MAJOR 0
+#define MUTLASS_MINOR 2
+#define MUTLASS_PATCH 0
+
+#ifdef MUTLASS_VERSIONS_GENERATED
+#include "mutlass/version_extended.h"
+#else
+#define MUTLASS_BUILD 0
+#define MUTLASS_REVISION ""
+#endif
+
+#define MUTLASS_VERSION ((MUTLASS_MAJOR)*100 + (MUTLASS_MINOR)*10 + MUTLASS_PATCH)
+
+namespace mutlass {
+
+  inline constexpr uint32_t getVersion() {
+    return MUTLASS_VERSION;
+  }
+  inline constexpr uint32_t getVersionMajor() {
+    return MUTLASS_MAJOR;
+  }
+  inline constexpr uint32_t getVersionMinor() {
+    return MUTLASS_MINOR;
+  }
+  inline constexpr uint32_t getVersionPatch() {
+    return MUTLASS_PATCH;
+  }
+  inline constexpr uint32_t getVersionBuild() {
+    return MUTLASS_BUILD + 0;
+  }
+
+  inline std::string getVersionString() {
+    std::string version = "@MUTLASS_VERSION@";
+    if (getVersionBuild()) {
+      version += "." + std::to_string(getVersionBuild());
+    }
+    return version;
+  }
+
+  inline std::string getGitRevision() {
+    return "@MUTLASS_REVISION@";
+  }
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/include/mutlass/workspace.h b/vllm_musa/csrc/flash_attn/include/mutlass/workspace.h
new file mode 100644
index 000000000..d25ef5dc8
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/include/mutlass/workspace.h
@@ -0,0 +1,139 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 Utilities for initializing workspaces
+*/
+
+#pragma once
+
+#if !defined(__MUSACC_RTC__)
+#include "musa.h"
+#include "musa_runtime.h"
+
+#include "mutlass/trace.h"
+#endif
+
+#include "mutlass.h"
+#include "mutlass/musa_host_adapter.hpp"
+namespace mutlass {
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+static constexpr int MinWorkspaceAlignment = 16;
+
+#if !defined(__MUSACC_RTC__)
+static Status
+zero_workspace(void* workspace, size_t workspace_size, musaStream_t stream = nullptr) {
+  if (workspace_size > 0) {
+    if (workspace == nullptr) {
+      MUTLASS_TRACE_HOST("  error: device workspace must not be null");
+      return Status::kErrorWorkspaceNull;
+    }
+
+    MUTLASS_TRACE_HOST("  clearing workspace");
+    musaError_t result = musaMemsetAsync(workspace, 0, workspace_size, stream);
+    if (musaSuccess != result) {
+      result = musaGetLastError(); // to clear the error bit
+      MUTLASS_TRACE_HOST("  musaMemsetAsync() returned error " << musaGetErrorString(result));
+      return Status::kErrorInternal;
+    }
+  }
+
+  return Status::kSuccess;
+}
+#endif
+
+#if !defined(__MUSACC_RTC__)
+template <typename T>
+Status
+fill_workspace(void* workspace, T fill_value, size_t fill_count, musaStream_t stream = nullptr, MusaHostAdapter *musa_adapter = nullptr) {
+  static_assert(sizeof(T) == 4 || sizeof(T) == 2 || sizeof(T) == 1, "Unsupported fill type");
+  if (fill_count > 0) {
+    if (workspace == nullptr) {
+      MUTLASS_TRACE_HOST("  error: device workspace must not be null");
+      return Status::kErrorWorkspaceNull;
+    }
+
+    MUTLASS_TRACE_HOST("  filling workspace");
+    MUdeviceptr d_workspace = reinterpret_cast<MUdeviceptr>(workspace);
+
+#if defined(MUTLASS_ENABLE_MUSA_HOST_ADAPTER) && MUTLASS_ENABLE_MUSA_HOST_ADAPTER
+
+    //
+    // Use the musa host adapter
+    //
+    MUTLASS_ASSERT(musa_adapter);
+    if (musa_adapter) {
+      Status status = Status::kErrorInternal;
+
+      status = musa_adapter->memsetDevice(workspace, fill_value, fill_count, stream);
+
+      if (status!=Status::kSuccess) {
+        return Status::kErrorInternal;
+      }
+    }
+    else {
+      return Status::kErrorInternal;
+    }
+#else
+    MUresult result = MUSA_SUCCESS;
+    if (sizeof(T) == 4) {
+      result = muMemsetD32Async(d_workspace, reinterpret_cast<uint32_t&>(fill_value), fill_count, stream);
+    }
+    else if (sizeof(T) == 2) {
+      result = muMemsetD16Async(d_workspace, reinterpret_cast<uint16_t&>(fill_value), fill_count, stream);
+    }
+    else if (sizeof(T) == 1) {
+      result = muMemsetD8Async(d_workspace, reinterpret_cast<uint8_t&>(fill_value), fill_count, stream);
+    }
+
+    if (MUSA_SUCCESS != result) {
+      const char** error_string_ptr = nullptr;
+      (void) muGetErrorString(result, error_string_ptr);
+      if (error_string_ptr != nullptr) {
+        MUTLASS_TRACE_HOST("  muMemsetD" << sizeof(T) * 8 << "Async() returned error " << *error_string_ptr);
+      }
+      else {
+        MUTLASS_TRACE_HOST("  muMemsetD" << sizeof(T) * 8 << "Async() returned unrecognized error");
+      }
+      return Status::kErrorInternal;
+    }
+#endif
+  }
+
+  return Status::kSuccess;
+}
+#endif
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace mutlass
diff --git a/vllm_musa/csrc/flash_attn/v1_flash_attn_fwd.mu b/vllm_musa/csrc/flash_attn/v1_flash_attn_fwd.mu
new file mode 100644
index 000000000..31421b920
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/v1_flash_attn_fwd.mu
@@ -0,0 +1,234 @@
+#include <torch/torch.h>
+#include <iostream>
+#include "v1_mp22_impl/attention_kernel.hpp"
+#include "v1_mp22_impl/kernel_traits.hpp"
+
+#include "../type_convert.muh"
+#include "../dispatch_utils.h"
+
+#define DISPATCH_HEAD_DIM_QK_VO(d_qk, d_vo, ...)                                                   \
+    TORCH_CHECK((d_qk) == (d_vo), "QK head dim (", d_qk, ") must match VO head dim (", d_vo, ")"); \
+    if ((d_qk) == 64)                                                                              \
+    {                                                                                              \
+        constexpr int HeadDimQK = 64;                                                              \
+        constexpr int HeadDimVO = 64;                                                              \
+        __VA_ARGS__();                                                                             \
+    }                                                                                              \
+    else if ((d_qk) == 128)                                                                        \
+    {                                                                                              \
+        constexpr int HeadDimQK = 128;                                                             \
+        constexpr int HeadDimVO = 128;                                                             \
+        __VA_ARGS__();                                                                             \
+    }                                                                                              \
+    else if ((d_qk) == 256)                                                                        \
+    {                                                                                              \
+        constexpr int HeadDimQK = 256;                                                             \
+        constexpr int HeadDimVO = 256;                                                             \
+        __VA_ARGS__();                                                                             \
+    }                                                                                              \
+    else                                                                                           \
+    {                                                                                              \
+        TORCH_CHECK(false, "Unsupported head dim: ", d_qk,                                         \
+                    ". Only 64, 128, 256 are supported.");                                         \
+    }
+
+template <typename T, bool return_lse, int HeadDimQK, int HeadDimVO>
+void varlen_flash_attn_fwd_kernel_launcher(
+    const T *Q, const T *K, const T *V,
+    T *out, T *P, float *Lse,
+    int q, int kv, int d_qk, int d_vo,
+    int h, int h_kv, int b, bool causal, mute::Shape<int, int, int> qkv_elem,
+    int q_seq_stride, int q_num_head_stride,
+    int k_num_page_stride, int k_page_size_stride, int k_num_head_stride,
+    int v_num_page_stride, int v_page_size_stride, int v_num_head_stride,
+    int o_seq_stride, int o_num_head_stride,
+    int lse_num_head_stride,
+    int *cu_q_seqlen, int *cu_kv_seqlen, int *block_table, int block_table_stride_batch, float sm_scale, int max_seqlen_q, int max_seqlen_kv)
+{
+    constexpr int TileM = 128;
+    constexpr int TileN = 32;
+    const musaStream_t stream = at::musa::getCurrentMUSAStream();
+    VariableLength varlen_q{max_seqlen_q, cu_q_seqlen};
+    VariableLength varlen_kv{max_seqlen_kv, cu_kv_seqlen};
+
+    int h_r = h / h_kv;
+    auto problem_shape = mute::make_shape(varlen_q, varlen_kv, d_qk, d_vo, mute::make_shape(h_r, h_kv), b);
+
+    using KernelTraits = AttentionForwardKernelTraits<HeadDimQK, HeadDimVO, TileM, TileN,
+                                                      true, true, true, true, true, T, return_lse>;
+
+    using AttentionParams = AttentionParams<KernelTraits>;
+    using Arguments = typename AttentionParams::Arguments;
+    Arguments arguments = {Q, K, V, out, Lse, P};
+
+    auto params = AttentionParams::to_underlying_arguments(problem_shape, arguments, qkv_elem,
+                                                           q_seq_stride, q_num_head_stride,
+                                                           k_num_page_stride, k_page_size_stride, k_num_head_stride,
+                                                           v_num_page_stride, v_page_size_stride, v_num_head_stride,
+                                                           o_seq_stride, o_num_head_stride,
+                                                           lse_num_head_stride,
+                                                           block_table, block_table_stride_batch, sm_scale);
+
+    dim3 block_dims(KernelTraits::NumThreads);
+    dim3 grid_dims(mute::ceil_div(mute::size<0>(problem_shape), KernelTraits::CTA_Q),
+                   mute::size<4>(problem_shape),
+                   mute::size<5>(problem_shape));
+
+    // musaDeviceSynchronize();
+    // auto error = musaGetLastError();
+    // if (error != musaSuccess)
+    // {
+    //     std::cerr << "Error before kernel launch: " << musaGetErrorString(error) << std::endl;
+    //     return;
+    // }
+
+    if (causal)
+    {
+        mp22_varlen_flash_atten_forward<KernelTraits, true, true><<<grid_dims, block_dims, 0, stream>>>(params);
+    }
+    else
+    {
+        mp22_varlen_flash_atten_forward<KernelTraits, true, false><<<grid_dims, block_dims, 0, stream>>>(params);
+    }
+
+    // musaDeviceSynchronize();
+    // error = musaGetLastError();
+    // if (error != musaSuccess)
+    // {
+    //     std::cerr << "Kernel execution error: " << musaGetErrorString(error) << std::endl;
+    // }
+}
+
+void mp22_flash_attn_varlen_func(torch::Tensor &Q,                  // [seqs_len, num_heads, head_size]
+                                 torch::Tensor &K,                  // [num_block, page_size, num_heads, head_size]
+                                 torch::Tensor &V,                  // [num_block, page_size, num_heads, head_size]
+                                 torch::Tensor &out,                // [seqs_len, num_heads, head_size]
+                                 c10::optional<torch::Tensor> &lse, // [ num_heads, seqs_len]
+                                 torch::Tensor &cu_q_seqlen,        // such as [0, 19, 44]
+                                 torch::Tensor &cu_kv_seqlen,
+                                 torch::Tensor &block_table,
+                                 float sm_scale,
+                                 int max_seqlen_q,
+                                 int max_seqlen_kv,
+                                 bool causal,
+                                 bool return_softmax_lse)
+{
+    TORCH_CHECK(Q.scalar_type() == at::ScalarType::BFloat16 ||
+                    Q.scalar_type() == at::ScalarType::Half,
+                "Unsupported data type: ", Q.scalar_type());
+
+    int current_arch = at::musa::getMUSAArch();
+    TORCH_CHECK(current_arch == 220, "flash_attn_varlen_fun not support on MUSA arch ", current_arch);
+
+    int q_elem = Q.storage().nbytes() / Q.element_size();
+    int k_elem = K.storage().nbytes() / K.element_size();
+    int v_elem = V.storage().nbytes() / V.element_size();
+    auto qkv_elem = mute::make_shape(q_elem, k_elem, v_elem);
+
+    int batch = cu_q_seqlen.size(0) - 1;
+    int num_head = Q.size(1);
+    int num_head_kv = K.size(2);
+    int d_qk = Q.size(2);
+    int d_vo = V.size(3);
+    void *P = nullptr;
+
+    int page_size = K.size(1);
+    TORCH_CHECK(page_size == 32, "mp22 varlen flash attn BLOCK_SIZE only support 32.");
+
+    float *lse_ptr = nullptr;
+    int lse_num_head_stride = 1;
+    if (return_softmax_lse)
+    {
+        TORCH_CHECK(lse.has_value(), "lse tensor is required");
+        lse_ptr = static_cast<float *>(lse->data_ptr());
+        lse_num_head_stride = lse->stride(0);
+    }
+
+    DISPATCH_HEAD_DIM_QK_VO(d_qk, d_vo, [&]
+                            {
+        if (return_softmax_lse) {
+            if (Q.scalar_type() == at::ScalarType::BFloat16) {
+                varlen_flash_attn_fwd_kernel_launcher<mutlass::bfloat16_t, true, HeadDimQK, HeadDimVO>(
+                    static_cast<mutlass::bfloat16_t *>(Q.data_ptr()),
+                    static_cast<mutlass::bfloat16_t *>(K.data_ptr()),
+                    static_cast<mutlass::bfloat16_t *>(V.data_ptr()),
+                    static_cast<mutlass::bfloat16_t *>(out.data_ptr()),
+                    static_cast<mutlass::bfloat16_t *>(P),
+                    lse_ptr,
+                    max_seqlen_q, max_seqlen_kv, d_qk, d_vo,
+                    num_head, num_head_kv, batch, causal, qkv_elem,
+                    Q.stride(0), Q.stride(1),
+                    K.stride(0), K.stride(1), K.stride(2),
+                    V.stride(0), V.stride(1), V.stride(2),
+                    out.stride(0), out.stride(1),
+                    lse_num_head_stride,
+                    static_cast<int *>(cu_q_seqlen.data_ptr()),
+                    static_cast<int *>(cu_kv_seqlen.data_ptr()),
+                    static_cast<int *>(block_table.data_ptr()),
+                    block_table.stride(0),
+                    sm_scale, max_seqlen_q, max_seqlen_kv);
+            } else { // Half
+                varlen_flash_attn_fwd_kernel_launcher<mutlass::half_t, true, HeadDimQK, HeadDimVO>(
+                    static_cast<mutlass::half_t *>(Q.data_ptr()),
+                    static_cast<mutlass::half_t *>(K.data_ptr()),
+                    static_cast<mutlass::half_t *>(V.data_ptr()),
+                    static_cast<mutlass::half_t *>(out.data_ptr()),
+                    static_cast<mutlass::half_t *>(P),
+                    lse_ptr,
+                    max_seqlen_q, max_seqlen_kv, d_qk, d_vo,
+                    num_head, num_head_kv, batch, causal, qkv_elem,
+                    Q.stride(0), Q.stride(1),
+                    K.stride(0), K.stride(1), K.stride(2),
+                    V.stride(0), V.stride(1), V.stride(2),
+                    out.stride(0), out.stride(1),
+                    lse_num_head_stride,
+                    static_cast<int *>(cu_q_seqlen.data_ptr()),
+                    static_cast<int *>(cu_kv_seqlen.data_ptr()),
+                    static_cast<int *>(block_table.data_ptr()),
+                    block_table.stride(0),
+                    sm_scale, max_seqlen_q, max_seqlen_kv);
+            }
+        } else {
+            if (Q.scalar_type() == at::ScalarType::BFloat16) {
+                varlen_flash_attn_fwd_kernel_launcher<mutlass::bfloat16_t, false, HeadDimQK, HeadDimVO>(
+                    static_cast<mutlass::bfloat16_t *>(Q.data_ptr()),
+                    static_cast<mutlass::bfloat16_t *>(K.data_ptr()),
+                    static_cast<mutlass::bfloat16_t *>(V.data_ptr()),
+                    static_cast<mutlass::bfloat16_t *>(out.data_ptr()),
+                    static_cast<mutlass::bfloat16_t *>(P),
+                    lse_ptr,
+                    max_seqlen_q, max_seqlen_kv, d_qk, d_vo,
+                    num_head, num_head_kv, batch, causal, qkv_elem,
+                    Q.stride(0), Q.stride(1),
+                    K.stride(0), K.stride(1), K.stride(2),
+                    V.stride(0), V.stride(1), V.stride(2),
+                    out.stride(0), out.stride(1),
+                    lse_num_head_stride,
+                    static_cast<int *>(cu_q_seqlen.data_ptr()),
+                    static_cast<int *>(cu_kv_seqlen.data_ptr()),
+                    static_cast<int *>(block_table.data_ptr()),
+                    block_table.stride(0),
+                    sm_scale, max_seqlen_q, max_seqlen_kv);
+            } else { // Half
+                varlen_flash_attn_fwd_kernel_launcher<mutlass::half_t, false, HeadDimQK, HeadDimVO>(
+                    static_cast<mutlass::half_t *>(Q.data_ptr()),
+                    static_cast<mutlass::half_t *>(K.data_ptr()),
+                    static_cast<mutlass::half_t *>(V.data_ptr()),
+                    static_cast<mutlass::half_t *>(out.data_ptr()),
+                    static_cast<mutlass::half_t *>(P),
+                    lse_ptr,
+                    max_seqlen_q, max_seqlen_kv, d_qk, d_vo,
+                    num_head, num_head_kv, batch, causal, qkv_elem,
+                    Q.stride(0), Q.stride(1),
+                    K.stride(0), K.stride(1), K.stride(2),
+                    V.stride(0), V.stride(1), V.stride(2),
+                    out.stride(0), out.stride(1),
+                    lse_num_head_stride,
+                    static_cast<int *>(cu_q_seqlen.data_ptr()),
+                    static_cast<int *>(cu_kv_seqlen.data_ptr()),
+                    static_cast<int *>(block_table.data_ptr()),
+                    block_table.stride(0),
+                    sm_scale, max_seqlen_q, max_seqlen_kv);
+            }
+        } });
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/flash_attn/v1_mp22_impl/attention_kernel.hpp b/vllm_musa/csrc/flash_attn/v1_mp22_impl/attention_kernel.hpp
new file mode 100644
index 000000000..3ba9af41a
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/v1_mp22_impl/attention_kernel.hpp
@@ -0,0 +1,608 @@
+#pragma once
+
+#include <mutlass/numeric_conversion.h>
+
+#include "attention_params.hpp"
+#include "../impl/utils.hpp"
+#include "../impl/online_softmax.hpp"
+
+using namespace mute;
+
+
+template <
+  class KernelTraits,
+  bool EnableSoftmax,
+  bool EnableCausal
+>
+__global__ __launch_bounds__(KernelTraits::NumThreads, 1)
+void mp22_varlen_flash_atten_forward(const typename AttentionParams<KernelTraits>::Params params) {
+  #if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 220)
+  // printf("liliang in kernel \n");
+  static_assert(KernelTraits::IsQInReg, "Must keep Q in registers now.");
+  using ElementL = typename KernelTraits::ElementLse;
+
+  using ElementAccum = typename KernelTraits::ElementAccumulator;
+  constexpr int FragmentSize = KernelTraits::FragmentSize;
+
+  __shared__ char __attribute__((aligned(128))) smem[sizeof(typename KernelTraits::SharedStorage)];
+  auto& shared_storage = *reinterpret_cast<typename KernelTraits::SharedStorage*>(smem);
+
+  auto [coord_q, coord_h, coord_b] = static_cast<uint3>(blockIdx);
+
+  int thread_idx = threadIdx.x;
+
+  auto tile_shape_Q = select<0, 2>(typename KernelTraits::TileShape_QKD{});
+  auto tile_shape_K = select<1, 2>(typename KernelTraits::TileShape_QKD{});
+  auto tile_shape_V = select<2, 1>(typename KernelTraits::TileShape_PDV{});
+  auto tile_shape_O = select<0, 1>(typename KernelTraits::TileShape_PDV{});
+  auto tile_shape_P = select<0, 1>(typename KernelTraits::TileShape_QKD{});
+
+  // int actual_seqlen_q = 10;
+  int actual_seqlen_q = get<0>(params.problem_shape).cumulative_length[coord_b + 1] - get<0>(params.problem_shape).cumulative_length[coord_b];
+
+
+  // if (threadIdx.x == 0) {
+  //   printf("liliang0 coord_q = %d, actual_seqlen_q = %d block_table = %p\n", coord_q, actual_seqlen_q, params.block_table);
+  // }
+  if (coord_q * KernelTraits::CTA_Q >= actual_seqlen_q) return;
+  // int page_offset = 0;
+  int page_offset = params.block_table[coord_b * params.block_table_stride_batch];
+  int num_heads_kv = get<4, 1>(params.problem_shape);
+  int cur_seqlen_kv = get<1>(params.problem_shape).cumulative_length[coord_b];
+  int q_ptr_offset = get<0>(params.problem_shape).cumulative_length[coord_b] * get<0>(params.stride_Q);
+  int l_ptr_offset = get<0>(params.problem_shape).cumulative_length[coord_b] * get<0>(params.stride_L);
+  int k_num_head_offset = (coord_h / mute::size<4,0>(params.problem_shape)) * get<2, 1>(params.stride_K);
+  int v_num_head_offset = (coord_h / mute::size<4,0>(params.problem_shape)) * get<2, 1>(params.stride_V);
+  int64_t k_ptr_offset =  page_offset* get<3>(params.stride_K) + k_num_head_offset; //TODO
+  int64_t v_ptr_offset = page_offset * get<3>(params.stride_V) + v_num_head_offset;
+  int o_ptr_offset = get<0>(params.problem_shape).cumulative_length[coord_b] * get<0>(params.stride_O);
+  // if (threadIdx.x == 0) {
+  //   printf("liliang1 coord_h = %d, num_heads_kv = %d block_table = %p\n", coord_h, num_heads_kv, params.block_table);
+  // }
+
+
+  // (actual_seqlen_q, D_QK, (H_R, H_KV))
+  auto shape_Q = make_shape(actual_seqlen_q, get<2>(params.problem_shape), get<4>(params.problem_shape));
+  auto Shape_L = make_shape(actual_seqlen_q, mute::_1{}, get<4>(params.problem_shape));
+  // (actual_seqlen_k, D_QK, (H_R, H_KV))
+  // (BlockN, D_QK, (H_R, H_KV))
+  auto shape_K = make_shape(select<1>(typename KernelTraits::TileShape_QKD{}), get<2>(params.problem_shape), get<4>(params.problem_shape));
+  // (actual_seqlen_k, D_VO, (H_R, H_KV))
+  // (BlockN, D_VO, (H_R, H_KV))
+  auto shape_V = make_shape(select<1>(typename KernelTraits::TileShape_QKD{}), get<3>(params.problem_shape), get<4>(params.problem_shape));
+  // (actual_seqlen_q, D_VO, (H_R, H_KV))
+  auto shape_O = make_shape(actual_seqlen_q, get<3>(params.problem_shape), get<4>(params.problem_shape));
+
+  Tensor mQ = make_tensor(make_gmem_ptr(params.ptr_Q + q_ptr_offset), shape_Q, select<0, 1, 2>(params.stride_Q));
+  Tensor mO = make_tensor(make_gmem_ptr(params.ptr_O + o_ptr_offset), shape_O, select<0, 1, 2>(params.stride_O));
+  Tensor mL = make_tensor(make_gmem_ptr(params.ptr_L + l_ptr_offset), Shape_L, params.stride_L);
+
+  Tensor gQ = local_tile(mQ, tile_shape_Q, make_coord(coord_q, _0{}, coord_h));   // (CTA_Q, HEAD_DIM_QK)
+  Tensor gK = make_tensor(make_gmem_ptr(params.ptr_K + k_ptr_offset), tile_shape_K, select<0, 1>(params.stride_K));
+  Tensor gV = make_tensor(make_gmem_ptr(params.ptr_V + v_ptr_offset), tile_shape_V, select<0, 1>(params.stride_V));
+  Tensor gO = local_tile(mO, tile_shape_O, make_coord(coord_q, _0{}, coord_h));   // (CTA_Q, CTA_KV, kv)
+  Tensor gL = local_tile(mL, tile_shape_O, make_coord(coord_q, _0{}, coord_h));
+  // if (threadIdx.x == 0) {
+  //   printf("liliang coord_q = %d, coord_h = %d \n", coord_q, coord_h); print(gO); print("\n");
+  // }
+
+  Tensor sQ  = make_tensor(make_smem_ptr(shared_storage.smem_q.data()), typename KernelTraits::SmemLayoutQ{});
+  Tensor sK  = make_tensor(make_smem_ptr(shared_storage.smem_k.data()), typename KernelTraits::SmemLayoutK{});
+  Tensor sV  = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), typename KernelTraits::SmemLayoutV{});
+  Tensor sVt = make_tensor(make_smem_ptr(shared_storage.smem_v.data()), typename KernelTraits::SmemLayoutVt{});
+  Tensor sP  = make_tensor(make_smem_ptr(shared_storage.smem_p.data()), typename KernelTraits::SmemLayoutP{});
+
+  typename KernelTraits::GmemTiledCopyQ gmem_tiled_copy_Q;
+  typename KernelTraits::GmemTiledCopyK gmem_tiled_copy_K;
+  typename KernelTraits::GmemTiledCopyV gmem_tiled_copy_V;
+
+  ThrCopy gmem_thr_copy_Q = gmem_tiled_copy_Q.get_thread_slice(thread_idx);
+  ThrCopy gmem_thr_copy_K = gmem_tiled_copy_K.get_thread_slice(thread_idx);
+  ThrCopy gmem_thr_copy_V = gmem_tiled_copy_V.get_thread_slice(thread_idx);
+
+  Tensor tQgQ = gmem_thr_copy_Q.partition_S(gQ); // (QCPY, QCPY_M, QCPY_N)
+  Tensor tQsQ = gmem_thr_copy_Q.partition_D(sQ); // (QCPY, QCPY_M, QCPY_N)
+  Tensor tQrQ = make_fragment_like(tQsQ);
+
+  Tensor tKgK = gmem_thr_copy_K.partition_S(gK); // (KCPY, KCPY_M, KCPY_N)
+  Tensor tKsK = gmem_thr_copy_K.partition_D(sK); // (KCPY, KCPY_M, KCPY_N)
+  Tensor tKrK = make_fragment_like(tKsK);
+
+  Tensor tVgV = gmem_thr_copy_V.partition_S(gV); // (VCPY, VCPY_M, VCPY_N)
+  Tensor tVsV = gmem_thr_copy_V.partition_D(sV); // (VCPY, VCPY_M, VCPY_N)
+  Tensor tVrV = make_fragment_like(tVsV);
+
+  typename KernelTraits::TiledMmaQK tiled_mma_qk;
+  typename KernelTraits::TiledMmaPV tiled_mma_pv;
+
+  auto thr_mma_qk = tiled_mma_qk.get_thread_slice(thread_idx);
+  auto thr_mma_pv = tiled_mma_pv.get_thread_slice(thread_idx);
+
+  auto gmem_tiled_copy_L = make_tiled_copy_C_atom(
+    Copy_Atom<UniversalCopy<uint_bit_t<sizeof_bits_v<ElementL>>>, ElementL>{},
+    tiled_mma_qk
+  );
+  ThrCopy gmem_thr_copy_L = gmem_tiled_copy_L.get_thread_slice(thread_idx);
+  Tensor tSgL = gmem_thr_copy_L.partition_D(gL);
+
+
+  auto temp_shape = mute::Shape<Int<size<0>(tile_shape_P)>, _1>{};
+  auto temp_layout = make_layout(make_shape(size<0>(tile_shape_P), 1), Stride<_1, _0>{});
+  Tensor cL   = make_identity_tensor(temp_shape);
+  Tensor tCcL = thr_mma_qk.partition_C(cL.compose(temp_layout));
+
+  // QK MMA AB
+  Tensor tPrQ = thr_mma_qk.partition_fragment_A(sQ); // (MMA, MMA_M, MMA_K)
+  Tensor tPrK = thr_mma_qk.partition_fragment_B(sK); // (MMA, MMA_M, MMA_K)
+
+  // PV MMA AB
+  Tensor tOrP  = thr_mma_pv.partition_fragment_A(sP);
+  Tensor tOrVt = thr_mma_pv.partition_fragment_B(sVt);
+
+  // Allocate accum
+  Tensor accum_p = partition_fragment_C(tiled_mma_qk, tile_shape_P);
+  Tensor accum_o = partition_fragment_C(tiled_mma_pv, tile_shape_O);
+
+  // Allocate softmax
+  Tensor accum_softmax = make_fragment_like(accum_p);
+  Tensor accum_cvt     = make_fragment_like<typename KernelTraits::ElementKV>(accum_softmax);
+
+  Tensor tAcc_frg = recast<mutlass::Array<ElementAccum, FragmentSize>>(accum_softmax);
+  Tensor tCvt_frg = recast<mutlass::Array<typename KernelTraits::ElementKV, FragmentSize>>(accum_cvt);
+
+  // STS P
+  auto tiled_copy_r2s = make_tiled_copy_C(typename KernelTraits::R2SCopyAtomP{}, tiled_mma_qk);
+  auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(thread_idx);
+  Tensor tPrP = thr_copy_r2s.retile_S(accum_cvt);
+  Tensor tPsP = thr_copy_r2s.partition_D(sP);
+
+  // LDS QKPV
+  auto thr_copy_Q = make_tiled_copy_A(typename KernelTraits::SmemCopyAtomQ{}, tiled_mma_qk).get_thread_slice(thread_idx);
+  auto thr_copy_K = make_tiled_copy_B(typename KernelTraits::SmemCopyAtomK{}, tiled_mma_qk).get_thread_slice(thread_idx);
+  auto thr_copy_P = make_tiled_copy_A(typename KernelTraits::SmemCopyAtomP{}, tiled_mma_pv).get_thread_slice(thread_idx);
+  auto thr_copy_V = make_tiled_copy_B(typename KernelTraits::SmemCopyAtomV{}, tiled_mma_pv).get_thread_slice(thread_idx);
+
+  // QK GEMM
+  Tensor tPsQ           = thr_copy_Q.partition_S(sQ);
+  Tensor tPrQ_copy_view = thr_copy_Q.retile_D(tPrQ);
+  MUTE_STATIC_ASSERT_V(size<1>(tPsQ) == size<1>(tPrQ_copy_view));
+
+  Tensor tPsK           = thr_copy_K.partition_S(sK);
+  Tensor tPrK_copy_view = thr_copy_K.retile_D(tPrK);
+  MUTE_STATIC_ASSERT_V(size<1>(tPsK) == size<1>(tPrK_copy_view));
+
+  // Softmax
+  constexpr int RowsPerThread = size<0>(tile_shape_Q) / tile_size<0>(tiled_mma_qk) * 4;
+  Tensor alpha       = make_tensor<ElementAccum>(Shape<Int<RowsPerThread>>{});
+  Tensor lse         = make_fragment_like(alpha);
+  Tensor row_sum     = make_fragment_like(alpha);
+  Tensor row_max     = make_fragment_like(alpha);
+  Tensor old_row_max = make_fragment_like(alpha);
+
+  if constexpr (EnableSoftmax) {
+    fill(lse,         ElementAccum(0.0f));
+    fill(alpha,       ElementAccum(0.0f));
+    fill(old_row_max, -mutlass::platform::numeric_limits<ElementAccum>::infinity());
+  } else {
+    fill(lse, ElementAccum(1.0f));
+  }
+
+  // PV GEMM
+  Tensor tOsP           = thr_copy_P.partition_S(sP);
+  Tensor tOrP_copy_view = thr_copy_P.retile_D(tOrP);
+
+  Tensor tOsVt           = thr_copy_V.partition_S(sVt);
+  Tensor tOrV_copy_view  = thr_copy_V.retile_D(tOrVt);
+
+  // clear(accum_p);
+  clear(accum_o);
+
+  //
+  // PREDICATES
+  //
+
+  // Allocate predicate tensors for Q/K/V
+  auto residue_q = actual_seqlen_q - size<0>(gQ) * coord_q;
+  auto residue_kv = (cur_seqlen_kv % size<0>(gK) == 0) ? size<0>(gK) : cur_seqlen_kv % size<0>(gK);
+
+  // Deal with kv oob
+  Tensor cP = make_identity_tensor(tile_shape_P);
+  Tensor tPcP = thr_mma_qk.partition_C(cP);
+  Tensor tPcP_view = make_tensor(tPcP.data(), convert_layout_acc_rowcol(tPcP.layout()));
+
+  // Output Tensor
+  Tensor tOgO = thr_mma_pv.partition_C(gO);
+  Tensor cO   = make_identity_tensor(tile_shape_O);
+  Tensor tCcO = thr_mma_pv.partition_C(cO);
+  // Assume head dim match
+  auto oob_mn = mute::make_tuple(residue_q, size<1>(tile_shape_O));
+
+
+  Tensor cQ = make_identity_tensor(make_shape(size<0>(sQ), size<1>(sQ)));
+  Tensor tQcQ = gmem_thr_copy_Q.partition_S(cQ);
+  Tensor tQpQ = make_tensor<bool>(make_shape(size<1>(tQsQ), size<2>(tQsQ)), Stride<_1,_0>{}); // let's only check Head
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < size<0>(tQpQ); ++i) {
+    // int cur_head = cur_head_id * ValidBlockH + get<0>(tQcQ(0,i,0));
+    // tQpQ(i, 0) = (cur_head < (cur_head_id + 1) * ValidBlockH) && (cur_head < QHeadNum);
+    int cur_seq_q = coord_q * KernelTraits::CTA_Q + get<0>(tQcQ(0,i,0));
+    tQpQ(i, 0) = (cur_seq_q < (coord_q + 1) * KernelTraits::CTA_Q) && (cur_seq_q < actual_seqlen_q);
+  }
+
+  Tensor cK = make_identity_tensor(make_shape(size<0>(sK), size<1>(sK)));
+  Tensor tKcK = gmem_thr_copy_K.partition_S(cK);
+  Tensor tKpK = make_tensor<bool>(make_shape(size<1>(tKcK), size<2>(tKcK)), Stride<_1, _0>{}); // let's only check Seq
+  MUTLASS_PRAGMA_UNROLL
+  for (int i = 0; i < size<0>(tKpK); ++i) {
+    tKpK(i, 0) = get<0>(tKcK(0,i,0)) < cur_seqlen_kv;
+  }
+  // if (blockIdx.x == 0) {
+  //   for (int i = 0; i < size<0>(tKpK); ++i) {
+  //     printf("liliang threadIdx.x = %d, tKpK(%d, 0) = %d\n", threadIdx.x, i, int(tKpK(i, 0)));
+  //   }
+  // }
+
+  // Debug save
+  // Tensor tPgP = thr_copy_r2s.partition_D(gP); // (CPY, CPY_M, CPY_N, kv)
+
+  if (thread0()) {
+    //print("row_sum:"); print(row_sum); print("\n");
+    // print("cK:"); print(cK); print("\n");
+    // print("tKcK:"); print(tKcK); print("\n");
+    // print("tKpK:"); print(tKpK); print("\n");
+    // print("gK:"); print(gK); print("\n");
+    //print("mV:"); print(mV); print("\n");
+    //print("gV:"); print(gV); print("\n");
+    //print("gP:"); print(gP); print("\n");
+    //print("tPgP:"); print(tPgP); print("\n");
+    //
+    //print("tPcS"); print(tPcS); print("\n");
+    //
+    //print("shapeK:"); print(shape_K); print("\n");
+    //print("residue_kv:"); print(residue_kv); print("\n");
+    //
+    //print("actual seqlen q:"); print(actual_seqlen_q); print("\n");
+    //print("actual seqlen kv:"); print(actual_seqlen_kv); print("\n");
+
+    //print("tPcP_view"); print(tPcP_view); print("\n");
+    //print("v:"); print(tPcP_view(0, 1)); print("\n");
+    //print("v:"); print(tPcP_view(0, 2)); print("\n");
+    // print("sQ:"); print(sQ); print("\n");
+    // print("sK:"); print(sK); print("\n");
+    //print("sV:"); print(sV); print("\n");
+    //print("sP:"); print(sP); print("\n");
+    //print("tQgQ:"); print(tQgQ); print("\n");
+    //print("tQsQ:"); print(tQsQ); print("\n");
+    // print("gmem_tiled_copy_K:"); print(gmem_tiled_copy_K); print("\n");
+    // print("tKgK:"); print(tKgK); print("\n");
+    // print("tKsK:"); print(tKsK); print("\n");
+    // print("tKrK:"); print(tKrK); print("\n");
+    // print("accum_p:"); print(accum_p); print("\n");
+    //print("accum_o:"); print(accum_o); print("\n");
+    //print("tPrQ:"); print(tPrQ); print("\n");
+    // print("tPrK:"); print(tPrK); print("\n");
+    // print("tPsK:"); print(tPsK); print("\n");
+    // print("tPsQ:"); print(tPsQ); print("\n");
+    // print("tPrQ_copy_view:"); print(tPrQ_copy_view); print("\n");
+    // print("tPrQ:"); print(tPrQ); print("\n");
+  }
+
+  //if (threadIdx.x == 0) {
+  //  print("residue_kv:"); print(residue_kv); print("\n");
+  //}
+
+  //
+  // Prologue
+  //
+
+  clear(tKrK);
+  clear(tVrV);
+
+  // LDG Q
+  if constexpr (KernelTraits::UseRobustLoadQ) {
+    copy_if(gmem_tiled_copy_Q.with(params.desc_Q), tQpQ, tQgQ, tQrQ);
+  } else {
+    copy_if(gmem_tiled_copy_Q, tQpQ, tQgQ, tQrQ);
+  }
+  // STS Q
+  // if (thread(32) || thread(64)) {
+  //   printf("liliang tQgQ: "); print(tQgQ(0)); print("\n");
+  //   printf("liliang tQrQ: "); print(tQrQ(0)); print("\n");
+  // }
+  // if (threadIdx.x == 0) {
+  //   printf("liliang coord_h  = %d ", coord_h);print(tQrQ(0));print("\n");
+  // }
+  copy(tQrQ, tQsQ);
+
+  __syncthreads();
+
+  // LDS Q, keep Q in register
+  copy(tPsQ, tPrQ_copy_view);
+
+  // auto k_tile_iter = mute::make_coord_iterator(shape<2>(gK));
+  // auto v_tile_iter = k_tile_iter;
+  int sliding_window = 0;
+  int mask_m_offset = cur_seqlen_kv - actual_seqlen_q + sliding_window ;
+
+  int orig_kv_tile_count = mute::ceil_div(cur_seqlen_kv, size<0>(tile_shape_K));
+  int kv_tile_count = orig_kv_tile_count;
+
+  if constexpr (EnableCausal) {
+    kv_tile_count = std::min(orig_kv_tile_count,
+                        int(mute::ceil_div((coord_q + 1) * size<0>(tile_shape_Q) + mask_m_offset + 1, size<0>(gK))));
+
+  }
+  //   if (thread(0)) {
+  //   printf("liliang cur_seqlen_kv = %d, kv_tile_count = %d, orig_kv_tile_count = %d mask_m_offset = %d\n",
+  //          cur_seqlen_kv, kv_tile_count, orig_kv_tile_count, mask_m_offset);
+  // }
+
+  // LDG K
+  if constexpr (KernelTraits::UseRobustLoadK) {
+    copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK, tKgK, tKrK);
+  } else {
+    copy_if(gmem_tiled_copy_K, tKpK, tKgK, tKrK);
+  }
+  // ++k_tile_iter;
+
+  if constexpr (KernelTraits::ReuseSmemQ) {
+    __syncthreads();
+  }
+  // if (threadIdx.x == 0) {
+  //   printf("liliang coord_h  = %d ", coord_h);print(tKrK(0));print("\n");
+  // }
+
+  copy(tKrK, tKsK);
+  __syncthreads();
+
+      //   #pragma unroll
+      // for (int i = 0; i < size(tPrQ); ++i) {
+      //   if (tPrQ(i) != 1.0f) {
+      //     printf("liliang0 blockIdx = %d, %d, %d, threadIdx.x = %d i = %d, tPrQ: ", blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, i); print(tPrQ(i)); print("\n");
+      //   }
+      // }
+
+  // P = QK
+  // asm("CTRL.NOP");
+  clear(accum_p);
+  // fill(accum_p, ElementAccum(0.0f)); // Initialize accum_p to 128.0f
+
+  // asm("FOP.MOV %0, C0" : "=r"(accum_p));
+  // #pragma unroll
+  // for (int i = 0; i < size(accum_p); ++i) {
+  //   accum_p(i) = 0;
+  //   if (accum_p(i) != 0.0f) {
+  //     printf("liliang2 blockIdx = %d, %d, %d, threadIdx.x = %d i = %d, accum_p: ", blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, i); print(accum_p(i)); print("\n");
+  //   }
+  // }
+
+  gemm<KernelTraits::IsQInReg>(accum_p, tPrQ, tPrK, tPsQ, tPsK, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+  // asm("CTRL.NOP");
+      // #pragma unroll
+      // for (int i = 0; i < size(accum_p); ++i) {
+      //   if (accum_p(i) != 128.0f) {
+      //     printf("liliang0 blockIdx = %d, %d, %d, threadIdx.x = %d i = %d, accum_p: ", blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, i); print(accum_p(i)); print("\n");
+      //   }
+      // }
+      // if (thread(0)) {
+      //   printf("liliang tPrQ: "); print(tPrQ(0)); print("\n");
+      //   auto temp = recast<uint16_t>(tPrQ);
+      //   // auto byte_ptr = static_cast<const uint16_t>(tPrQ);
+      //   // a
+      //       #pragma unroll
+      //       for (size_t i = 0; i < 16; ++i) {
+      //           printf("temp = %d ", (temp[0] >> i) & 1);
+      //       }
+      //       printf("\n");
+      // }
+
+    // if (threadIdx.x < 4 || (threadIdx.x >=64 && threadIdx.x < 68)) {
+    //   printf("liliang1 thread = %d tPrQ: ", threadIdx.x); print(tPrQ(0));print(" ");print(tPrQ(1)); print("\n");
+    // }
+    // if (threadIdx.x < 1) {
+    //   printf("liliang1 thread = %d accum_p: ", threadIdx.x); print(accum_p(0));print(" ");print(accum_p(1)); print("\n");
+    // }
+  // Apply Causal Mask
+  if constexpr (EnableCausal) {
+    apply_causal_mask(accum_p, tPcP, coord_q * size<0>(gQ) + mask_m_offset, 0, false);
+  }
+      if(thread(0)) {
+      // print("liliang thr_mma_qk: "); print(thr_mma_qk); print("\n");
+      // print("liliang accum_p: "); print(accum_p); print("\n");
+      // print("liliang cP: "); print(cP); print("\n");
+      // print("liliang tPcP: "); print(tPcP); print("\n");
+      // print("liliang gV: "); print(gV); print("\n");
+      // print("liliang v_ptr_offset: "); print(v_ptr_offset); print("\n");
+      // print("liliang tVgV: "); print(tVgV); print("\n");
+    }
+  // if (blockIdx.x == 0 && threadIdx.x < 8) {
+    // #pragma unroll
+    // for (int i = 0; i < 8; ++i) {
+    //   printf("liliang  threadIdx.x = %d accum_p(%d) = %f\n", threadIdx.x, i, accum_p(i));
+    // }
+  // }
+
+  // Softmax(P)
+  if constexpr (EnableSoftmax) {
+    online_softmax_scale(accum_p, accum_softmax, params.rln2_scale, tPcP, residue_kv, orig_kv_tile_count == 1);
+    //online_softmax_scale(accum_p, accum_softmax, params.rln2_scale);
+    online_softmax_row_max<ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+  } else {
+    online_softmax_scale(accum_p, accum_softmax, float(1.0f));
+  }
+  int page_size = 32; //TODO:
+  for (int loop = 1; loop < kv_tile_count; ++loop) {
+    int page_idx = loop * get<0>(tile_shape_K) / page_size;
+    int page_size_idx = (loop * get<0>(tile_shape_K) % page_size);
+    k_ptr_offset = params.block_table[coord_b * params.block_table_stride_batch + page_idx] * get<3>(params.stride_K) + page_size_idx * get<0>(params.stride_K) + k_num_head_offset;
+    gK = make_tensor(make_gmem_ptr(params.ptr_K + k_ptr_offset), tile_shape_K, select<0, 1>(params.stride_K));
+    tKgK = gmem_thr_copy_K.partition_S(gK);
+
+    // LDG V
+    if constexpr (KernelTraits::UseRobustLoadV) {
+      copy_if(gmem_tiled_copy_V.with(params.desc_V), tKpK, tVgV, tVrV);
+    } else {
+      copy_if(gmem_tiled_copy_V, tKpK, tVgV, tVrV);
+    }
+
+    // if(thread(0)) {
+    //   print("liliang params.ptr_V: "); print(params.ptr_V); print("\n");
+    //   print("liliang tVgV: "); print(tVgV(0)); print("\n");
+    //   print("liliang gV: "); print(gV); print("\n");
+    //   print("liliang v_ptr_offset: "); print(v_ptr_offset); print("\n");
+    // }
+
+    MUTLASS_PRAGMA_UNROLL
+    for (int i = 0; i < size<0>(tKpK); ++i) {
+      tKpK(i, 0) = loop * get<0>(tile_shape_K) + get<0>(tKcK(0,i,0)) < cur_seqlen_kv;
+    }
+
+    if constexpr (KernelTraits::UseRobustLoadK) {
+      copy_if(gmem_tiled_copy_K.with(params.desc_K), tKpK, tKgK, tKrK);
+    } else {
+      copy_if(gmem_tiled_copy_K, tKpK, tKgK, tKrK);
+    }
+    // ++k_tile_iter;
+    copy(tKrK, tKsK);
+    __syncthreads();
+
+    clear(accum_p);
+    gemm<KernelTraits::IsQInReg>(accum_p, tPrQ, tPrK, tPsQ, tPsK, tiled_mma_qk, thr_copy_Q, thr_copy_K);
+
+      // #pragma unroll
+      // for (int i = 0; i < size(accum_p); ++i) {
+      //   if (accum_p(i) != 128.0f) {
+      //     printf("liliang1 blockIdx = %d, %d, %d, threadIdx.x = %d i = %d, accum_p: ", blockIdx.x, blockIdx.y, blockIdx.z, threadIdx.x, i); print(accum_p(i)); print("\n");
+      //   }
+      // }
+
+    if constexpr (EnableSoftmax) {
+      // do softmax on accum_softmax
+      online_softmax_exp2(accum_softmax);
+      update_accum(accum_o, alpha);
+    }
+
+    // tCvt_frg -> accum_cvt, tAcc_frg -> accum_softmax
+    MUTE_UNROLL
+    for (int i = 0; i < size(tCvt_frg); ++i) {
+      tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::ElementKV, ElementAccum, FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+    }
+
+    // STS P (tPrP -> accum_cvt)
+    copy(tiled_copy_r2s, tPrP, tPsP);
+
+    // ++v_tile_iter;
+    // STS V
+    v_ptr_offset = params.block_table[coord_b * params.block_table_stride_batch + page_idx] * get<3>(params.stride_V) + page_size_idx * get<0>(params.stride_V) + v_num_head_offset;
+    gV = make_tensor(make_gmem_ptr(params.ptr_V + v_ptr_offset), tile_shape_V, select<0, 1>(params.stride_V));
+    tVgV = gmem_thr_copy_V.partition_S(gV); // (VCPY, VCPY_M, VCPY_N)
+    copy(tVrV, tVsV);
+
+    __syncthreads();
+
+    // O = PV
+    gemm(accum_o, tOrP, tOrVt, tOsP, tOsVt, tiled_mma_pv, thr_copy_P, thr_copy_V);
+    // if (thread(0)) {
+    //   print("liliang accum_o: "); print(accum_o(0)); print("\n");
+    // }
+
+    if constexpr (EnableSoftmax) {
+      // do softmax
+      online_softmax_row_sum<ElementAccum>(accum_softmax, row_sum, lse, alpha);
+
+      // Apply Causal Mask
+      if constexpr (EnableCausal) {
+        apply_causal_mask(accum_p, tPcP, coord_q * size<0>(gQ) + mask_m_offset, loop * size<0>(gK), true);
+      }
+      online_softmax_scale(accum_p, accum_softmax, params.rln2_scale, tPcP, residue_kv, loop == orig_kv_tile_count - 1);
+      online_softmax_row_max<ElementAccum>(accum_softmax, row_max, old_row_max, alpha);
+    } else {
+      online_softmax_scale(accum_p, accum_softmax, float(1.0f));
+    }
+  }
+
+  // Last tile
+  if constexpr (EnableSoftmax) {
+    // do softmax
+    online_softmax_exp2(accum_softmax);
+    update_accum(accum_o, alpha);
+  }
+
+  // tCvt_frg -> accum_cvt, tAcc_frg -> accum_softmax
+  MUTE_UNROLL
+  for (int i = 0; i < size(tCvt_frg); ++i) {
+    tCvt_frg(i) = mutlass::NumericArrayConverter<typename KernelTraits::ElementKV, ElementAccum, FragmentSize, mutlass::FloatRoundStyle::round_to_nearest>{}(tAcc_frg(i));
+  }
+
+  // Temp save
+  //copy(tiled_copy_r2s, tPrP, tPgP(_,_,_,0));
+
+  // int page_idx = kv_tile_count * get<0>(tile_shape_K) / page_size;
+  // int page_size_idx = kv_tile_count * get<0>(tile_shape_K) % page_size;
+  // v_ptr_offset = params.block_table[coord_b * params.block_table_stride_batch + page_idx] * get<3>(params.stride_V) + page_size_idx * get<0>(params.stride_V) + (coord_h / mute::size<4,0>(params.problem_shape)) * get<2, 1>(params.stride_V);
+  // gV = make_tensor(make_gmem_ptr(params.ptr_V + v_ptr_offset), tile_shape_V, select<0, 1>(params.stride_V));
+  // tVgV = gmem_thr_copy_V.partition_S(gV); // (VCPY, VCPY_M, VCPY_N)
+
+  // LDG V
+  if constexpr (KernelTraits::UseRobustLoadV) {
+    copy_if(gmem_tiled_copy_V.with(params.desc_V), tKpK, tVgV, tVrV);
+  } else {
+    copy_if(gmem_tiled_copy_V, tKpK, tVgV, tVrV);
+  }
+
+  // STS V
+  copy(tVrV, tVsV);
+
+  // STS P (tPrP -> accum_cvt)
+  copy(tiled_copy_r2s, tPrP, tPsP);
+
+  __syncthreads();
+
+  // if (threadIdx.x == 0) {
+  // if (thread(0)) {
+  //   // print("liliang tiled_copy_r2s: "); print(tiled_copy_r2s); print("\n");
+  //   // print("liliang thr_copy_P: "); print(thr_copy_P); print("\n");
+  //   // print("liliang tPrP: "); print(tPrP); print("\n");
+  //   // print("liliang tPsP: "); print(tPsP); print("\n");
+  //   // print("liliang tPrP: "); print(tPrP(0)); print("\n");
+  //   print("liliang mO: "); print(mO); print("\n");
+  //   print("liliang gO: "); print(gO); print("\n");
+  //   print("liliang tOgO: "); print(tOgO); print("\n");
+  //   print("liliang accum_o: "); print(accum_o); print("\n");
+  //   // printf("liliang coord_h  = %d ", coord_h);print(tPrP(0));print(" ");print(accum_softmax(0));print("\n");
+  // }
+
+  // O = PV
+  gemm(accum_o, tOrP, tOrVt, tOsP, tOsVt, tiled_mma_pv, thr_copy_P, thr_copy_V);
+
+    // if (thread(0)) {
+    //   print("liliang1 tOrP: "); print(tOrP(0)); print("\n");
+    //   print("liliang1 accum_o: "); print(accum_o(0)); print("\n");
+    // }
+  if constexpr (EnableSoftmax) {
+    online_softmax_row_sum<ElementAccum>(accum_softmax, row_sum, lse, alpha);
+
+    apply_softmax_normalizer(accum_o, lse);
+  //   if (threadIdx.x == 0) {
+  //   printf("liliang accum_o:  "); print(accum_o(0)); print(" lse = ");print(lse(0));print("\n");
+  // }
+  }
+
+  if constexpr(KernelTraits::ReturnL) {
+    #pragma unroll
+    for(int i = 0; i < size(lse); ++i) {
+      if (elem_less(get<0>(tCcL(i * size<0, 0>(tCcL))), get<0>(oob_mn))) {
+        tSgL(i) = (log2(lse(i)) + row_max(i)) / float(1.4426950408889634074);
+      }
+    }
+  }
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(accum_o); ++i) {
+    if (elem_less(tCcO(i), oob_mn)) {
+      tOgO(i) = typename KernelTraits::ElementQ(accum_o(i));
+      // tOgO(i) = typename KernelTraits::ElementQ(1.0);
+    }
+  }
+  #endif
+}
diff --git a/vllm_musa/csrc/flash_attn/v1_mp22_impl/attention_params.hpp b/vllm_musa/csrc/flash_attn/v1_mp22_impl/attention_params.hpp
new file mode 100644
index 000000000..9a7d1415e
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/v1_mp22_impl/attention_params.hpp
@@ -0,0 +1,164 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+#include "../impl/mutlass_extension.hpp"
+
+
+struct VariableLength {
+  int max_length;
+  int* cumulative_length = nullptr;
+
+  MUTE_HOST_DEVICE operator int() const {
+    return max_length;
+  }
+};
+
+namespace mute {
+template <>
+struct is_integral<VariableLength> : true_type {};
+
+MUTE_HOST_DEVICE
+void print(VariableLength _) {
+  printf("Varlen<%d, %p>", _.max_length, _.cumulative_length);
+}
+
+} // namespace mute
+
+template <class KernelTraits>
+struct AttentionParams {
+  static constexpr bool IsVarlen = KernelTraits::IsVarlen;
+
+  using ElementQ  = typename KernelTraits::ElementQ;
+  using ElementKV = typename KernelTraits::ElementKV;
+  using ElementAccumulator = typename KernelTraits::ElementAccumulator;
+  using ElementLse = typename KernelTraits::ElementLse;
+
+  using ElementO = ElementQ;
+  using ElementP = ElementKV;
+
+  using IndexType = int32_t;
+
+  // Q, K, D_QK, D_VO, (H_R, H_KV), B
+  using ProblemShapeFixed  = mute::Shape<int, int, int, int, mute::Shape<int, int>, int>;
+  using ProblemShapeVarlen = mute::Shape<VariableLength, VariableLength, int, int, mute::Shape<int, int>, int>;
+  using ProblemShape = mute::conditional_t<IsVarlen,
+                                           ProblemShapeVarlen,
+                                           ProblemShapeFixed>;
+
+  // Q, D_QK, (H_R, H_KV), B
+  using StrideQ = mute::Stride<int, mute::_1, mute::Stride<int, int>, int>;
+
+  // K, D_QK, (H_R, H_KV), B
+  using StrideK = mute::Stride<int, mute::_1, mute::Stride<mute::_0, int>, int>;
+  using StrideL = mute::Stride<mute::_1, mute::_0, mute::Stride<int, int>>;
+
+  // V, D_VO, (H_R, H_KV), B
+  using StrideV = StrideK;
+
+  // Q, D_VO, (H_R, H_KV), B
+  using StrideO = StrideQ;
+
+  // Debug
+  // Q, K, (H_R, H_KV), B
+  using StrideP  = StrideQ;
+
+  struct Arguments {
+    ElementQ const* ptr_Q;
+    ElementKV const* ptr_K;
+    ElementKV const* ptr_V;
+
+    ElementO* ptr_O;
+    ElementLse* ptr_L;
+
+    // Debug
+    ElementP* ptr_P;
+  };
+
+
+  struct Params {
+    ProblemShape problem_shape;
+
+    ElementQ const* ptr_Q;
+    StrideQ stride_Q;
+    mute::RobustDescriptor desc_Q;
+    ElementKV const* ptr_K;
+    StrideK stride_K;
+    mute::RobustDescriptor desc_K;
+    ElementKV const* ptr_V;
+    StrideV stride_V;
+    mute::RobustDescriptor desc_V;
+
+    ElementO* ptr_O;
+    StrideO stride_O;
+
+    ElementLse* ptr_L;
+    StrideL stride_L;
+
+    float rln2_scale;
+    int* block_table;
+    int block_table_stride_batch;
+
+    // Debug
+    ElementP* ptr_P;
+    StrideP stride_P;
+  };
+
+  static Params
+  to_underlying_arguments(ProblemShape problem_shape, Arguments const& args,
+                          mute::Shape<int, int, int> qkv_elem,
+    IndexType q_seq_stride, IndexType q_num_head_stride,
+    IndexType k_num_page_stride, IndexType k_page_size_stride, IndexType k_num_head_stride,
+    IndexType v_num_page_stride, IndexType v_page_size_stride, IndexType v_num_head_stride,
+    IndexType o_seq_stride, IndexType o_num_head_stride,
+    IndexType lse_num_head_stride,
+    int* block_table, int block_table_stride_batch, float sm_scale
+                          ) {
+
+    int DimQK = mute::size<2>(problem_shape);
+    int H_KV  = mute::size<4,1>(problem_shape);
+    int H_R  = mute::size<4,0>(problem_shape);
+
+
+    StrideQ stride_Q = mute::make_stride(q_seq_stride, mute::_1{},
+                      mute::make_stride(q_num_head_stride, q_num_head_stride * H_R), 0); // TODO:
+    StrideK stride_K = mute::make_stride(k_page_size_stride, mute::_1{},
+                      mute::make_stride(mute::_0{}, k_num_head_stride), k_num_page_stride);
+    StrideV stride_V = mute::make_stride(v_page_size_stride, mute::_1{},
+                      mute::make_stride(mute::_0{}, v_num_head_stride), v_num_page_stride);
+    StrideO stride_O = mute::make_stride(o_seq_stride, mute::_1{},
+                      mute::make_stride(o_num_head_stride, o_num_head_stride * H_R), 0);
+    StrideL stride_L = mute::make_stride(mute::_1{}, mute::_0{}, mute::make_stride(lse_num_head_stride, lse_num_head_stride * H_R)); // TODO:
+
+    mute::RobustDescriptor desc_Q = mute::make_robust_desc(args.ptr_Q, mute::get<0>(qkv_elem));
+    mute::RobustDescriptor desc_K = mute::make_robust_desc(args.ptr_K, mute::get<1>(qkv_elem));
+    mute::RobustDescriptor desc_V = mute::make_robust_desc(args.ptr_V, mute::get<2>(qkv_elem));
+
+    constexpr float kLog2e = float(1.4426950408889634074);
+    const float softmax_scale = sm_scale;
+    float rln2_scale = softmax_scale * kLog2e;
+
+    StrideP stride_P;
+
+    return {
+      problem_shape,
+      args.ptr_Q,
+      stride_Q,
+      desc_Q,
+      args.ptr_K,
+      stride_K,
+      desc_K,
+      args.ptr_V,
+      stride_V,
+      desc_V,
+      args.ptr_O,
+      stride_O,
+      args.ptr_L,
+      stride_L,
+      rln2_scale,
+      block_table,
+      block_table_stride_batch,
+      args.ptr_P,
+      stride_P
+    };
+  }
+};
diff --git a/vllm_musa/csrc/flash_attn/v1_mp22_impl/kernel_traits.hpp b/vllm_musa/csrc/flash_attn/v1_mp22_impl/kernel_traits.hpp
new file mode 100644
index 000000000..e3edf2755
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/v1_mp22_impl/kernel_traits.hpp
@@ -0,0 +1,204 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mute/atom/mma_atom.hpp>
+#include <mute/atom/copy_atom.hpp>
+
+#include <mutlass/mutlass.h>
+#include <mutlass/gemm/gemm.h>
+#include <mutlass/gemm/collective/collective_builder.hpp>
+
+#include "../impl/mutlass_extension.hpp"
+
+// using namespace mute;
+
+template <
+  class ElementQ,
+  class ElementKV,
+  class SmemLayoutQ,
+  class SmemLayoutK,
+  class SmemLayoutV,
+  class SmemLayoutP
+>
+struct SharedStorage {
+  mute::array_aligned<ElementQ,  mute::cosize_v<SmemLayoutQ>> smem_q;
+  mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutK>> smem_k;
+  mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutV>> smem_v;
+  mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutP>> smem_p; //TODO:
+};
+
+
+template <
+  class ElementQ,
+  class ElementKV,
+  class SmemLayoutQ,
+  class SmemLayoutK,
+  class SmemLayoutV,
+  class SmemLayoutP
+>
+struct SharedStorageReuseQ {
+  union {
+    mute::array_aligned<ElementQ,  mute::cosize_v<SmemLayoutQ>> smem_q;
+    struct {
+      mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutK>> smem_k;
+      mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutV>> smem_v;
+      mute::array_aligned<ElementKV, mute::cosize_v<SmemLayoutP>> smem_p;
+    };
+  };
+};
+
+
+template <class TileN, class PermuteTileShape = TileN>
+constexpr auto
+get_permute_tile_for_operand_B() {
+  using MmaAtomN = mute::_16;
+  static_assert(TileN{} % PermuteTileShape{} == 0);
+  static_assert(PermuteTileShape{} % MmaAtomN{} == 0);
+
+  using RepeatInsidePermuteTile = decltype(PermuteTileShape{} / MmaAtomN{});
+  using RepeatAcrossPermuteTile = decltype(TileN{} / PermuteTileShape{});
+
+  using PermuteTileN = mute::Layout<mute::Shape <               MmaAtomN, RepeatInsidePermuteTile, RepeatAcrossPermuteTile>,
+                              mute::Stride<RepeatInsidePermuteTile,                      mute::_1,        PermuteTileShape>>;
+
+  return PermuteTileN{};
+}
+
+template <
+  int HEAD_DIM_QK_,
+  int HEAD_DIM_VO_,
+  int CTA_Q_,
+  int CTA_KV_,
+  bool IsVarlen_,
+  bool IsQInReg_,
+  bool UseRobustLoadQ_,
+  bool UseRobustLoadK_,
+  bool UseRobustLoadV_,
+  class ElementQ_,
+  bool ReturnL_,
+  class ElementKV_ = ElementQ_,
+  class IdType_ = int32_t
+>
+struct AttentionForwardKernelTraits {
+  using ElementQ  = ElementQ_;
+  using ElementKV = ElementKV_;
+  using ElementAccumulator = float;
+  using ElementLse = float;
+
+  static constexpr int CTA_Q       = CTA_Q_;
+  static constexpr int CTA_KV      = CTA_KV_;
+  static constexpr int HEAD_DIM_QK = HEAD_DIM_QK_;
+  static constexpr int HEAD_DIM_VO = HEAD_DIM_VO_;
+  static constexpr bool IsVarlen = IsVarlen_;
+  static constexpr bool IsQInReg = IsQInReg_;
+  static constexpr bool ReuseSmemQ = IsQInReg;
+  static constexpr bool UseRobustLoadQ = UseRobustLoadQ_;
+  static constexpr bool UseRobustLoadK = UseRobustLoadK_;
+  static constexpr bool UseRobustLoadV = UseRobustLoadV_;
+  static constexpr bool ReturnL = ReturnL_;
+
+  // TODO: now only pack T0V0/TV01, try to pack more regs
+  static constexpr int FragmentSize = 2;
+
+  using TileShape_QKD = mute::Shape<mute::Int<CTA_Q>, mute::Int<CTA_KV>, mute::Int<HEAD_DIM_QK>>;
+  using TileShape_PDV = mute::Shape<mute::Int<CTA_Q>, mute::Int<HEAD_DIM_VO>, mute::Int<CTA_KV>>;
+
+  // We extend in M-dim to avoid smem reduce
+  using AtomLayoutMNK = mute::Layout<mute::Shape<mute::Int<CTA_Q/32>, mute::_1, mute::_1>>;
+
+  // Internal Mma Stride type
+  using MmaStrideMNMajor = mute::Stride<mute::_1, int>;
+  using MmaStrideKMajor  = mute::Stride<int, mute::_1>;
+
+  using PermutedTileN = decltype(get_permute_tile_for_operand_B<mute::Int<CTA_KV>, mute::Int<32>>());
+  using PermuteLayoutMNK = mute::Tile<mute::Underscore, PermutedTileN, mute::Underscore>;
+
+  using TiledMmaQK = decltype(mute::make_tiled_mma(
+    mutlass::gemm::collective::detail::mp22_mma_operation_select<
+                                        ElementQ, MmaStrideKMajor,
+                                        ElementKV,MmaStrideKMajor>(),
+    AtomLayoutMNK{},
+    PermuteLayoutMNK{}));
+
+  using TiledMmaPV =decltype(mute::make_tiled_mma(
+    mutlass::gemm::collective::detail::mp22_mma_operation_select<
+                                        ElementKV, MmaStrideKMajor,
+                                        ElementKV, MmaStrideMNMajor>(),
+    AtomLayoutMNK{}));
+
+  static constexpr int NumThreads = size(TiledMmaQK{});
+  static_assert(size(TiledMmaQK{}) == size(TiledMmaPV{}));
+
+  using SmemLayoutAtomQ = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementQ,  MmaStrideKMajor>());
+  using SmemLayoutQ = decltype(tile_to_shape(SmemLayoutAtomQ{}, mute::select<0,2>(TileShape_QKD{})));
+
+  using SmemLayoutAtomK = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutK = decltype(tile_to_shape(SmemLayoutAtomK{}, mute::select<1,2>(TileShape_QKD{})));
+
+  using SmemLayoutAtomP = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutP = decltype(tile_to_shape(SmemLayoutAtomP{}, mute::select<0,1>(TileShape_QKD{})));
+
+  using SmemLayoutAtomV = decltype(
+    mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementKV, MmaStrideKMajor>());
+  using SmemLayoutV  = decltype(tile_to_shape(SmemLayoutAtomV{}, mute::select<2,1>(TileShape_PDV{})));
+  using SmemLayoutVt = decltype(composition(SmemLayoutV{},
+                        make_layout(mute::select<1,2>(TileShape_PDV{}), mute::LayoutRight{})));
+
+  using SmemCopyAtomQ = mute::Copy_Atom<mute::DefaultCopy, ElementQ>;
+  using SmemCopyAtomK = mute::Copy_Atom<mute::DefaultCopy, ElementKV>;
+  using SmemCopyAtomP = mute::Copy_Atom<mute::DefaultCopy, ElementKV>;
+  using SmemCopyAtomV = mute::Copy_Atom<mute::DefaultCopy, ElementKV>;
+
+  using R2SCopyAtomP  = mute::Copy_Atom<mute::UniversalCopy<mute::uint_byte_t<FragmentSize * sizeof(ElementKV)>>, ElementKV>;
+
+  // TileCopy
+  static constexpr int AlignmentQ = std::min(128 / mute::sizeof_bits_v<ElementQ>, size<0>(TileShape_QKD{}) * size<2>(TileShape_QKD{}) / NumThreads);
+  static constexpr int ThreadsMajorQ = mute::gcd(HEAD_DIM_QK / AlignmentQ, NumThreads);
+  static constexpr int ThreadsMinorQ = NumThreads / ThreadsMajorQ;
+  using GmemCopyQThrLayout = mute::Layout<mute::Shape <mute::Int<ThreadsMinorQ>, mute::Int<ThreadsMajorQ>>,
+                                    mute::Stride<mute::Int<ThreadsMajorQ>,                mute::_1>>;
+  using GmemCopyQValLayout = mute::Layout<mute::Shape<mute::_1, mute::Int<AlignmentQ>>>;
+  using GmemCopyOpQ = mute::conditional_t<UseRobustLoadQ,
+            mute::MP22_ROBUST_LOAD<mute::uint_bit_t<AlignmentQ * mute::sizeof_bits_v<ElementQ>>>,
+            mute::UniversalCopy<mute::uint_bit_t<AlignmentQ * mute::sizeof_bits_v<ElementQ>>>>;
+
+  using GmemTiledCopyQ = decltype(make_tiled_copy(
+                            mute::Copy_Atom<GmemCopyOpQ, ElementQ>{},
+                            GmemCopyQThrLayout{}, GmemCopyQValLayout{}));
+
+  static constexpr int AlignmentK = std::min(128 / mute::sizeof_bits_v<ElementKV>, size<1>(TileShape_QKD{}) * size<2>(TileShape_QKD{}) / NumThreads);
+  static constexpr int ThreadsMajorK = mute::gcd(HEAD_DIM_QK / AlignmentK, NumThreads);
+  static constexpr int ThreadsMinorK = NumThreads / ThreadsMajorK;
+  using GmemCopyKThrLayout = mute::Layout<mute::Shape <mute::Int<ThreadsMinorK>, mute::Int<ThreadsMajorK>>,
+                                    mute::Stride<mute::Int<ThreadsMajorK>,                mute::_1>>;
+  using GmemCopyKValLayout = mute::Layout<mute::Shape<mute::_1, mute::Int<AlignmentK>>>;
+  using GmemCopyOpK = mute::conditional_t<UseRobustLoadK,
+            mute::MP22_ROBUST_LOAD<mute::uint_bit_t<AlignmentK * mute::sizeof_bits_v<ElementKV>>>,
+            mute::UniversalCopy<mute::uint_bit_t<AlignmentK * mute::sizeof_bits_v<ElementKV>>>>;
+
+
+  using GmemTiledCopyK = decltype(make_tiled_copy(
+                            mute::Copy_Atom<GmemCopyOpK, ElementKV>{},
+                            GmemCopyKThrLayout{}, GmemCopyKValLayout{}));
+
+  static constexpr int AlignmentV = std::min(128 / mute::sizeof_bits_v<ElementKV>, size<1>(TileShape_PDV{}) * size<2>(TileShape_PDV{}) / NumThreads);
+  static constexpr int ThreadsMajorV = mute::gcd(HEAD_DIM_VO / AlignmentV, NumThreads);
+  static constexpr int ThreadsMinorV = NumThreads / ThreadsMajorV;
+  using GmemCopyVThrLayout = mute::Layout<mute::Shape <mute::Int<ThreadsMinorV>, mute::Int<ThreadsMajorV>>,
+                                    mute::Stride<mute::Int<ThreadsMajorV>,                mute::_1>>;
+  using GmemCopyVValLayout = mute::Layout<mute::Shape<mute::_1, mute::Int<AlignmentV>>>;
+  using GmemCopyOpV = mute::conditional_t<UseRobustLoadV,
+            mute::MP22_ROBUST_LOAD<mute::uint_bit_t<AlignmentV * mute::sizeof_bits_v<ElementKV>>>,
+            mute::UniversalCopy<mute::uint_bit_t<AlignmentV * mute::sizeof_bits_v<ElementKV>>>>;
+
+  using GmemTiledCopyV = decltype(make_tiled_copy(
+                            mute::Copy_Atom<GmemCopyOpV, ElementKV>{},
+                            GmemCopyVThrLayout{}, GmemCopyVValLayout{}));
+
+  using SharedStorage = mute::conditional_t<ReuseSmemQ,
+      SharedStorageReuseQ<ElementQ, ElementKV, SmemLayoutQ, SmemLayoutK, SmemLayoutV, SmemLayoutP>,
+      SharedStorage<ElementQ, ElementKV, SmemLayoutQ, SmemLayoutK, SmemLayoutV, SmemLayoutP>>;
+};
diff --git a/vllm_musa/csrc/flash_attn/varlen_fa_padding_kernel.mu b/vllm_musa/csrc/flash_attn/varlen_fa_padding_kernel.mu
new file mode 100644
index 000000000..d01a94192
--- /dev/null
+++ b/vllm_musa/csrc/flash_attn/varlen_fa_padding_kernel.mu
@@ -0,0 +1,563 @@
+#include <musa_runtime.h>
+#include <vector>
+#include <musa_fp16.h>
+#include <iostream>
+#include <random>
+#include <cassert>
+#include <chrono>
+#include <torch/torch.h>
+#include <torch_musa/csrc/core/MUSAGuard.h>
+
+using vec_128bit  = uint32_t __attribute__((vector_size(16)));
+
+template <typename T, int BLOCKDIM_X, int BLOCKDIM_Y, bool IsPaddingKV>
+__global__ void padding_kernel(
+    const T* input1,         //  [sum_seq, head_dim, num_heads]
+    const T* input2,         //  [sum_seq, head_dim, num_heads]
+    T* output1,              //  [max_seq, head_dim, num_heads, batch_size]
+    T* output2,
+    const int* accum_length,//  [seq_0, seq_1,...]
+    int max_seq,
+    int head_dim,
+    int num_heads,
+    int batch_size,
+    int seq_stride,int head_stride,int seq_out_stride,int head_out_stride,int bs_out_stride
+) {
+    const int input_base_idx = blockIdx.y * head_stride + (accum_length[blockIdx.z] + blockIdx.x * BLOCKDIM_Y + threadIdx.y) * seq_stride;
+    const int out_base_idx = blockIdx.z * bs_out_stride + blockIdx.y*head_out_stride;
+
+    const int seq_length = accum_length[blockIdx.z+1] - accum_length[blockIdx.z];
+
+    const int seq_out_idx = blockIdx.x * BLOCKDIM_Y + threadIdx.y;
+
+    bool is_oob = threadIdx.x * 16 / sizeof(T) >= head_dim;
+
+    if (seq_out_idx >= max_seq || is_oob) {
+        return;
+    }
+
+    int output_pos = out_base_idx + seq_out_idx * seq_out_stride;
+    const vec_128bit* input_vec1  = reinterpret_cast<const vec_128bit*>(input1 + input_base_idx);
+    const vec_128bit* input_vec2  = reinterpret_cast<const vec_128bit*>(input2 + input_base_idx);
+
+    vec_128bit* output_vec_1 = reinterpret_cast<vec_128bit*>(output1 + output_pos);
+    vec_128bit* output_vec_2 = reinterpret_cast<vec_128bit*>(output2 + output_pos);
+
+
+    if (seq_out_idx < seq_length) {
+        output_vec_1[threadIdx.x] = input_vec1[threadIdx.x];
+        if constexpr (IsPaddingKV) {
+            output_vec_2[threadIdx.x] = input_vec2[threadIdx.x];
+        }
+    } else {
+        vec_128bit Zeros = {};
+        output_vec_1[threadIdx.x] =  Zeros;
+        if constexpr (IsPaddingKV) {
+            output_vec_2[threadIdx.x] =  Zeros;
+        }
+    }
+}
+
+template <typename T, int BLOCKDIM_X, int BLOCKDIM_Y>
+__global__ void unpadding_kernel(
+    const T* input,         //  [max_seq, head_dim, num_heads, batch_size]
+    T* output,              //  [sum_seq, head_dim, num_heads]
+    const int* accum_length,//  [seq_0, seq_1,...]
+    int sum_seq,            // (seq_lens)）
+    int max_seq,
+    int head_dim,
+    int num_heads,
+    int batch_size,
+    int seq_stride,int head_stride,int bs_stride, int seq_out_stride,int head_out_stride
+) {
+    const int out_base_idx = blockIdx.y * head_out_stride + (accum_length[blockIdx.z] + blockIdx.x * BLOCKDIM_Y + threadIdx.y) * seq_out_stride;
+    const int in_base_idx = blockIdx.z * bs_stride + blockIdx.y*head_stride;
+
+    const int seq_length = accum_length[blockIdx.z+1] - accum_length[blockIdx.z];
+
+    const int seq_in_idx = blockIdx.x * BLOCKDIM_Y + threadIdx.y;
+    bool is_oob = threadIdx.x * 16 / sizeof(T) >= head_dim;
+
+    if (seq_in_idx >= max_seq || is_oob) {
+        return;
+    }
+
+    int input_pos = in_base_idx + seq_in_idx * seq_stride;
+    vec_128bit* output_vec  = reinterpret_cast<vec_128bit*>(output + out_base_idx);
+
+    const vec_128bit* input_vec = reinterpret_cast<const vec_128bit*>(input + input_pos);
+
+    if (seq_in_idx < seq_length) {
+        output_vec[threadIdx.x] = input_vec[threadIdx.x];
+    }
+}
+
+template<class T>
+void launch_padding_kernel(
+    const T* input_q,
+    const T* input_k,
+    const T* input_v,
+    T* output_q,
+    T* output_k,
+    T* output_v,
+    const int* accum_lens_q,
+    const int* accum_lens_k,
+    int seqlen_q,
+    int seqlen_kv,
+    int head_dim_q,
+    int head_dim_k,
+    int head_dim_v,
+    int num_heads_q,
+    int num_heads_k,
+    int num_heads_v,
+    int batch_size,
+    int q_seq_stride, int q_head_stride,int q_seq_out_stride,int q_head_out_stride,int q_bs_out_stride,
+    int k_seq_stride, int k_head_stride,int k_seq_out_stride, int k_head_out_stride,int k_bs_out_stride,
+    int v_seq_stride, int v_head_stride,int v_seq_out_stride, int v_head_out_stride,int v_bs_out_stride,
+    musaStream_t stream = 0
+) {
+    if (head_dim_q > 128) {
+
+        dim3 block(32, 32);
+        dim3 grid_q(
+            (seqlen_q + block.y - 1) / block.y,
+            num_heads_q,
+            batch_size
+        );
+        dim3 grid_k(
+            (seqlen_kv + block.y - 1) / block.y,
+            num_heads_k,
+            batch_size
+        );
+        dim3 grid_v(
+            (seqlen_kv + block.y - 1) / block.y,
+            num_heads_v,
+            batch_size
+        );
+        padding_kernel<T, 32, 32, false><<<grid_q, block, 0, stream>>>(
+        input_q, nullptr, output_q, nullptr, accum_lens_q,
+        seqlen_q, head_dim_q, num_heads_q, batch_size,
+        q_seq_stride,q_head_stride,q_seq_out_stride,q_head_out_stride,q_bs_out_stride
+        );
+        padding_kernel<T, 32, 32, false><<<grid_k, block, 0, stream>>>(
+        input_k, nullptr, output_k, nullptr,accum_lens_k,
+        seqlen_kv, head_dim_k, num_heads_k, batch_size,
+        k_seq_stride, k_head_stride,k_seq_out_stride, k_head_out_stride,k_bs_out_stride
+        );
+        padding_kernel<T, 32, 32, false><<<grid_v, block, 0, stream>>>(
+        input_v, nullptr, output_v, nullptr,accum_lens_k,
+        seqlen_kv, head_dim_v, num_heads_v, batch_size,
+        v_seq_stride, v_head_stride,v_seq_out_stride, v_head_out_stride,v_bs_out_stride
+        );
+
+    }
+    else if (head_dim_q > 64) {
+        dim3 block(16, 32);
+        dim3 grid_q(
+            (seqlen_q+ block.y - 1) / block.y,
+            num_heads_q,
+            batch_size
+        );
+        dim3 grid_k(
+            (seqlen_kv + block.y - 1) / block.y,
+            num_heads_k,
+            batch_size
+        );
+        dim3 grid_v(
+            (seqlen_kv + block.y - 1) / block.y,
+            num_heads_v,
+            batch_size
+        );
+        padding_kernel<T, 16, 32, false><<<grid_q, block, 0, stream>>>(
+        input_q, nullptr, output_q, nullptr, accum_lens_q,
+        seqlen_q, head_dim_q, num_heads_q, batch_size,
+        q_seq_stride,q_head_stride,q_seq_out_stride,q_head_out_stride,q_bs_out_stride
+        );
+        padding_kernel<T, 16, 32, false><<<grid_k, block, 0, stream>>>(
+        input_k, nullptr, output_k, nullptr,accum_lens_k,
+        seqlen_kv, head_dim_k, num_heads_k, batch_size,
+        k_seq_stride, k_head_stride,k_seq_out_stride, k_head_out_stride,k_bs_out_stride
+        );
+        padding_kernel<T, 16, 32, false><<<grid_v, block, 0, stream>>>(
+        input_v, nullptr, output_v, nullptr,accum_lens_k,
+        seqlen_kv, head_dim_v, num_heads_v, batch_size,
+        v_seq_stride, v_head_stride,v_seq_out_stride, v_head_out_stride,v_bs_out_stride
+        );
+    } else if (head_dim_q > 32) {
+        dim3 block(8, 64);
+        dim3 grid_q(
+            (seqlen_q + block.y - 1) / block.y,
+            num_heads_q,
+            batch_size
+        );
+        dim3 grid_k(
+            (seqlen_kv + block.y - 1) / block.y,
+            num_heads_k,
+            batch_size
+        );
+        dim3 grid_v(
+            (seqlen_kv +  block.y - 1) / block.y,
+            num_heads_v,
+            batch_size
+        );
+        padding_kernel<T, 8, 64, false><<<grid_q, block, 0, stream>>>(
+        input_q, nullptr, output_q, nullptr, accum_lens_q,
+        seqlen_q, head_dim_q, num_heads_q, batch_size,
+        q_seq_stride,q_head_stride,q_seq_out_stride,q_head_out_stride,q_bs_out_stride
+        );
+        padding_kernel<T, 8, 64, false><<<grid_k, block, 0, stream>>>(
+        input_k, nullptr, output_k, nullptr,accum_lens_k,
+        seqlen_kv, head_dim_k, num_heads_k, batch_size,
+        k_seq_stride, k_head_stride,k_seq_out_stride, k_head_out_stride,k_bs_out_stride
+        );
+        padding_kernel<T, 8, 64, false><<<grid_v, block, 0, stream>>>(
+        input_v, nullptr, output_v, nullptr,accum_lens_k,
+        seqlen_kv, head_dim_v, num_heads_v, batch_size,
+        v_seq_stride, v_head_stride,v_seq_out_stride, v_head_out_stride,v_bs_out_stride
+        );
+    } else {
+        printf("no support on pad kernel\n");
+    }
+}
+
+template<class T>
+void launch_unpadding_kernel(
+    const T* input,
+    T* output,
+    const int* accum_lens,
+    int sum_seq,
+    int max_seq,
+    int head_dim,
+    int num_heads,
+    int batch_size,
+    int seq_out_stride, int head_out_stride,int seq_stride,int head_stride,int bs_stride,
+    musaStream_t stream = 0
+) {
+    if (head_dim > 128) {
+
+        dim3 block(32, 32);
+        dim3 grid(
+            (max_seq + block.y - 1) / block.y,
+            num_heads,
+            batch_size
+        );
+        unpadding_kernel<T, 32, 32><<<grid, block, 0, stream>>>(
+        input, output,accum_lens,
+        sum_seq, max_seq, head_dim, num_heads, batch_size,
+        seq_stride, head_stride,bs_stride, seq_out_stride, head_out_stride
+        );
+    }
+    else if (head_dim > 64) {
+        dim3 block(16, 32);
+        dim3 grid(
+            (max_seq + block.y - 1) / block.y,
+            num_heads,
+            batch_size
+        );
+        unpadding_kernel<T, 16, 32><<<grid, block, 0, stream>>>(
+        input, output, accum_lens,
+        sum_seq, max_seq, head_dim, num_heads, batch_size,
+        seq_stride, head_stride,bs_stride, seq_out_stride, head_out_stride
+        );
+    } else if (head_dim > 32) {
+        dim3 block(8, 64);
+        dim3 grid(
+            (max_seq + block.y - 1) / block.y,
+            num_heads,
+            batch_size
+        );
+        unpadding_kernel<T, 8, 64><<<grid, block, 0, stream>>>(
+        input, output,accum_lens,
+        sum_seq, max_seq, head_dim, num_heads, batch_size,
+        seq_stride, head_stride,bs_stride, seq_out_stride, head_out_stride
+        );
+    } else {
+        printf("no support on unpad kernel\n");
+    }
+}
+
+void varlen_fa_seqlen_pad(
+    torch::Tensor &q_in,
+    torch::Tensor &k_in,
+    torch::Tensor &v_in,
+    torch::Tensor &q_out,
+    torch::Tensor &k_out,
+    torch::Tensor &v_out,
+    torch::Tensor &accum_lens_q,
+    torch::Tensor &accum_lens_k,
+    int sum_seq,
+    int max_seq,
+    int batch_size) {
+
+    int q_seq_stride = q_in.strides()[0];
+    int q_head_stride = q_in.strides()[1];
+
+    int q_seq_out_stride = q_out.strides()[2];
+    int q_head_out_stride = q_out.strides()[1];
+    int q_bs_out_stride = q_out.strides()[0];
+
+
+    int k_seq_stride = k_in.strides()[0];
+    int k_head_stride = k_in.strides()[1];
+
+    int k_seq_out_stride = k_out.strides()[2];
+    int k_head_out_stride = k_out.strides()[1];
+    int k_bs_out_stride = k_out.strides()[0];
+
+    int v_seq_stride = v_in.strides()[0];
+    int v_head_stride = v_in.strides()[1];
+
+    int v_seq_out_stride = v_out.strides()[2];
+    int v_head_out_stride = v_out.strides()[1];
+    int v_bs_out_stride = v_out.strides()[0];
+
+
+
+    int head_dim_q = q_in.sizes()[2];
+    int head_dim_k = k_in.sizes()[2];
+    int head_dim_v = v_in.sizes()[2];
+
+    int num_heads_q = q_in.sizes()[1];
+    int num_heads_k = k_in.sizes()[1];
+    int num_heads_v = v_in.sizes()[1];
+
+    int seqlen_q = q_out.sizes()[2];
+    int seqlen_kv = k_out.sizes()[2];
+
+    const at::musa::OptionalMUSAGuard device_guard(device_of(q_in));
+    auto stream = c10::musa::getCurrentMUSAStream().stream();
+
+    if (q_in.scalar_type() == at::ScalarType::BFloat16)
+    {
+        launch_padding_kernel<__mt_bfloat16>(
+            static_cast<const __mt_bfloat16*>(q_in.data_ptr()),
+            static_cast<const __mt_bfloat16*>(k_in.data_ptr()),
+            static_cast<const __mt_bfloat16*>(v_in.data_ptr()),
+            static_cast<__mt_bfloat16*>(q_out.data_ptr()),
+            static_cast<__mt_bfloat16*>(k_out.data_ptr()),
+            static_cast<__mt_bfloat16*>(v_out.data_ptr()),
+            static_cast<int*>(accum_lens_q.data_ptr()),
+            static_cast<int*>(accum_lens_k.data_ptr()),
+            seqlen_q,
+            seqlen_kv,
+            head_dim_q,
+            head_dim_k,
+            head_dim_v,
+            num_heads_q,
+            num_heads_k,
+            num_heads_v,
+            batch_size,
+            q_seq_stride,q_head_stride,q_seq_out_stride,q_head_out_stride,q_bs_out_stride,
+            k_seq_stride, k_head_stride,k_seq_out_stride,k_head_out_stride, k_bs_out_stride,
+            v_seq_stride, v_head_stride,v_seq_out_stride,v_head_out_stride, v_bs_out_stride, stream);
+    }
+    else if (q_in.scalar_type() == at::ScalarType::Half)
+    {
+        launch_padding_kernel<__half>(
+            static_cast<const __half*>(q_in.data_ptr()),
+            static_cast<const __half*>(k_in.data_ptr()),
+            static_cast<const __half*>(v_in.data_ptr()),
+            static_cast<__half*>(q_out.data_ptr()),
+            static_cast<__half*>(k_out.data_ptr()),
+            static_cast<__half*>(v_out.data_ptr()),
+            static_cast<int*>(accum_lens_q.data_ptr()),
+            static_cast<int*>(accum_lens_k.data_ptr()),
+            seqlen_q,
+            seqlen_kv,
+            head_dim_q,
+            head_dim_k,
+            head_dim_v,
+            num_heads_q,
+            num_heads_k,
+            num_heads_v,
+            batch_size,
+            q_seq_stride,q_head_stride,q_seq_out_stride,q_head_out_stride,q_bs_out_stride,
+            k_seq_stride, k_head_stride,k_seq_out_stride,k_head_out_stride, k_bs_out_stride,
+            v_seq_stride, v_head_stride,v_seq_out_stride,v_head_out_stride, v_bs_out_stride, stream);
+    }
+    else if (q_in.scalar_type() == at::ScalarType::Float)
+    {
+        launch_padding_kernel<float>(
+            static_cast<const float*>(q_in.data_ptr()),
+            static_cast<const float*>(k_in.data_ptr()),
+            static_cast<const float*>(v_in.data_ptr()),
+            static_cast<float*>(q_out.data_ptr()),
+            static_cast<float*>(k_out.data_ptr()),
+            static_cast<float*>(v_out.data_ptr()),
+            static_cast<int*>(accum_lens_q.data_ptr()),
+            static_cast<int*>(accum_lens_k.data_ptr()),
+            seqlen_q,
+            seqlen_kv,
+            head_dim_q,
+            head_dim_k,
+            head_dim_v,
+            num_heads_q,
+            num_heads_k,
+            num_heads_v,
+            batch_size,
+            q_seq_stride,q_head_stride,q_seq_out_stride,q_head_out_stride,q_bs_out_stride,
+            k_seq_stride, k_head_stride,k_seq_out_stride,k_head_out_stride, k_bs_out_stride,
+            v_seq_stride, v_head_stride,v_seq_out_stride,v_head_out_stride, v_bs_out_stride, stream);
+    }
+    else
+    {
+        TORCH_CHECK(false, "only support Half and BFloat16 dtype");
+    }
+}
+
+void varlen_fa_seqlen_unpad(
+    torch::Tensor &in,
+    torch::Tensor &out,
+    torch::Tensor &accum_lens,
+    int sum_seq,
+    int max_seq,
+    int head_dim,
+    int num_heads,
+    int batch_size) {
+
+    int seq_out_stride = out.strides()[0];
+    int head_out_stride = out.strides()[1];
+
+    int seq_stride = in.strides()[2];
+    int head_stride = in.strides()[1];
+    int bs_stride = in.strides()[0];
+
+    const at::musa::OptionalMUSAGuard device_guard(device_of(in));
+    auto stream = c10::musa::getCurrentMUSAStream().stream();
+
+    if (in.scalar_type() == at::ScalarType::BFloat16)
+    {
+        launch_unpadding_kernel<__mt_bfloat16>(
+            static_cast<const __mt_bfloat16*>(in.data_ptr()),
+            static_cast<__mt_bfloat16*>(out.data_ptr()),
+            static_cast<int*>(accum_lens.data_ptr()),
+            sum_seq,
+            max_seq,
+            head_dim,
+            num_heads,
+            batch_size,
+            seq_out_stride,head_out_stride,seq_stride,head_stride,bs_stride, stream);
+    }
+    else if (in.scalar_type() == at::ScalarType::Half)
+    {
+        launch_unpadding_kernel<__half>(
+            static_cast<const __half*>(in.data_ptr()),
+            static_cast<__half*>(out.data_ptr()),
+            static_cast<int*>(accum_lens.data_ptr()),
+            sum_seq,
+            max_seq,
+            head_dim,
+            num_heads,
+            batch_size,
+            seq_out_stride,head_out_stride,seq_stride,head_stride,bs_stride, stream);
+    }
+    else if (in.scalar_type() == at::ScalarType::Float)
+    {
+        launch_unpadding_kernel<float>(
+            static_cast<const float*>(in.data_ptr()),
+            static_cast<float*>(out.data_ptr()),
+            static_cast<int*>(accum_lens.data_ptr()),
+            sum_seq,
+            max_seq,
+            head_dim,
+            num_heads,
+            batch_size,
+            seq_out_stride,head_out_stride,seq_stride,head_stride,bs_stride, stream);
+    }
+    else
+    {
+        TORCH_CHECK(false, "only support Half and BFloat16 dtype");
+    }
+}
+
+constexpr int VEC_SIZE = 4;
+template <int BLOCK_THREADS>
+__global__ void varlen_fa_lse_seqlen_unpad_kernel(
+    const float* __restrict__ input,     // [batch_size, num_heads, max_seq]
+    float* __restrict__ output,           // [num_heads, sum_seq]
+    const int* __restrict__ accum_length, // [batch_size+1]
+    const int max_seq,
+    const int num_heads,
+    const int batch_size,
+    const int head_stride,       // input.stride(1)
+    const int head_out_stride    // output.stride(0)
+) {
+    const int batch_idx = blockIdx.z;
+    const int head_idx = blockIdx.y;
+    const int tid = threadIdx.x;
+
+    if (batch_idx >= batch_size || head_idx >= num_heads)
+        return;
+
+    const int seq_start = accum_length[batch_idx];
+    const int seq_end = accum_length[batch_idx + 1];
+    const int seq_len = seq_end - seq_start;
+
+    const float* in_ptr = input +
+                         batch_idx * num_heads * head_stride +
+                         head_idx * head_stride;
+
+    float* out_ptr = output + head_idx * head_out_stride + seq_start;
+
+    const int num_vec = (seq_len + VEC_SIZE - 1) / VEC_SIZE;
+
+    for (int vec_idx = tid; vec_idx < num_vec; vec_idx += BLOCK_THREADS) {
+        float4 vec_val;
+        const int elem_offset = vec_idx * VEC_SIZE;
+        const int remaining = seq_len - elem_offset;
+
+        if (remaining >= VEC_SIZE) {
+            vec_val = reinterpret_cast<const float4*>(in_ptr)[vec_idx];
+        } else {
+            #pragma unroll
+            for (int i = 0; i < VEC_SIZE; ++i) {
+                reinterpret_cast<float*>(&vec_val)[i] =
+                    (i < remaining) ? in_ptr[elem_offset + i] : 0.0f;
+            }
+        }
+
+        if (elem_offset < seq_len) {
+            const int store_size = min(VEC_SIZE, remaining);
+            if (store_size == VEC_SIZE) {
+                reinterpret_cast<float4*>(out_ptr)[vec_idx] = vec_val;
+            } else {
+                #pragma unroll
+                for (int i = 0; i < store_size; ++i) {
+                    out_ptr[elem_offset + i] = reinterpret_cast<float*>(&vec_val)[i];
+                }
+            }
+        }
+    }
+}
+
+void varlen_fa_lse_seqlen_unpad(
+    torch::Tensor &in,          // [batch_size, num_heads, max_seq]
+    torch::Tensor &out,         // [num_heads, sum_seq]
+    torch::Tensor &accum_lens,  // [batch_size+1]
+    int sum_seq,
+    int max_seq,
+    int num_heads,
+    int batch_size
+) {
+    const int head_stride = in.stride(1);
+    const int head_out_stride = out.stride(0);
+
+    constexpr int BLOCK_THREADS = 256;
+
+    dim3 grid(1, num_heads, batch_size);
+    dim3 block(BLOCK_THREADS);
+
+    const at::musa::OptionalMUSAGuard device_guard(in.device());
+    auto stream = c10::musa::getCurrentMUSAStream().stream();
+
+    varlen_fa_lse_seqlen_unpad_kernel<BLOCK_THREADS><<<grid, block, 0, stream>>>(
+        in.data_ptr<float>(),
+        out.data_ptr<float>(),
+        accum_lens.data_ptr<int>(),
+        max_seq,
+        num_heads,
+        batch_size,
+        head_stride,
+        head_out_stride
+    );
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/layernorm_kernels.mu b/vllm_musa/csrc/layernorm_kernels.mu
new file mode 100644
index 000000000..6eb78175c
--- /dev/null
+++ b/vllm_musa/csrc/layernorm_kernels.mu
@@ -0,0 +1,193 @@
+#include "type_convert.muh"
+#include "dispatch_utils.h"
+
+#include "musa.h"
+
+namespace vllm_musa {
+
+// TODO(woosuk): Further optimize this kernel.
+template <typename scalar_t>
+__global__ void rms_norm_kernel(
+    scalar_t* __restrict__ out,           // [..., hidden_size]
+    const scalar_t* __restrict__ input,   // [..., hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float epsilon, const int num_tokens, const int hidden_size) {
+  __shared__ float s_variance;
+  float variance = 0.0f;
+
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    const float x = (float)input[blockIdx.x * hidden_size + idx];
+    variance += x * x;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    float x = (float)input[blockIdx.x * hidden_size + idx];
+    out[blockIdx.x * hidden_size + idx] =
+        ((scalar_t)(x * s_variance)) * weight[idx];
+  }
+}
+
+/* Function specialization in the case of FP16/BF16 tensors.
+   Additional optimizations we can make in this case are
+   packed and vectorized operations, which help with the
+   memory latency bottleneck. */
+template <typename scalar_t, int width>
+__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
+fused_add_rms_norm_kernel(
+    scalar_t* __restrict__ input,         // [..., hidden_size]
+    scalar_t* __restrict__ residual,      // [..., hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float epsilon, const int num_tokens, const int hidden_size) {
+  // Sanity checks on our vector struct and type-punned pointer arithmetic
+  static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
+  static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
+
+  const int vec_hidden_size = hidden_size / width;
+  __shared__ float s_variance;
+  float variance = 0.0f;
+  /* These and the argument pointers are all declared `restrict` as they are
+     not aliased in practice. Argument pointers should not be dereferenced
+     in this kernel as that would be undefined behavior */
+  auto* __restrict__ input_v =
+      reinterpret_cast<_f16Vec<scalar_t, width>*>(input);
+  auto* __restrict__ residual_v =
+      reinterpret_cast<_f16Vec<scalar_t, width>*>(residual);
+  auto* __restrict__ weight_v =
+      reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight);
+
+  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
+    int id = blockIdx.x * vec_hidden_size + idx;
+    _f16Vec<scalar_t, width> temp = input_v[id];
+    temp += residual_v[id];
+    variance += temp.sum_squares();
+    residual_v[id] = temp;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
+    int id = blockIdx.x * vec_hidden_size + idx;
+    _f16Vec<scalar_t, width> temp = residual_v[id];
+    temp *= s_variance;
+    temp *= weight_v[idx];
+    input_v[id] = temp;
+  }
+}
+
+/* Generic fused_add_rms_norm_kernel
+   The width field is not used here but necessary for other specializations.
+ */
+template <typename scalar_t, int width>
+__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
+fused_add_rms_norm_kernel(
+    scalar_t* __restrict__ input,         // [..., hidden_size]
+    scalar_t* __restrict__ residual,      // [..., hidden_size]
+    const scalar_t* __restrict__ weight,  // [hidden_size]
+    const float epsilon, const int num_tokens, const int hidden_size) {
+  __shared__ float s_variance;
+  float variance = 0.0f;
+
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    scalar_t z = input[blockIdx.x * hidden_size + idx];
+    z += residual[blockIdx.x * hidden_size + idx];
+    float x = (float)z;
+    variance += x * x;
+    residual[blockIdx.x * hidden_size + idx] = z;
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStore;
+  variance = BlockReduce(reduceStore).Reduce(variance, cub::Sum{}, blockDim.x);
+
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / hidden_size + epsilon);
+  }
+  __syncthreads();
+
+  for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
+    float x = (float)residual[blockIdx.x * hidden_size + idx];
+    input[blockIdx.x * hidden_size + idx] =
+        ((scalar_t)(x * s_variance)) * weight[idx];
+  }
+}
+
+}  // namespace vllm_musa
+
+void rms_norm(torch::Tensor& out,     // [..., hidden_size]
+              torch::Tensor& input,   // [..., hidden_size]
+              torch::Tensor& weight,  // [hidden_size]
+              double epsilon) {
+  int hidden_size = input.size(-1);
+  int num_tokens = input.numel() / hidden_size;
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min(hidden_size, 1024));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
+    vllm_musa::rms_norm_kernel<scalar_t><<<grid, block, 0, stream>>>(
+        out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
+        weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
+  });
+}
+
+#define LAUNCH_FUSED_ADD_RMS_NORM(width)                                       \
+  VLLM_DISPATCH_FLOATING_TYPES(                                                \
+      input.scalar_type(), "fused_add_rms_norm_kernel", [&] {                  \
+        vllm_musa::fused_add_rms_norm_kernel<scalar_t, width>                       \
+            <<<grid, block, 0, stream>>>(input.data_ptr<scalar_t>(),           \
+                                         residual.data_ptr<scalar_t>(),        \
+                                         weight.data_ptr<scalar_t>(), epsilon, \
+                                         num_tokens, hidden_size);             \
+      });
+
+void fused_add_rms_norm(torch::Tensor& input,     // [..., hidden_size]
+                        torch::Tensor& residual,  // [..., hidden_size]
+                        torch::Tensor& weight,    // [hidden_size]
+                        double epsilon) {
+  int hidden_size = input.size(-1);
+  int num_tokens = input.numel() / hidden_size;
+
+  dim3 grid(num_tokens);
+  /* This kernel is memory-latency bound in many scenarios.
+     When num_tokens is large, a smaller block size allows
+     for increased block occupancy on CUs and better latency
+     hiding on global mem ops. */
+  const int max_block_size = (num_tokens < 256) ? 1024 : 256;
+  dim3 block(std::min(hidden_size, max_block_size));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  /*If the tensor types are FP16/BF16, try to use the optimized kernel
+    with packed + vectorized ops.
+    Max optimization is achieved with a width-8 vector of FP16/BF16s
+    since we can load at most 128 bits at once in a global memory op.
+    However, this requires each tensor's data to be aligned to 16
+    bytes.
+   */
+  auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
+  auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
+  auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
+  bool ptrs_are_aligned =
+      inp_ptr % 16 == 0 && res_ptr % 16 == 0 && wt_ptr % 16 == 0;
+  if (ptrs_are_aligned && hidden_size % 8 == 0) {
+    LAUNCH_FUSED_ADD_RMS_NORM(8);
+  } else {
+    LAUNCH_FUSED_ADD_RMS_NORM(0);
+  }
+}
diff --git a/vllm_musa/csrc/layernorm_kernels_musa.mu b/vllm_musa/csrc/layernorm_kernels_musa.mu
new file mode 100644
index 000000000..a764d97ea
--- /dev/null
+++ b/vllm_musa/csrc/layernorm_kernels_musa.mu
@@ -0,0 +1,589 @@
+/* Copyright @2020-2024 Moore Threads Technology Co., Ltd("Moore Threads"). All
+ * rights reserved.
+ *
+ * This software ("this software and its documentations" or "the software") is
+ * protected by Copyright and the information contained herein is confidential.
+ *
+ * The software contained herein is PROPRIETARY to Moore Threads and is being
+ * provided under the terms and conditions of a form of Moore Threads software
+ * license agreement by and between Moore Threads and Licensee ("License
+ * Agreement") or electronically accepted by Licensee. Notwithstanding any
+ * terms or conditions to the contrary in the License Agreement, copy or
+ * disclosure of the software to any third party without the express written
+ * consent of Moore Threads is prohibited.
+ *
+ * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE LICENSE
+ * AGREEMENT, MOORE THREADS MAKES NO REPRESENTATION ABOUT ANY WARRANTIES,
+ * INCLUDING BUT NOT LIMITED TO THE SUITABILITY OF THE SOFTWARE FOR ANY
+ * PURPOSE. IT IS PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF
+ * ANY KIND. MOORE THREADS DISCLAIMS ALL WARRANTIES WITH REGARD TO THE
+ * SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
+ * NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
+ * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
+ * LICENSE AGREEMENT, IN NO EVENT SHALL MOORE THREADS BE LIABLE FOR ANY
+ * SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
+ * DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
+ * WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
+ * ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
+ * OF THE SOFTWARE.
+ */
+
+ #include "musa.h"
+ #include <iostream>
+ #include <vector>
+ #include <cmath>
+ #include <musa_runtime.h>
+ #include <musa_fp16.h>
+ #include "musa_bf16.h"
+ #include <musa_robust.h>
+ #include <torch/torch.h>
+
+
+ #define CHECK_MUSA(call)                                                      \
+   {                                                                           \
+     musaError_t err = call;                                                   \
+     if (err != musaSuccess) {                                                 \
+       std::cerr << "MUSA error in " << __FILE__ << ":" << __LINE__ << " - "  \
+                 << musaGetErrorString(err) << std::endl;                      \
+       exit(1);                                                                \
+     }                                                                         \
+   }
+
+ typedef __half float16_t;
+ typedef __mt_bfloat16 bfloat16_t;
+
+ #define DEVICE_INLINE __device__ __forceinline__
+ template <typename T, int width>
+ __device__ __forceinline__ T mudnn_shfl_down_sync(T val, unsigned int delta) {
+ // #if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 220))  // MUSIFY_EXCL_LINE
+   return __shfl_down_sync(0xffffffff, val, delta, width);
+ }
+
+ __device__ __host__ __forceinline__ constexpr int ceil_div(int a, int b) {
+   return (a + b - 1) / b;
+ }
+
+ __device__ __host__ __forceinline__ constexpr int64_t ceil_div(int64_t a,
+                                                                int64_t b) {
+   return (a + b - 1) / b;
+ }
+
+ #define WARP_THREADS 32
+ #define SMEM_STOP (WARP_THREADS / 2)
+ #define SHFL_START min(WARP_THREADS / 2, BLOCK_X / 2)
+ #define __SYNCTHREADS_LM __syncthreads_lm()
+ #define MACRO_UNROLL _Pragma("unroll")
+ #define LD_BYP_SLC(_BITS, _BYTES)                                   \
+   VecType dst;                                                      \
+   const BaseType* addr = ptr + idx;                                 \
+   asm volatile("LSU.LD.B" #_BITS " %0, %1, _, " #_BYTES             \
+                ", 1, 1, inner_persist=0, "                          \
+                "outer_persist=2, chrnt=l2_l3, slc=byp, persist=0, " \
+                "stride_add_first=0"                                 \
+                : "=R"(dst)                                          \
+                : "R"(addr));                                        \
+   return dst;
+
+     #define ATTR_ALIGNED(v) __attribute__((aligned(v)))
+     #define SELF_VEC_DEF(BASE_TYPE, VEC_TYPE_V2, VEC_TYPE_V4)                  \
+       struct ATTR_ALIGNED(sizeof(BASE_TYPE) * 2) VEC_TYPE_V2 {                 \
+         __device__ VEC_TYPE_V2() {}                                            \
+         __device__ VEC_TYPE_V2(const VEC_TYPE_V2& t) {                         \
+           this->x = t.x;                                                       \
+           this->y = t.y;                                                       \
+         }                                                                      \
+         BASE_TYPE x, y;                                                        \
+       };                                                                       \
+                                                                                \
+       __device__ __forceinline__ VEC_TYPE_V2 make_##VEC_TYPE_V2(BASE_TYPE x,   \
+                                                                 BASE_TYPE y) { \
+         VEC_TYPE_V2 t;                                                         \
+         t.x = x, t.y = y;                                                      \
+         return t;                                                              \
+       }                                                                        \
+                                                                                \
+       struct ATTR_ALIGNED(sizeof(BASE_TYPE) * 4) VEC_TYPE_V4 {                 \
+         __device__ VEC_TYPE_V4() {}                                            \
+         __device__ VEC_TYPE_V4(const VEC_TYPE_V4& t) {                         \
+           this->x = t.x;                                                       \
+           this->y = t.y;                                                       \
+           this->z = t.z;                                                       \
+           this->w = t.w;                                                       \
+         }                                                                      \
+         BASE_TYPE x, y, z, w;                                                  \
+       };                                                                       \
+                                                                                \
+       __device__ __forceinline__ VEC_TYPE_V4 make_##VEC_TYPE_V4(               \
+           BASE_TYPE x, BASE_TYPE y, BASE_TYPE z, BASE_TYPE w) {                \
+         VEC_TYPE_V4 t;                                                         \
+         t.x = x, t.y = y, t.z = z, t.w = w;                                    \
+         return t;                                                              \
+       }
+
+     SELF_VEC_DEF(float16_t, Half2, Half4)
+     // #if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+     SELF_VEC_DEF(bfloat16_t, Bhalf2, Bhalf4)
+     // #endif
+
+     #define GEN_VECTYPE(_CTYPE, _VECTYPE, _BYTES, _VLEN) \
+       struct ATTR_ALIGNED(_BYTES) _VECTYPE {             \
+         __device__ _VECTYPE() {}                         \
+         __device__ _VECTYPE(const _VECTYPE& t) {         \
+           MACRO_UNROLL                                   \
+           for (int i = 0; i < _VLEN; i++) {              \
+             this->arr[i] = t.arr[i];                     \
+           }                                              \
+         }                                                \
+         _CTYPE arr[_VLEN];                               \
+       }
+
+ GEN_VECTYPE(float16_t, Half8, 16, 8);
+ GEN_VECTYPE(bfloat16_t, Bhalf8, 16, 8);
+ template <typename type>
+ class Dtype;
+
+ #define INST(_type, _vec2, _vec4)                                        \
+   template <>                                                            \
+   class Dtype<_type> {                                                   \
+    public:                                                               \
+     using Scalar = _type;                                                \
+     using Vec2 = _vec2;                                                  \
+     using Vec4 = _vec4;                                                  \
+     static __device__ __forceinline__ Vec2 make_vec2(_type x, _type y) { \
+       return make_##_vec2(x, y);                                         \
+     }                                                                    \
+     static __device__ __forceinline__ Vec4 make_vec4(_type x, _type y,   \
+                                                      _type z, _type w) { \
+       return make_##_vec4(x, y, z, w);                                   \
+     }                                                                    \
+   }
+
+ INST(float, float2, float4);
+ INST(bfloat16_t, Bhalf2, Bhalf4);
+
+ template <typename T, int bits = 16 * 8>
+ struct VecType;
+
+ template <typename T>
+ struct DeduceVectorizedType {
+   using Type = T;
+ };
+
+ template <>
+ struct DeduceVectorizedType<half> {
+   using Type = _Float16;
+ };
+
+ template <>
+ struct DeduceVectorizedType<bfloat16_t> {
+   using Type = _Float16;
+ };
+
+ #define DEF_VECT(_CTYPE, _VECTYPE)                                            \
+     template <>                                                               \
+     struct VecType<_CTYPE, sizeof(_VECTYPE) * 8> {                            \
+     static constexpr int vec_bytes = sizeof(_VECTYPE);                        \
+     static constexpr int bit_per_byte = 8;                                    \
+     using BaseType = _CTYPE;                                                  \
+     using RobustTypePtr = __musa::robust_ptr<_CTYPE>;                         \
+     using Ttype = _VECTYPE;                                                   \
+     static constexpr int bits = vec_bytes * bit_per_byte;                     \
+     static constexpr int vlen = bits / (sizeof(BaseType) * bit_per_byte);     \
+     using VectorizedType = typename DeduceVectorizedType<BaseType>::Type;     \
+     typedef VectorizedType VxTy __attribute__((vector_size(vec_bytes)));      \
+     template <typename OffsetType>                                            \
+     static __device__ __forceinline__ VecType load(const BaseType* ptr,       \
+                                                     OffsetType idx) {         \
+         return *(VecType*)(ptr + idx);                                        \
+     }                                                                         \
+     template <typename OffsetType>                                            \
+     static __device__ __forceinline__ VecType                                 \
+     load_byp_slc(const BaseType* ptr, OffsetType idx) {                       \
+         if constexpr (vec_bytes == 16) {                                      \
+         LD_BYP_SLC(128, 16);                                                  \
+         } else if constexpr (vec_bytes == 8) {                                \
+         LD_BYP_SLC(64, 8);                                                    \
+         } else if constexpr (vec_bytes == 4) {                                \
+         LD_BYP_SLC(32, 4);                                                    \
+         } else if constexpr (vec_bytes == 2) {                                \
+         LD_BYP_SLC(16, 2);                                                    \
+         } else {                                                              \
+         LD_BYP_SLC(8, 1);                                                     \
+         }                                                                     \
+     }                                                                         \
+     template <typename OffsetType>                                            \
+     static __device__ __forceinline__ VecType                                 \
+     robust_load(const RobustTypePtr ptr, OffsetType idx) {                    \
+         return __musa::robust_load<VecType, BaseType>(ptr, idx);              \
+     }                                                                         \
+                                                                               \
+     template <typename OffsetType>                                            \
+     static __device__ __forceinline__ void store(BaseType* ptr,               \
+                                                     OffsetType idx,           \
+                                                     const VecType& dst) {     \
+         *(VecType*)(ptr + idx) = dst;                                         \
+     }                                                                         \
+     template <typename OffsetType>                                            \
+     static __device__ __forceinline__ void robust_store(RobustTypePtr ptr,    \
+                                                         OffsetType idx,       \
+                                                         const VecType& dst) { \
+         __musa::robust_store<VecType, BaseType>(dst, ptr, idx);               \
+     }                                                                         \
+                                                                               \
+     __device__ VecType() {                                                    \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         this->val_.elem[i] = 0;                                               \
+         }                                                                     \
+     }                                                                         \
+     __device__ VecType(const VecType& t) {                                    \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         this->val_.elem[i] = t.val_.elem[i];                                  \
+         }                                                                     \
+     }                                                                         \
+     __device__ VecType& operator=(const VecType& t) {                         \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         this->val_.elem[i] = t.val_.elem[i];                                  \
+         }                                                                     \
+         return *this;                                                         \
+     }                                                                         \
+     __device__ VecType(_CTYPE val) {                                          \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         this->val_.elem[i] = val;                                             \
+         }                                                                     \
+     }                                                                         \
+     template <typename SrcVecType>                                            \
+     friend __device__ VecType operator+(VecType lhs, const SrcVecType& rhs) { \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         lhs.val_.elem[i] += static_cast<BaseType>(rhs.val_.elem[i]);          \
+         }                                                                     \
+         return lhs;                                                           \
+     }                                                                         \
+     friend __device__ VecType operator+(VecType lhs, const _CTYPE& rhs) {     \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         lhs.val_.elem[i] += rhs;                                              \
+         }                                                                     \
+         return lhs;                                                           \
+     }                                                                         \
+     friend __device__ VecType operator-(VecType lhs, const VecType& rhs) {    \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         lhs.val_.elem[i] -= rhs.val_.elem[i];                                 \
+         }                                                                     \
+         return lhs;                                                           \
+     }                                                                         \
+     friend __device__ VecType operator*(VecType lhs, const VecType& rhs) {    \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         lhs.val_.elem[i] *= rhs.val_.elem[i];                                 \
+         }                                                                     \
+         return lhs;                                                           \
+     }                                                                         \
+     template <typename Func>                                                  \
+     __device__ VecType& apply() {                                             \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         this->val_.elem[i] = Func::apply(this->val_.elem[i]);                 \
+         }                                                                     \
+         return *this;                                                         \
+     }                                                                         \
+     template <typename SrcVecType>                                            \
+     static __device__ VecType cvt(const SrcVecType& src) {                    \
+         VecType dst;                                                          \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+         dst.val_.elem[i] = (BaseType)(src.val_.elem[i]);                      \
+         }                                                                     \
+         return dst;                                                           \
+     }                                                                         \
+     union U {                                                                 \
+         __device__ U() {                                                      \
+         MACRO_UNROLL                                                          \
+         for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {          \
+             this->elem[i] = 0;                                                \
+         }                                                                     \
+         }                                                                     \
+         Ttype storage;                                                        \
+         BaseType elem[sizeof(Ttype) / sizeof(BaseType)];                      \
+         VxTy vt_elem;                                                         \
+     };                                                                        \
+     U val_{};                                                                 \
+     }
+ DEF_VECT(float16_t, float16_t);
+ DEF_VECT(float16_t, Half2);
+ DEF_VECT(bfloat16_t, bfloat16_t);
+ DEF_VECT(bfloat16_t, Bhalf2);
+ DEF_VECT(bfloat16_t, Bhalf8);
+ DEF_VECT(float16_t, Half8);
+ DEF_VECT(float, float4);
+
+ enum class VarUpdateMode { WELFORD, WELFORD_ONLY_MEAN, CHAN, CHAN_ONLY_MEAN };
+
+ static __device__ __forceinline__ float fast_rcpf(float x) {
+   float y = __frcp_rn(x);
+   y = y * (2.0 - x * y);
+   return y;
+ }
+
+ static __device__ __forceinline__ float fast_divf(float a, float b) {
+   return a * fast_rcpf(b);
+ }
+
+ static __device__ __forceinline__ float fast_rsqrtf(float a) {
+   float x = 0.5 * a;
+   float y = __frsqrt_rn(a);
+   y = y * (1.5 - x * y * y);
+   return y;
+ }
+
+ template <typename T, VarUpdateMode Mode>
+ struct VarUpdate;
+
+ template <typename T>
+ struct VarUpdate<T, VarUpdateMode::WELFORD_ONLY_MEAN> {
+     DEVICE_INLINE void apply(T curr, T* mu, T* cnt) {
+     *cnt += 1;
+     T delta = curr - *mu;
+     *mu += fast_divf(delta, *cnt);
+     }
+ };
+
+ template <typename T>
+ struct VarUpdate<T, VarUpdateMode::CHAN_ONLY_MEAN> {
+   DEVICE_INLINE void apply(T mu_B, T cnt_B, T* mu, T* cnt) {
+     if (cnt_B > 0) {
+       T n_AB = cnt_B + (*cnt);
+       T delta = mu_B - (*mu);
+       *mu += delta * fast_divf(cnt_B, n_AB);
+       *cnt = n_AB;
+     }
+   }
+ };
+
+ template <typename ComputeType, int BLOCK_X, int BLOCK_Y,int Vlen>
+ struct AllReduceOp {
+   DEVICE_INLINE void apply(ComputeType* sum, int tx, int ty) {
+     __shared__ ComputeType __attribute__((aligned(16)))
+     smem[BLOCK_X * BLOCK_Y * Vlen];
+     ComputeType* smem_sum = &smem[0];
+
+       static_assert(Vlen == 1,
+                     "Axis COLUMN doesn't support vlen greater than 1");
+ #pragma unroll
+       for (int offset = BLOCK_X / 2; offset > SMEM_STOP; offset /= 2) {
+         if (tx >= offset && tx < 2 * offset) {
+           smem_sum[ty * BLOCK_X + tx] = *sum;
+         }
+         __SYNCTHREADS_LM;
+         if (tx < offset) {
+           *sum += smem_sum[ty * BLOCK_X + tx + offset];
+         }
+       }
+ #if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ >= 220))
+ #pragma unroll
+       for (int offset = SHFL_START; offset > 0; offset /= 2) {
+         *sum += mudnn_shfl_down_sync<ComputeType, 32>(*sum, offset);
+       }
+ #endif
+       if (tx == 0) {
+         smem_sum[ty * BLOCK_X + tx] = *sum;
+       }
+       __SYNCTHREADS_LM;
+       *sum = smem_sum[ty * BLOCK_X];
+   }
+ };
+
+  template <typename SrcDtype, typename ComputeType, int BLOCK_X, int BLOCK_Y, int vlen>
+  __global__ void LayerNormGlobalKernelVlen(
+     SrcDtype* input, SrcDtype* residual, const SrcDtype* weight,
+      const size_t M, const size_t N, const ComputeType eps) {
+    size_t tx = threadIdx.x;
+    size_t ty = threadIdx.y;
+    size_t m_idx = blockIdx.x * blockDim.y + ty;
+    size_t n_idx = tx * vlen;
+    size_t n_step = (size_t)blockDim.x * vlen;
+
+    using SrcVec = VecType<SrcDtype, vlen * sizeof(SrcDtype) * 8>;
+
+    ComputeType var = 0;
+    const SrcDtype* __restrict p_src = input + m_idx * N;
+    SrcDtype* __restrict p_res = residual + m_idx * N;  // residual ptr
+
+    // TODO(wuke): use robust_load, robust_store
+    bool m_valid = m_idx < M;
+    if (m_valid) {
+      for (size_t j = n_idx; j < N; j += n_step) {
+        SrcVec curr, res_vec, fused_vec;
+        #if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 220))
+            curr = *(SrcVec *)(p_src+j);
+            res_vec = *(SrcVec *)(p_res+j);
+        #elif ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310))
+            curr = SrcVec::load_byp_slc(p_src, j);
+            res_vec = SrcVec::load_byp_slc(p_res, j);
+        #endif
+        #pragma unroll
+        for (int k = 0; k < vlen; k++) {
+         fused_vec.val_.elem[k] = curr.val_.elem[k] + res_vec.val_.elem[k];
+         var += (ComputeType)fused_vec.val_.elem[k] * (ComputeType)fused_vec.val_.elem[k];
+        }
+        *(SrcVec*)(p_res + j) = fused_vec;
+      }
+    }
+    AllReduceOp<ComputeType, BLOCK_X, BLOCK_Y, 1> all_reduce_op;
+    all_reduce_op.apply(&var, tx, ty);
+    if (m_valid) {
+      ComputeType inv_var = fast_rsqrtf(var / N + eps);
+      SrcDtype* __restrict p_dst = input + m_idx * N;
+      bool with_weight = (weight != NULL);
+      if (with_weight) {
+        for (size_t j = n_idx; j < N; j += n_step) {
+          SrcVec fused_vec, weight_val, dst;
+          #if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 220))
+            fused_vec = *(SrcVec *)(p_res+j);
+            weight_val = *(SrcVec *)(weight+j);
+          #elif ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310))
+            fused_vec = SrcVec::load_byp_slc(p_res, j);
+            weight_val = SrcVec::load_byp_slc(weight, j);
+          #endif
+  #pragma unroll
+          for (int k = 0; k < vlen; k++) {
+             dst.val_.elem[k] = (SrcDtype)((ComputeType)fused_vec.val_.elem[k] * inv_var *
+                             (ComputeType)weight_val.val_.elem[k]);
+          }
+          *(SrcVec*)(p_dst + j) = dst;
+        }
+      }
+    }
+  }
+
+ template <typename DstType, typename SrcType>
+ void golden_rmsnorm(DstType* output, const SrcType* input,
+                     const DstType* gamma, int M, int N, float epsilon) {
+     for (int i = 0; i < M; ++i) {
+         float square_sum = 0.0f;
+         for (int j = 0; j < N; ++j) {
+             float val = static_cast<float>(input[i * N + j]);
+             square_sum += val * val;
+         }
+         float rms = std::sqrt(square_sum / N + epsilon);
+
+         for (int j = 0; j < N; ++j) {
+             float val = static_cast<float>(input[i * N + j]) / rms;
+             float scaled = val * float(gamma[j]);
+             output[i * N + j] = static_cast<DstType>(scaled);
+         }
+     }
+ }
+
+ template <typename DstType, typename SrcType>
+ void golden_fused_add_rmsnorm(
+   DstType* output, const SrcType* input,const SrcType* residual,
+   const DstType* weight, size_t num_tokens,size_t hidden_size,float epsilon) {
+
+   for (size_t token_idx = 0; token_idx < num_tokens; ++token_idx) {
+     float sum_squares = 0.0f;
+
+     // 1. Fused input = input + residual
+     std::vector<float> fused(hidden_size);
+     for (size_t h = 0; h < hidden_size; ++h) {
+       float val = static_cast<float>(input[token_idx * hidden_size + h]) +
+                   static_cast<float>(residual[token_idx * hidden_size + h]);
+       fused[h] = val;
+       sum_squares += val * val;
+     }
+
+     // 2. RMSNorm
+     float rms = std::sqrt(sum_squares / hidden_size + epsilon);
+     float inv_rms = 1.0f / rms;
+
+     // 3. Apply scale and bias (if present)
+     for (size_t h = 0; h < hidden_size; ++h) {
+       float normed = fused[h] * inv_rms;
+       if (!weight) normed *= static_cast<float>(weight[h]);
+       output[token_idx * hidden_size + h] = static_cast<DstType>(normed);
+     }
+   }
+ }
+
+
+ #define CALL_KERN(_SRC_DTYPE,_KERN, _BLKX, _BLKY, _VLEN)                      \
+ {                                                                             \
+     const uint32_t block_x = _BLKX;                                           \
+     const uint32_t block_y = _BLKY;                                           \
+     const uint32_t nr_blocks = ceil_div(m, (size_t)block_y);                  \
+     dim3 block_size{block_x, block_y, 1};                                     \
+     dim3 grid_size{nr_blocks, 1, 1};                                          \
+     LayerNorm##_KERN##KernelVlen<_SRC_DTYPE, float,                           \
+                                      block_x, block_y, _VLEN>                 \
+             <<<grid_size, block_size, 0, stream>>>(                           \
+                 static_cast<_SRC_DTYPE*>(input),                              \
+                 static_cast<_SRC_DTYPE*>(residual),                           \
+                 static_cast<_SRC_DTYPE*>(weight),                             \
+                 m, n, static_cast<float>(epsilon));                           \
+ }
+
+ #define DISPATCH_KERNEL(_KERN, _BLKX, _BLKY)                                  \
+    if constexpr (std::is_same_v<SrcDtype,float16_t>) {                        \
+       CALL_KERN(float16_t, _KERN, _BLKX, _BLKY, 8);                           \
+    } else if constexpr (std::is_same_v<SrcDtype, bfloat16_t>) {               \
+      CALL_KERN(bfloat16_t, _KERN, _BLKX, _BLKY, 8);                                                                                          \
+    }
+
+ template <typename SrcDtype>
+ void rms_fused_add_rms_norm(SrcDtype* input, SrcDtype* residual, SrcDtype* weight, int m, int n, double epsilon) {
+   auto stream = c10::musa::getCurrentMUSAStream().stream();
+   const int vlen = 2;
+   DISPATCH_KERNEL(Global, 1024, 1);
+ }
+
+void musa_fused_add_rms_norm(
+    torch::Tensor &input,
+    torch::Tensor &residual,
+    torch::Tensor &weight,
+    double epsilon) {
+
+    int m = input.size(0);
+    int n = input.size(1);
+
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+
+    if (input.scalar_type() == at::ScalarType::BFloat16)
+    {
+        rms_fused_add_rms_norm<__mt_bfloat16>(
+            static_cast<__mt_bfloat16*>(input.data_ptr()),
+            static_cast<__mt_bfloat16*>(residual.data_ptr()),
+            static_cast<__mt_bfloat16*>(weight.data_ptr()),
+            m,
+            n,
+            epsilon);
+    }
+    else if (input.scalar_type() == at::ScalarType::Half)
+    {
+        rms_fused_add_rms_norm<__half>(
+            static_cast<__half*>(input.data_ptr()),
+            static_cast<__half*>(residual.data_ptr()),
+            static_cast<__half*>(weight.data_ptr()),
+            m,
+            n,
+            epsilon);
+    }
+    else if (input.scalar_type() == at::ScalarType::Float)
+    {
+        rms_fused_add_rms_norm<float>(
+            static_cast<float*>(input.data_ptr()),
+            static_cast<float*>(residual.data_ptr()),
+            static_cast<float*>(weight.data_ptr()),
+            m,
+            n,
+            epsilon);
+    }
+    else
+    {
+        TORCH_CHECK(false, "only support Float32, Half and BFloat16 dtype");
+    }
+}
diff --git a/vllm_musa/csrc/linear/ai_qint8.h b/vllm_musa/csrc/linear/ai_qint8.h
new file mode 100644
index 000000000..32a76ab2e
--- /dev/null
+++ b/vllm_musa/csrc/linear/ai_qint8.h
@@ -0,0 +1,55 @@
+/* Copyright @2020-2024 Moore Threads Technology Co., Ltd("Moore Threads"). All
+ * rights reserved.
+ *
+ * This software ("this software and its documentations" or "the software") is
+ * protected by Copyright and the information contained herein is confidential.
+ *
+ * The software contained herein is PROPRIETARY to Moore Threads and is being
+ * provided under the terms and conditions of a form of Moore Threads software
+ * license agreement by and between Moore Threads and Licensee ("License
+ * Agreement") or electronically accepted by Licensee. Notwithstanding any
+ * terms or conditions to the contrary in the License Agreement, copy or
+ * disclosure of the software to any third party without the express written
+ * consent of Moore Threads is prohibited.
+ *
+ * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE LICENSE
+ * AGREEMENT, MOORE THREADS MAKES NO REPRESENTATION ABOUT ANY WARRANTIES,
+ * INCLUDING BUT NOT LIMITED TO THE SUITABILITY OF THE SOFTWARE FOR ANY
+ * PURPOSE. IT IS PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF
+ * ANY KIND. MOORE THREADS DISCLAIMS ALL WARRANTIES WITH REGARD TO THE
+ * SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
+ * NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
+ * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
+ * LICENSE AGREEMENT, IN NO EVENT SHALL MOORE THREADS BE LIABLE FOR ANY
+ * SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
+ * DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
+ * WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
+ * ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
+ * OF THE SOFTWARE.
+ */
+#ifndef INCLUDE_MUDNN_API_AI_QINT8_H_
+#define INCLUDE_MUDNN_API_AI_QINT8_H_
+
+class qint8_t {
+  int8_t _;
+
+ public:
+  __device__ int8_t as_int8() const { return _; }
+
+  __host__ __device__ qint8_t(int8_t val) : _(val) {}
+  __host__ __device__ qint8_t() : _(0) {}
+  __host__ __device__ qint8_t(const qint8_t& b) { _ = b._; }
+  __host__ __device__ qint8_t(const int& b) { _ = b; }
+  __host__ __device__ qint8_t(const char& b) { _ = b; }
+  __host__ __device__ operator char() const { return _; }
+  __host__ __device__ qint8_t operator+=(const qint8_t& b) { return _ += b._; }
+  __host__ __device__ qint8_t operator-=(const qint8_t& b) { return _ -= b._; }
+  __host__ __device__ qint8_t operator*=(const qint8_t& b) { return _ *= b._; }
+
+  bool operator<(const qint8_t& b) const { return _ < b._; }
+  bool operator>(const qint8_t& b) const { return _ > b._; }
+  bool operator==(const qint8_t& b) const { return _ == b._; }
+  bool operator!=(const qint8_t& b) const { return _ != b._; }
+};
+
+#endif  // INCLUDE_MUDNN_API_AI_QINT8_H_
diff --git a/vllm_musa/csrc/linear/common.mu b/vllm_musa/csrc/linear/common.mu
new file mode 100644
index 000000000..4d657c8d3
--- /dev/null
+++ b/vllm_musa/csrc/linear/common.mu
@@ -0,0 +1,7 @@
+#include "common.muh"
+
+void LaunchKernel(std::function<void()> func, musaStream_t stream,
+                  const std::string name, const uint64_t bytes,
+                  const uint64_t operations, const uint64_t ddr_bytes) {
+  func();
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/linear/common.muh b/vllm_musa/csrc/linear/common.muh
new file mode 100644
index 000000000..ac0454637
--- /dev/null
+++ b/vllm_musa/csrc/linear/common.muh
@@ -0,0 +1,92 @@
+#pragma once
+#include "dtype.muh"
+#include <functional>
+#include <iostream>
+#include <musa_runtime.h>
+#include <sstream>
+#include <string>
+
+using namespace musa::dnn;
+
+void LaunchKernel(std::function<void()> func, musaStream_t stream,
+                  const std::string name, const uint64_t bytes,
+                  const uint64_t operations, const uint64_t ddr_bytes = 0);
+
+/// Returns the ceiling of (a / b)
+__device__ __host__ __forceinline__ constexpr int ceil_div(int a, int b) {
+  return (a + b - 1) / b;
+}
+
+template <typename T> __device__ __host__ inline T sigmoid(T x) {
+  return 1.f / (1.f + __expf(-x));
+}
+
+enum class AlignAttr { GENERAL, ALIGN_VEC };
+
+struct BatchGemvParam {
+  bool l_is_v{false};        // left tensor is regarded as vector
+  bool trans_mat{false};     // matrix transpose property
+  bool broadcast_v{false};   // vector need broadcast
+  bool broadcast_mat{false}; // matrix need broadcast
+  bool bias_scalar{false};   // fused bias is a scalar
+  int batch{1};
+  int mat_m{1};
+  int k{1};
+  int ldm{0};                // leading dimension of matrix
+  int inc_x{0};              // stride of each one-dimensional array vector
+  int inc_c{0};              // incy: stride of each one-dimensional array input
+  int inc_d{0};              // incy: stride of each one-dimensional array out
+  int64_t batch_stride_m{0}; // batch stride of matrix
+  int64_t batch_stride_x{0}; // batch stride of vector
+  int64_t batch_stride_c{0}; // batch stride of input
+  int64_t batch_stride_d{0}; // batch stride of out
+};
+
+template <typename T> struct InferComputeDtype {
+  using ComputeDtype = T;
+};
+
+template <> struct InferComputeDtype<float16_t> {
+  using ComputeDtype = float;
+};
+
+template <> struct InferComputeDtype<bfloat16_t> {
+  using ComputeDtype = float;
+};
+
+// template <> struct InferComputeDtype<fp8_e4m3_t> {
+//   using ComputeDtype = float;
+// };
+
+// template <> struct InferComputeDtype<fp8_e5m2_t> {
+//   using ComputeDtype = float;
+// };
+
+#define MUDNN_ATTR_INLINE inline __attribute__((always_inline))
+
+enum class StatusX {
+#define MUDNN_ITEM(x) x,
+#include "status.md"
+#undef MUDNN_ITEM
+};
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ == 310 // MUSIFY_EXCL_LINE
+#define __SYNCTHREADS_LM __syncthreads_lm()
+#else
+#define __SYNCTHREADS_LM __syncthreads()
+#endif
+
+#ifndef MUSA_CHECK
+#define MUSA_CHECK(expr)                                                       \
+  do {                                                                         \
+    musaError_t _err = (expr);                                                 \
+    if (_err != musaSuccess) {                                                 \
+      std::stringstream ss;                                                    \
+      ss << " musa runtime failed " << __FILE_NAME__ << ":" << __LINE__        \
+         << ": err " << _err << " = " << musaGetErrorString(_err)              \
+         << std::endl;                                                         \
+      std::cerr << ss.str() << std::endl;                                      \
+      exit(0);                                                                 \
+    }                                                                          \
+  } while (0)
+#endif
\ No newline at end of file
diff --git a/vllm_musa/csrc/linear/dtype.muh b/vllm_musa/csrc/linear/dtype.muh
new file mode 100644
index 000000000..7b0ac1965
--- /dev/null
+++ b/vllm_musa/csrc/linear/dtype.muh
@@ -0,0 +1,614 @@
+/* Copyright @2020-2024 Moore Threads Technology Co., Ltd("Moore Threads"). All
+ * rights reserved.
+ *
+ * This software ("this software and its documentations" or "the software") is
+ * protected by Copyright and the information contained herein is confidential.
+ *
+ * The software contained herein is PROPRIETARY to Moore Threads and is being
+ * provided under the terms and conditions of a form of Moore Threads software
+ * license agreement by and between Moore Threads and Licensee ("License
+ * Agreement") or electronically accepted by Licensee. Notwithstanding any
+ * terms or conditions to the contrary in the License Agreement, copy or
+ * disclosure of the software to any third party without the express written
+ * consent of Moore Threads is prohibited.
+ *
+ * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE LICENSE
+ * AGREEMENT, MOORE THREADS MAKES NO REPRESENTATION ABOUT ANY WARRANTIES,
+ * INCLUDING BUT NOT LIMITED TO THE SUITABILITY OF THE SOFTWARE FOR ANY
+ * PURPOSE. IT IS PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF
+ * ANY KIND. MOORE THREADS DISCLAIMS ALL WARRANTIES WITH REGARD TO THE
+ * SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
+ * NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
+ * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
+ * LICENSE AGREEMENT, IN NO EVENT SHALL MOORE THREADS BE LIABLE FOR ANY
+ * SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
+ * DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
+ * WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
+ * ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
+ * OF THE SOFTWARE.
+ */
+
+#pragma once
+#include <musa_runtime.h>
+#include <stdint.h>
+
+// #include "mudnn_config.h"
+// #include "mudnn/api/impl_base.h"
+#include "ai_qint8.h"
+#include "integer_subbyte.h"
+
+#ifdef __MTGPU__
+#include <musa_robust.h>
+#endif
+
+// //! NOTE: cpplint can not accept function with inline asm in namespace
+// __device__ inline int8_t gpu_pack_f32_int8(float x) {
+//   int8_t rst;
+// #if defined(__MUSA_ARCH__)
+//   asm("MOV.F32_TO_S8888_RTE %0 %1" : "=R"(rst) : "R"(x));
+// #else
+//   int i32 = (int)round(x);
+//   i32 = i32 > 127 ? 127 : i32;
+//   i32 = i32 < -128 ? -128 : i32;
+//   rst = (int8_t)i32;
+// #endif
+//   return rst;
+// }
+
+#include "musa_fp16.h"
+#include "musa_bf16.h"
+#include "musa_fp8.h"
+#include <torch/all.h>
+
+namespace musa {
+namespace dnn {
+
+template <typename T> struct sizeof_bits {
+  static constexpr size_t value = sizeof(T) * 8;
+};
+
+
+template <int Bits, bool Signed>
+struct sizeof_bits<integer_subbyte<Bits, Signed>> {
+  static constexpr size_t value = Bits;
+};
+
+template <typename T>
+static constexpr int sizeof_bits_v = sizeof_bits<T>::value;
+
+typedef __half float16_t;
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+typedef __mt_bfloat16 bfloat16_t;
+#else
+typedef int16_t bfloat16_t;
+#endif
+
+
+
+// #if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+// #define __MUSA_HOSTDEVICE__ __host__ __device__
+// template <typename T> struct FDtypeWrapper {
+//   T _x;
+//   __MUSA_HOSTDEVICE__ FDtypeWrapper() = default;
+//   __MUSA_HOSTDEVICE__ FDtypeWrapper(const float x) {
+//     if constexpr (std::is_same<__mt_fp8_e4m3, T>::value) {
+//       // hw instruction only support satfinity mode, add if to support normal
+//       // mode
+//       uint32_t f_abs_bits = (*(uint32_t *)&x) & (uint32_t)0x7fffffff;
+//       //
+//       https://github.com/pytorch/pytorch/blob/16de25b1dc3af7cc3efdb24e5637a62b8c329881/c10/util/Float8_e4m3fn.h#L183-L186
+//       constexpr uint32_t fp8_max = 0x43e80001;
+//       if (f_abs_bits >= fp8_max) {
+//         _x.__x = 0x7f;
+//         _x.__x |= UINT8_C(((*(uint32_t *)&x) & UINT32_C(0x80000000)) >> 24);
+//         //
+//       } else {
+//         _x = __mt_fp8_e4m3(x);
+//       }
+//     } else if constexpr (std::is_same<__mt_fp8_e5m2, T>::value) {
+//       // hw instruction only support satfinity mode, add if to support normal
+//       // mode
+//       uint32_t f_abs_bits = (*(uint32_t *)&x) & (uint32_t)0x7fffffff;
+//       //
+//       https://github.com/pytorch/pytorch/blob/16de25b1dc3af7cc3efdb24e5637a62b8c329881/c10/util/Float8_e5m2.h#L64
+//       constexpr uint32_t fp8_max = 0x47700000;
+//       constexpr uint32_t fp32_inf = UINT32_C(255) << 23;
+//       if (f_abs_bits >= fp8_max) {
+//         _x.__x = f_abs_bits > fp32_inf ? UINT8_C(0x7F) : UINT8_C(0x7C);
+//         _x.__x |= UINT8_C(((*(uint32_t *)&x) & UINT32_C(0x80000000)) >> 24);
+//       } else {
+//         _x = __mt_fp8_e5m2(x);
+//       }
+//     } else {
+//       _x = T(x);
+//     }
+//   }
+//   __MUSA_HOSTDEVICE__ operator float() const {
+//     if constexpr (std::is_same<__mt_fp8_e4m3, T>::value) {
+//       __half temp;
+//       asm("CVT.RTE.DENORM.BST1 %0.f16, %1.e4m3\n" : "=R"(temp) : "R"(_x));
+//       return (float)temp;
+//     } else if constexpr (std::is_same<__mt_fp8_e5m2, T>::value) {
+//       __half temp;
+//       asm("CVT.RTE.DENORM.BST1 %0.f16, %1.e5m2\n" : "=R"(temp) : "R"(_x));
+//       return (float)temp;
+//     } else {
+//       return (float)_x;
+//     }
+//   }
+//   __MUSA_HOSTDEVICE__ FDtypeWrapper &operator+=(const float x) {
+//     _x = T((float)_x + x);
+//     return *this;
+//   }
+//   __MUSA_HOSTDEVICE__ FDtypeWrapper &operator-=(const float x) {
+//     _x = T((float)_x - x);
+//     return *this;
+//   }
+//   __MUSA_HOSTDEVICE__ FDtypeWrapper &operator*=(const float x) {
+//     _x = T((float)_x * x);
+//     return *this;
+//   }
+//   __MUSA_HOSTDEVICE__ FDtypeWrapper &operator/=(const float x) {
+//     _x = T((float)_x / x);
+//     return *this;
+//   }
+// };
+// #undef __MUSA_HOSTDEVICE__
+// typedef FDtypeWrapper<__mt_fp8_e4m3> fp8_e4m3_t;
+// typedef FDtypeWrapper<__mt_fp8_e5m2> fp8_e5m2_t;
+// #else
+// typedef int8_t fp8_e4m3_t;
+// typedef uint8_t fp8_e5m2_t;
+// #endif
+// #define MUDNN_ITEM(x, y, z, w, t) \
+//   template <> struct TensorImpl::DtypeTraits<TensorImpl::Type::x> { \
+//     using Ctype = w; \
+//   };
+// #include "data_type.md"
+// #undef MUDNN_ITEM
+
+#define MACRO_UNROLL _Pragma("unroll")
+
+#define ATTR_ALIGNED(v) __attribute__((aligned(v)))
+#define SELF_VEC_DEF(BASE_TYPE, VEC_TYPE_V2, VEC_TYPE_V4)                      \
+  struct ATTR_ALIGNED(sizeof(BASE_TYPE) * 2) VEC_TYPE_V2 {                     \
+    __device__ VEC_TYPE_V2() {}                                                \
+    __device__ VEC_TYPE_V2(const VEC_TYPE_V2 &t) {                             \
+      this->x = t.x;                                                           \
+      this->y = t.y;                                                           \
+    }                                                                          \
+    BASE_TYPE x, y;                                                            \
+  };                                                                           \
+                                                                               \
+  __device__ __forceinline__ VEC_TYPE_V2 make_##VEC_TYPE_V2(BASE_TYPE x,       \
+                                                            BASE_TYPE y) {     \
+    VEC_TYPE_V2 t;                                                             \
+    t.x = x, t.y = y;                                                          \
+    return t;                                                                  \
+  }                                                                            \
+                                                                               \
+  struct ATTR_ALIGNED(sizeof(BASE_TYPE) * 4) VEC_TYPE_V4 {                     \
+    __device__ VEC_TYPE_V4() {}                                                \
+    __device__ VEC_TYPE_V4(const VEC_TYPE_V4 &t) {                             \
+      this->x = t.x;                                                           \
+      this->y = t.y;                                                           \
+      this->z = t.z;                                                           \
+      this->w = t.w;                                                           \
+    }                                                                          \
+    BASE_TYPE x, y, z, w;                                                      \
+  };                                                                           \
+                                                                               \
+  __device__ __forceinline__ VEC_TYPE_V4 make_##VEC_TYPE_V4(                   \
+      BASE_TYPE x, BASE_TYPE y, BASE_TYPE z, BASE_TYPE w) {                    \
+    VEC_TYPE_V4 t;                                                             \
+    t.x = x, t.y = y, t.z = z, t.w = w;                                        \
+    return t;                                                                  \
+  }
+
+SELF_VEC_DEF(float16_t, Half2, Half4)
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+SELF_VEC_DEF(bfloat16_t, Bhalf2, Bhalf4)
+#endif
+// #if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+// SELF_VEC_DEF(fp8_e5m2_t, E5m2x2, E5m2x4)
+// SELF_VEC_DEF(fp8_e4m3_t, E4m3x2, E4m3x4)
+// #endif
+
+#define GEN_VECTYPE(_CTYPE, _VECTYPE, _BYTES, _VLEN)                           \
+  struct ATTR_ALIGNED(_BYTES) _VECTYPE {                                       \
+    __device__ _VECTYPE() {}                                                   \
+    __device__ _VECTYPE(const _VECTYPE &t) {                                   \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < _VLEN; i++) {                                        \
+        this->arr[i] = t.arr[i];                                               \
+      }                                                                        \
+    }                                                                          \
+    _CTYPE arr[_VLEN];                                                         \
+  }
+GEN_VECTYPE(float16_t, Half8, 16, 8);
+GEN_VECTYPE(signed char, Char8, 8, 8);
+GEN_VECTYPE(uint8_t, Uint8, 8, 8);
+GEN_VECTYPE(int16_t, Short8, 16, 8);
+GEN_VECTYPE(int16_t, Short16, 32, 16);
+GEN_VECTYPE(uint16_t, Ushort8, 16, 8);
+GEN_VECTYPE(uint32_t, UInt8, 32, 8);
+GEN_VECTYPE(uint32_t, UInt16, 64, 16);
+GEN_VECTYPE(uint64_t, ULong8, 64, 8);
+GEN_VECTYPE(uint64_t, ULong16, 128, 16);
+GEN_VECTYPE(float, Float8, 32, 8);
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+GEN_VECTYPE(bfloat16_t, Bhalf8, 16, 8);
+GEN_VECTYPE(bfloat16_t, Bhalf16, 32, 16);
+GEN_VECTYPE(bfloat16_t, Bhalf32, 64, 32);
+#endif
+// #if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+// GEN_VECTYPE(fp8_e4m3_t, E4m3x8, 8, 8);
+// GEN_VECTYPE(fp8_e4m3_t, E4m3x16, 16, 16);
+// GEN_VECTYPE(fp8_e4m3_t, E4m3x32, 32, 32);
+// GEN_VECTYPE(fp8_e5m2_t, E5m2x8, 8, 8);
+// GEN_VECTYPE(fp8_e5m2_t, E5m2x16, 16, 16);
+// GEN_VECTYPE(fp8_e5m2_t, E5m2x32, 32, 32);
+// #endif
+GEN_VECTYPE(int32_t, Int8, 32, 8);
+GEN_VECTYPE(int64_t, Long8, 64, 8);
+GEN_VECTYPE(int64_t, Long16, 128, 16);
+GEN_VECTYPE(signed char, Char16, 16, 16);
+GEN_VECTYPE(float16_t, Half16, 32, 16);
+GEN_VECTYPE(float, Float16, 64, 16);
+GEN_VECTYPE(int32_t, Int16, 64, 16);
+
+template <typename type> class Dtype;
+
+#define INST(_type, _vec2, _vec4)                                              \
+  template <> class Dtype<_type> {                                             \
+  public:                                                                      \
+    using Scalar = _type;                                                      \
+    using Vec2 = _vec2;                                                        \
+    using Vec4 = _vec4;                                                        \
+    static __device__ __forceinline__ Vec2 make_vec2(_type x, _type y) {       \
+      return make_##_vec2(x, y);                                               \
+    }                                                                          \
+    static __device__ __forceinline__ Vec4 make_vec4(_type x, _type y,         \
+                                                     _type z, _type w) {       \
+      return make_##_vec4(x, y, z, w);                                         \
+    }                                                                          \
+  }
+
+INST(float, float2, float4);
+INST(float16_t, Half2, Half4);
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+INST(bfloat16_t, Bhalf2, Bhalf4);
+#endif
+// #if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+// INST(fp8_e5m2_t, E5m2x2, E5m2x4);
+// INST(fp8_e4m3_t, E4m3x2, E4m3x4);
+// #endif
+INST(int32_t, int2, int4);
+INST(uint32_t, uint2, uint4);
+INST(int8_t, char2, char4);
+INST(uint8_t, uchar2, uchar4);
+INST(int16_t, short2, short4);
+INST(uint16_t, ushort2, ushort4);
+INST(int64_t, long2, long4);
+INST(uint64_t, ulong2, ulong4);
+INST(double, double2, double4);
+
+#undef INST
+
+template <typename T> struct DeduceVectorizedType {
+  using Type = T;
+};
+
+template <> struct DeduceVectorizedType<bool> {
+  using Type = int8_t;
+};
+
+template <> struct DeduceVectorizedType<half> {
+  using Type = _Float16;
+};
+
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+template <> struct DeduceVectorizedType<bfloat16_t> {
+  using Type = _Float16;
+};
+#endif
+// #if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+// template <> struct DeduceVectorizedType<fp8_e4m3_t> {
+//   using Type = int8_t;
+// };
+// template <> struct DeduceVectorizedType<fp8_e5m2_t> {
+//   using Type = int8_t;
+// };
+// #endif
+template <> struct DeduceVectorizedType<qint8_t> {
+  using Type = int8_t;
+};
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ == 310 // MUSIFY_EXCL_LINE
+#define LD_BYP_SLC(_BITS, _BYTES)                                              \
+  VecType dst;                                                                 \
+  const BaseType *addr = ptr + idx;                                            \
+  asm volatile("LSU.LD.B" #_BITS " %0, %1, _, " #_BYTES                        \
+               ", 1, 1, inner_persist=0, "                                     \
+               "outer_persist=2, chrnt=l2_l3, slc=byp, persist=0, "            \
+               "stride_add_first=0"                                            \
+               : "=R"(dst)                                                     \
+               : "R"(addr));                                                   \
+  return dst;
+#else
+#define LD_BYP_SLC(_BITS, _BYTES) return *(VecType *)(ptr + idx);
+#endif
+
+template <typename T, int bits = 16 * 8> struct VecType;
+
+#define DEF_VECT(_CTYPE, _VECTYPE)                                             \
+  template <> struct VecType<_CTYPE, sizeof(_VECTYPE) * 8> {                   \
+    static constexpr int vec_bytes = sizeof(_VECTYPE);                         \
+    static constexpr int bit_per_byte = 8;                                     \
+    using BaseType = _CTYPE;                                                   \
+    using RobustTypePtr = __musa::robust_ptr<_CTYPE>;                          \
+    using Ttype = _VECTYPE;                                                    \
+    static constexpr int bits = vec_bytes * bit_per_byte;                      \
+    static constexpr int vlen = bits / (sizeof(BaseType) * bit_per_byte);      \
+    using VectorizedType = typename DeduceVectorizedType<BaseType>::Type;      \
+    typedef VectorizedType VxTy __attribute__((vector_size(vec_bytes)));       \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ VecType load(const BaseType *ptr,        \
+                                                   OffsetType idx) {           \
+      return *(VecType *)(ptr + idx);                                          \
+    }                                                                          \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ VecType                                  \
+    load_byp_slc(const BaseType *ptr, OffsetType idx) {                        \
+      if constexpr (vec_bytes == 16) {                                         \
+        LD_BYP_SLC(128, 16);                                                   \
+      } else if constexpr (vec_bytes == 8) {                                   \
+        LD_BYP_SLC(64, 8);                                                     \
+      } else if constexpr (vec_bytes == 4) {                                   \
+        LD_BYP_SLC(32, 4);                                                     \
+      } else if constexpr (vec_bytes == 2) {                                   \
+        LD_BYP_SLC(16, 2);                                                     \
+      } else {                                                                 \
+        LD_BYP_SLC(8, 1);                                                      \
+      }                                                                        \
+    }                                                                          \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ VecType                                  \
+    robust_load(const RobustTypePtr ptr, OffsetType idx) {                     \
+      return __musa::robust_load<VecType, BaseType>(ptr, idx);                 \
+    }                                                                          \
+                                                                               \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ void                                     \
+    store(BaseType *ptr, OffsetType idx, const VecType &dst) {                 \
+      *(VecType *)(ptr + idx) = dst;                                           \
+    }                                                                          \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ void                                     \
+    robust_store(RobustTypePtr ptr, OffsetType idx, const VecType &dst) {      \
+      __musa::robust_store<VecType, BaseType>(dst, ptr, idx);                  \
+    }                                                                          \
+                                                                               \
+    __device__ VecType() {                                                     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = 0;                                                \
+      }                                                                        \
+    }                                                                          \
+    __device__ VecType(const VecType &t) {                                     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = t.val_.elem[i];                                   \
+      }                                                                        \
+    }                                                                          \
+    __device__ VecType &operator=(const VecType &t) {                          \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = t.val_.elem[i];                                   \
+      }                                                                        \
+      return *this;                                                            \
+    }                                                                          \
+    __device__ VecType(_CTYPE val) {                                           \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = val;                                              \
+      }                                                                        \
+    }                                                                          \
+    template <typename SrcVecType>                                             \
+    friend __device__ VecType operator+(VecType lhs, const SrcVecType &rhs) {  \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        lhs.val_.elem[i] += static_cast<BaseType>(rhs.val_.elem[i]);           \
+      }                                                                        \
+      return lhs;                                                              \
+    }                                                                          \
+    friend __device__ VecType operator+(VecType lhs, const _CTYPE &rhs) {      \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        lhs.val_.elem[i] += rhs;                                               \
+      }                                                                        \
+      return lhs;                                                              \
+    }                                                                          \
+    friend __device__ VecType operator-(VecType lhs, const VecType &rhs) {     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        lhs.val_.elem[i] -= rhs.val_.elem[i];                                  \
+      }                                                                        \
+      return lhs;                                                              \
+    }                                                                          \
+    friend __device__ VecType operator*(VecType lhs, const VecType &rhs) {     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        lhs.val_.elem[i] *= rhs.val_.elem[i];                                  \
+      }                                                                        \
+      return lhs;                                                              \
+    }                                                                          \
+    template <typename Func> __device__ VecType &apply() {                     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = Func::apply(this->val_.elem[i]);                  \
+      }                                                                        \
+      return *this;                                                            \
+    }                                                                          \
+    template <typename SrcVecType>                                             \
+    static __device__ VecType cvt(const SrcVecType &src) {                     \
+      VecType dst;                                                             \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        dst.val_.elem[i] = (BaseType)(src.val_.elem[i]);                       \
+      }                                                                        \
+      return dst;                                                              \
+    }                                                                          \
+    union U {                                                                  \
+      __device__ U() {                                                         \
+        MACRO_UNROLL                                                           \
+        for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {           \
+          this->elem[i] = 0;                                                   \
+        }                                                                      \
+      }                                                                        \
+      Ttype storage;                                                           \
+      BaseType elem[sizeof(Ttype) / sizeof(BaseType)];                         \
+      VxTy vt_elem;                                                            \
+    };                                                                         \
+    U val_{};                                                                  \
+  }
+DEF_VECT(float16_t, float16_t);
+DEF_VECT(float16_t, Half2);
+DEF_VECT(float16_t, Half4);
+DEF_VECT(float16_t, Half8);
+DEF_VECT(float16_t, Half16);
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+DEF_VECT(bfloat16_t, bfloat16_t);
+DEF_VECT(bfloat16_t, Bhalf2);
+DEF_VECT(bfloat16_t, Bhalf4);
+DEF_VECT(bfloat16_t, Bhalf8);
+DEF_VECT(bfloat16_t, Bhalf16);
+DEF_VECT(bfloat16_t, Bhalf32);
+#endif
+// #if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+// DEF_VECT(fp8_e4m3_t, fp8_e4m3_t);
+// DEF_VECT(fp8_e4m3_t, E4m3x2);
+// DEF_VECT(fp8_e4m3_t, E4m3x4);
+// DEF_VECT(fp8_e4m3_t, E4m3x8);
+// DEF_VECT(fp8_e4m3_t, E4m3x16);
+// DEF_VECT(fp8_e4m3_t, E4m3x32);
+// DEF_VECT(fp8_e5m2_t, fp8_e5m2_t);
+// DEF_VECT(fp8_e5m2_t, E5m2x2);
+// DEF_VECT(fp8_e5m2_t, E5m2x4);
+// DEF_VECT(fp8_e5m2_t, E5m2x8);
+// DEF_VECT(fp8_e5m2_t, E5m2x16);
+// DEF_VECT(fp8_e5m2_t, E5m2x32);
+// #endif
+DEF_VECT(bool, char);
+DEF_VECT(bool, char2);
+DEF_VECT(bool, char3);
+DEF_VECT(bool, char4);
+DEF_VECT(bool, Char8);
+DEF_VECT(bool, Char16);
+DEF_VECT(int8_t, int8_t);
+DEF_VECT(int8_t, char2);
+DEF_VECT(int8_t, char3);
+DEF_VECT(int8_t, char4);
+DEF_VECT(int8_t, Char8);
+DEF_VECT(int8_t, Char16);
+DEF_VECT(qint8_t, int8_t);
+DEF_VECT(qint8_t, char2);
+DEF_VECT(qint8_t, char3);
+DEF_VECT(qint8_t, char4);
+DEF_VECT(qint8_t, Char8);
+DEF_VECT(qint8_t, Char16);
+DEF_VECT(uint8_t, uint8_t);
+DEF_VECT(uint8_t, uchar2);
+DEF_VECT(uint8_t, uchar3);
+DEF_VECT(uint8_t, uchar4);
+DEF_VECT(uint8_t, uint4);
+DEF_VECT(uint8_t, Uint8);
+DEF_VECT(int16_t, int16_t);
+DEF_VECT(int16_t, short2);
+DEF_VECT(int16_t, short3);
+DEF_VECT(int16_t, short4);
+DEF_VECT(int16_t, Short8);
+DEF_VECT(int16_t, Short16);
+DEF_VECT(uint16_t, ushort);
+DEF_VECT(uint16_t, ushort2);
+DEF_VECT(uint16_t, ushort3);
+DEF_VECT(uint16_t, ushort4);
+DEF_VECT(uint16_t, Ushort8);
+DEF_VECT(int32_t, int);
+DEF_VECT(int32_t, int2);
+DEF_VECT(int32_t, int3);
+DEF_VECT(int32_t, int4);
+DEF_VECT(int32_t, Int8);
+DEF_VECT(int32_t, Int16);
+DEF_VECT(uint32_t, uint);
+DEF_VECT(uint32_t, uint2);
+DEF_VECT(uint32_t, uint3);
+DEF_VECT(uint32_t, uint4);
+DEF_VECT(uint32_t, UInt8);
+DEF_VECT(uint32_t, UInt16);
+DEF_VECT(uint64_t, uint64_t);
+DEF_VECT(uint64_t, ulong2);
+DEF_VECT(uint64_t, ulong3);
+DEF_VECT(uint64_t, ulong4);
+DEF_VECT(uint64_t, ULong8);
+DEF_VECT(uint64_t, ULong16);
+DEF_VECT(int64_t, int64_t);
+DEF_VECT(int64_t, long2);
+DEF_VECT(int64_t, long3);
+DEF_VECT(int64_t, long4);
+DEF_VECT(int64_t, Long8);
+DEF_VECT(int64_t, Long16);
+DEF_VECT(float, float);
+DEF_VECT(float, float2);
+DEF_VECT(float, float3);
+DEF_VECT(float, float4);
+DEF_VECT(float, Float8);
+DEF_VECT(float, Float16);
+
+DEF_VECT(double, double);
+DEF_VECT(double, double2);
+DEF_VECT(double, double3);
+DEF_VECT(double, double4);
+
+#undef DEF_VECT
+#undef MACRO_UNROLL
+template <typename T> struct ComputeDType {
+  using Type = typename std::conditional<
+      (sizeof(T) >= 4), T,
+      typename std::conditional<
+          std::is_integral<T>::value,
+          typename std::conditional<std::is_unsigned<T>::value, uint32_t,
+                                    int32_t>::type,
+          float>::type>::type;
+};
+template <> struct ComputeDType<bool> {
+  using Type = bool;
+};
+// qint type calculate as float
+template <> struct ComputeDType<qint8_t> {
+  using Type = float;
+};
+
+// static inline bool check_qint8_only_scale(CTR x) {
+//   if (x.type == TensorImpl::Type::QINT8) {
+//     size_t zp_size = x.quant_desc.zero_point.size();
+//     return zp_size == 0 || (zp_size == 1 && x.quant_desc.zero_point[0] == 0);
+//   }
+//   return true;
+// }
+
+template <typename T, typename... Types>
+bool check_qint8_only_scale(T x0, Types... xn) {
+  bool ok = check_qint8_only_scale(x0);
+  return ok && check_qint8_only_scale(xn...);
+}
+#define CHECK_QINT8(ARGS...)                                                   \
+  {                                                                            \
+    if (!check_qint8_only_scale(ARGS)) {                                       \
+      return fail::NOT_SUPPORTED() << "qint8 only support zero_point==0";      \
+    }                                                                          \
+  }
+
+
+} // namespace dnn
+} // namespace musa
diff --git a/vllm_musa/csrc/linear/gemv_swiglu.mu b/vllm_musa/csrc/linear/gemv_swiglu.mu
new file mode 100644
index 000000000..069b4071a
--- /dev/null
+++ b/vllm_musa/csrc/linear/gemv_swiglu.mu
@@ -0,0 +1,415 @@
+#include "common.muh"
+#include "dtype.muh"
+#include "gemv_swiglu.muh"
+#include <musa_runtime.h>
+#include <cassert>
+#include <mutex>
+
+
+#include "../type_convert.muh"
+#include "../dispatch_utils.h"
+
+
+// load vector into shared memory every blockDim.x elems
+// each warp load [BLOCK_K, BLOCK_M] of mat using DMA.LD.B128 into register
+// If BLOCK_K = 16 and BLOCK_M = 32, for example,
+// tid0-tid7  load 1st row,
+// tid8-tid15 load 2nd row, and so on...
+// col_continuous: matrix(k*m) is col-continuous
+template <typename T, AlignAttr ALIGN_ATTR, int BLOCK_K, int BLOCK_M, int iobit>
+__global__ void batch_gemv_swiglu_col_continuous_kernel(
+    T *d, const T *mat, const T *v, const int batch, const int m, const int k,
+    const int ldm, const int64_t batch_stride_m, const int inc_x,
+    const int64_t batch_stride_x, const int inc_d, const int64_t batch_stride_d,
+    float alpha, const int half_m) {
+  constexpr int bits_of_byte = 8;
+  constexpr int Vlen = iobit / (sizeof(T) * bits_of_byte);
+  using CDtype = typename InferComputeDtype<T>::ComputeDtype;
+  using vec = typename VecType<T, iobit>::Ttype;
+  using cvec = typename VecType<CDtype, 128>::Ttype;
+
+  int tid2 = threadIdx.x;
+  int bx = blockIdx.x; // block in vector
+  int by = blockIdx.y; // batch
+  if (bx * BLOCK_M >= half_m)
+    return;
+
+  // load mat using B128
+  int part_id = tid2 / (blockDim.x / 2);
+  int tid = tid2 % (blockDim.x / 2);
+  int slice_id = tid % (BLOCK_M / Vlen);
+  int lane_id = tid / (BLOCK_M / Vlen);
+  int m_id = bx * BLOCK_M + slice_id * Vlen;
+
+  CDtype res = CDtype(0);
+  extern __shared__ char shared_array[];
+  CDtype *smem = (CDtype *)shared_array;
+
+  CDtype cur_thread_sum[Vlen];
+#pragma unroll
+  for (int i = 0; i < Vlen; i++) {
+    cur_thread_sum[i] = (CDtype)0;
+  }
+
+  // #if defined(__MUSA_ARCH__) && __MUSA_ARCH__ >= 220
+  for (int row_id = lane_id; row_id < k; row_id += BLOCK_K) {
+    if (m_id < half_m) {
+      T m_reg[Vlen];
+      // base_mat_offset 去掉batch维度，删除by * batch_stride_m +
+      int64_t base_mat_offset = row_id * (int64_t)ldm +
+                                (int64_t)m_id + part_id * half_m;
+      if constexpr (ALIGN_ATTR == AlignAttr::ALIGN_VEC) {
+        *((vec *)m_reg) = *(vec *)(mat + base_mat_offset);
+      } else {
+        if (bx * BLOCK_M + BLOCK_M <= half_m) {
+          *((vec *)m_reg) = *(vec *)(mat + base_mat_offset);
+        } else {
+          m_reg[0] = *(mat + base_mat_offset);
+#pragma unroll
+          for (int p = 1; p < Vlen; p++) {
+            m_reg[p] = m_id + p < half_m ? *(mat + base_mat_offset + p) : T(0);
+          }
+        }
+      }
+
+      int64_t v_offset = row_id * (int64_t)inc_x + by * batch_stride_x;
+      T v_val = v[v_offset];
+#pragma unroll
+      for (int p = 0; p < Vlen; p++) {
+        cur_thread_sum[p] += (CDtype)v_val * (CDtype)m_reg[p];
+      }
+    }
+  }
+  // #endif
+
+  CDtype *smem_mat = &smem[0];
+  int sts_offset = lane_id * BLOCK_M * 2 + part_id * BLOCK_M + slice_id * Vlen;
+#pragma unroll
+  for (int p = 0; p < Vlen / 4; p++) {
+    *(cvec *)(&smem_mat[sts_offset + p * 4]) =
+        *(cvec *)(&cur_thread_sum[p * 4]);
+  }
+  __SYNCTHREADS_LM;
+  // return;
+
+#pragma unroll
+  for (int p = 0; p < Vlen / 4; p++) {
+#pragma unroll
+    for (int a = BLOCK_K / 2; a > 1; a >>= 1) {
+      if (lane_id < a) {
+        cvec w_reg = *(cvec *)(&smem_mat[sts_offset + p * 4]);
+        cvec r_reg = *(cvec *)(&smem_mat[sts_offset + p * 4 + a * BLOCK_M * 2]);
+        w_reg.x += r_reg.x;
+        w_reg.y += r_reg.y;
+        w_reg.z += r_reg.z;
+        w_reg.w += r_reg.w;
+        *(cvec *)(&smem_mat[sts_offset + p * 4]) = w_reg;
+      }
+      __SYNCTHREADS_LM;
+    }
+  }
+
+  m_id = bx * BLOCK_M + tid;
+  if (tid < BLOCK_M && m_id < half_m) {
+    CDtype v0 = alpha * (smem_mat[tid] + smem_mat[BLOCK_M * 2 + tid]);
+    CDtype v1 = alpha * (smem_mat[tid + BLOCK_M] +
+                         smem_mat[BLOCK_M * 2 + tid + BLOCK_M]);
+    res = v0 * sigmoid(v0) * v1;
+    int64_t out_offset = by * batch_stride_d + m_id * (int64_t)inc_d;
+    d[out_offset] = res;
+  }
+}
+
+template<int VLEN>
+class Block {
+public:
+  const int threads;
+  const int bm;
+  const int bk;
+
+  static constexpr float base_score = 1.f;
+  static constexpr float epilog_loss_ratio = 0.5;
+
+  Block() = delete;
+  Block(int _threads, int _bm, int _bk)
+      : threads(_threads), bm(_bm), bk(_bk){
+        assert(bk * (bm / VLEN) * 2 <= 1024);
+      }
+  std::string name() const {
+    return std::to_string(threads) + "_" + std::to_string(bm) + "_" + std::to_string(bk);
+  }
+  float estimate(const int batch, const int m, const int k,
+                 const int nr_mp) const {
+
+    int m_tile = ceil_div(m, bm);
+    // int k_tile = ceil_div(k, bk);
+    int wave = ceil_div(batch * m_tile, nr_mp);
+    float wave_ratio = (float)batch * m  / ((float)bm * wave * nr_mp);
+    // float epilog_loss = (float)bk / max(k, bk) * epilog_loss_ratio;
+    // float epilog_ratio = 1.f - epilog_loss;
+    float epilog_ratio = 1.f;
+    float normal_score = base_score * wave_ratio * epilog_ratio;
+    return normal_score;
+  }
+};
+
+static std::vector<Block<4>> block_list_vlen4;
+static std::once_flag block_list_init_flag;
+
+void InitBlockList() {
+  std::call_once(block_list_init_flag, []() {
+    block_list_vlen4.clear();
+    block_list_vlen4.emplace_back(1024, 16, 128);
+    block_list_vlen4.emplace_back(1024, 32, 64);
+    block_list_vlen4.emplace_back(1024, 64, 32);
+    block_list_vlen4.emplace_back(1024, 128, 16);
+  });
+}
+
+template <int VLEN>
+void get_best_block(const std::vector<Block<VLEN>>& block_list,
+                                  const int batch, const int m,
+                                  const int k, const int nr_mp ) {
+  float best_score = 0.f;
+  const Block<VLEN>* best_block = nullptr;
+  for (auto& block : block_list) {
+    auto score = block.estimate(batch, m, k, nr_mp);
+    if (score> best_score) {
+      best_score = score;
+      best_block = &block;
+    }
+  }
+  printf("best block %s score %f\n", best_block->name().c_str(), best_score);
+  // return {best_block, best_score};
+}
+
+#define BATCH_GEMV_SWIGLU_COLCONTINUOUS_LAUNCH(_ALIGN_ATTR, _IOBIT)            \
+  {                                                                            \
+    int smem_bytes = 2 * block_k * block_m * sizeof(float);                    \
+    batch_gemv_swiglu_col_continuous_kernel<T, _ALIGN_ATTR, block_k, block_m,  \
+                                            _IOBIT>                            \
+        <<<grid_size, block_size, smem_bytes, stream>>>(                       \
+            o, mat, v, batch, mat_m, mat_k, ldm, batch_stride_m, inc_v,        \
+            batch_stride_v, inc_o, batch_stride_o, alpha, half_mat_m);         \
+  }
+
+#define BATCH_GEMV_SWIGLU_COLCONTINUOUS_CALL(_IOBIT)                           \
+  if (aligned_vec) {                                                           \
+    BATCH_GEMV_SWIGLU_COLCONTINUOUS_LAUNCH(AlignAttr::ALIGN_VEC, _IOBIT)       \
+  } else {                                                                     \
+    BATCH_GEMV_SWIGLU_COLCONTINUOUS_LAUNCH(AlignAttr::GENERAL, _IOBIT)         \
+  }
+
+/*
+    only accept the shape
+     - matrix (k, m), the column is continuous
+     - vector (k)
+     - output of gemv   (m)
+     - output of swiglu (m/2)
+*/
+template <typename T>
+StatusX GemvSwiGLU(T *o, const T *mat, const T *v, const int batch,
+                   const int mat_m, const int mat_k, const int ldm,
+                   const int batch_stride_m, const int inc_v,
+                   const int batch_stride_v, const int inc_o,
+                   const int batch_stride_o, double alpha, musaStream_t &stream,
+                   int num_mp) {
+  // printf("%s\n", __FUNCTION__);
+
+  if (0 != mat_m % 2) {
+    printf("mat_m must be even.\n");
+    return StatusX::INVALID_PARAMETER;
+  }
+  const int half_mat_m = mat_m / 2;
+
+  dim3 block_size{128, 1, 1};
+  dim3 grid_size{1, static_cast<uint32_t>(batch), 1};
+
+  constexpr int iobit = 128;
+  constexpr int bits_of_byte = 8;
+  constexpr int vlen = iobit / (sizeof(T) * bits_of_byte);
+  constexpr int elems_per_cacheline = 128 / sizeof(T);
+
+  //   int device_id = 0;
+  //   int num_mp = 0;
+  //   musaGetDevice(&device_id);
+  //   musaDeviceProp device_prop;
+  //   if (musaSuccess == musaGetDeviceProperties(&device_prop, device_id)) {
+  //     num_mp = device_prop.multiProcessorCount;
+  //   }
+
+  bool aligned_vec = half_mat_m >= vlen && half_mat_m % vlen == 0;
+  int estimated_grid_x = ceil_div(half_mat_m, elems_per_cacheline);
+
+  //   constexpr int block_m = elems_per_cacheline;
+  //   constexpr int block_k = 32;
+  //   block_size.x = block_k * (block_m / vlen) * 2;
+  //   grid_size.x = ceil_div(half_mat_m, block_m);
+  //   BATCH_GEMV_SWIGLU_COLCONTINUOUS_CALL(iobit)
+
+  if (false) {
+    printf("num_mp %d\n", num_mp);
+    printf("estimated_grid_x %d\n", estimated_grid_x);
+  }
+
+  //   printf("elems_per_cacheline %d, block_m %d block_k %d vlen %d, "
+  //          "block_size.x %d grid_size.x %d\n",
+  //          elems_per_cacheline, block_m, block_k, vlen, block_size.x,
+  //          grid_size.x);
+
+#define BLOCK_STRATEGY(_BLK_M, _BLK_K)                                         \
+  constexpr int block_m = _BLK_M;                                              \
+  constexpr int block_k = _BLK_K;                                              \
+  block_size.x = block_k * (block_m / vlen) * 2;                               \
+  grid_size.x = ceil_div(half_mat_m, block_m);                                 \
+  BATCH_GEMV_SWIGLU_COLCONTINUOUS_CALL(iobit)
+
+  //   BLOCK_STRATEGY(elems_per_cacheline * 16, 4); //not good
+  //   BLOCK_STRATEGY(elems_per_cacheline * 8, 8);
+  // BLOCK_STRATEGY(32, 32);
+  // BLOCK_STRATEGY(32, 64);
+  // BLOCK_STRATEGY(64, 32);
+  // BLOCK_STRATEGY(128, 16);
+  // BLOCK_STRATEGY(16, 128);
+
+
+#define BLOCK_STRATEGY_M(_BLK_M)                                         \
+  constexpr int block_m = _BLK_M;                                        \
+  constexpr int block_k = 512/(block_m/vlen);                            \
+  block_size.x = block_k * (block_m / vlen) * 2;                         \
+  grid_size.x = ceil_div(half_mat_m, block_m);                           \
+  BATCH_GEMV_SWIGLU_COLCONTINUOUS_CALL(iobit)
+
+  // if (estimated_grid_x >= num_mp * 4) {
+  //   // printf("estimated_grid_x %d elems_per_cacheline %d\n", estimated_grid_x,elems_per_cacheline);
+  //   BLOCK_STRATEGY_M(elems_per_cacheline);
+  // } else if (estimated_grid_x >= num_mp * 2) {
+  //   BLOCK_STRATEGY_M(elems_per_cacheline);
+  // } else {
+  //   BLOCK_STRATEGY_M(elems_per_cacheline/2);
+  // }
+  // if (estimated_grid_x >= num_mp * 2) {
+  //   BLOCK_STRATEGY_M(elems_per_cacheline * 2);
+  // } else {
+  //   BLOCK_STRATEGY_M(elems_per_cacheline);
+  // }
+  // InitBlockList();
+  // get_best_block<vlen>(block_list_vlen4, batch, half_mat_m, mat_k, num_mp);
+
+  // printf("xxx\n");
+
+  // BLOCK_STRATEGY(elems_per_cacheline / 2, 128);
+  // BLOCK_STRATEGY(elems_per_cacheline, 64);
+
+  if (estimated_grid_x >= num_mp * 4) {
+    BLOCK_STRATEGY(elems_per_cacheline, 64);
+  } else if (estimated_grid_x >= num_mp * 2 && mat_k <= mat_m) {
+    BLOCK_STRATEGY(elems_per_cacheline / 2, 64);
+  } else {
+    BLOCK_STRATEGY(elems_per_cacheline / 2, 128);
+  }
+
+
+
+
+  MUSA_CHECK(musaGetLastError());
+  return StatusX::SUCCESS;
+}
+
+#define INST(T)                                                                \
+  template StatusX GemvSwiGLU(T *o, const T *mat, const T *v, const int batch, \
+                              const int mat_m, const int mat_k, const int ldm, \
+                              const int batch_stride_m, const int inc_v,       \
+                              const int batch_stride_v, const int inc_o,       \
+                              const int batch_stride_o, double alpha,          \
+                              musaStream_t &stream, int num_mp);
+
+INST(float)
+#undef INST
+
+void gemm_swiglu(torch::Tensor &output, // [batch, hidden_size/2]
+                 torch::Tensor &input,  // [batch, hidden_size]
+                 torch::Tensor &weight, // [hidden_size, hidden_size/2]
+                 double alpha) {
+
+  TORCH_CHECK(weight.scalar_type() == input.scalar_type(), "weight dtype must same as input");
+  // construct kernel input m=1
+  int batch = input.size(0);  // batch
+  int k = input.size(1);  // k
+  int n = weight.size(1); // n
+  int output_dim = n / 2;
+
+  int ldm = n;
+  float alpha_f = static_cast<float>(alpha);
+  int batch_stride_weight = k * n;
+  int inc_input = 1;
+  int batch_stride_input = k;
+  int inc_output = 1;
+  int batch_stride_output = output_dim;
+
+  const at::musa::OptionalMUSAGuard device_guard(device_of(input));
+  musaStream_t stream = at::musa::getCurrentMUSAStream();
+  // num_mp
+  int device;
+  musaGetDevice(&device);
+  musaDeviceProp device_prop;
+  musaGetDeviceProperties(&device_prop, device);
+  int num_mp = device_prop.multiProcessorCount;
+
+  switch (input.scalar_type()) {
+
+  case at::ScalarType::Float:
+    GemvSwiGLU<float>(output.data_ptr<float>(),
+              weight.data_ptr<float>(),
+              input.data_ptr<float>(),
+              batch,
+              n,
+              k,
+              ldm,
+              batch_stride_weight,
+              inc_input,
+              batch_stride_input,
+              inc_output,
+              batch_stride_output,
+              alpha_f,
+              stream,
+              num_mp);
+    break;
+  case at::ScalarType::Half:
+    GemvSwiGLU<float16_t>(static_cast<float16_t*>(output.data_ptr()),
+              static_cast<float16_t*>(weight.data_ptr()),
+              static_cast<float16_t*>(input.data_ptr()),
+              batch,
+              n,
+              k,
+              ldm,
+              batch_stride_weight,
+              inc_input,
+              batch_stride_input,
+              inc_output,
+              batch_stride_output,
+              alpha_f,
+              stream,
+              num_mp);
+    break;
+  case at::ScalarType::BFloat16:
+    GemvSwiGLU<bfloat16_t>(static_cast<bfloat16_t*>(output.data_ptr()),
+              static_cast<bfloat16_t*>(weight.data_ptr()),
+              static_cast<bfloat16_t*>(input.data_ptr()),
+              batch,
+              n,
+              k,
+              ldm,
+              batch_stride_weight,
+              inc_input,
+              batch_stride_input,
+              inc_output,
+              batch_stride_output,
+              alpha_f,
+              stream,
+              num_mp);
+    break;
+  default:
+    break;
+  }
+}
diff --git a/vllm_musa/csrc/linear/gemv_swiglu.muh b/vllm_musa/csrc/linear/gemv_swiglu.muh
new file mode 100644
index 000000000..7e8c78a09
--- /dev/null
+++ b/vllm_musa/csrc/linear/gemv_swiglu.muh
@@ -0,0 +1,34 @@
+#pragma once
+#include "common.muh"
+
+template <typename T>
+StatusX DispatchBatchGemvGeneral(T *o, const T *mat, const T *v,
+                                 const int batch, const int mat_m, const int k,
+                                 double alpha, musaStream_t &stream);
+
+template <typename T>
+StatusX _SwiGluRun(T *out, const T *in, const int batch, const int split_dim,
+                   const int nr_elem, musaStream_t &stream);
+
+/*
+C = MATMUL(A, B):
+    A shape: m, k
+    B shape: k, n
+    C shape: m, n
+X = SWIGLU(C)
+    C shape: m, n
+    X shape: m*n/2
+
+when m==1,
+    A is vector
+    B is matrix
+    X = GEMV_SIWGLU(B, A, n, k)
+    X shape: n/2
+*/
+template <typename T>
+StatusX GemvSwiGLU(T *o, const T *mat, const T *v, const int batch,
+                   const int mat_m, const int mat_k, const int ldm,
+                   const int batch_stride_m, const int inc_v,
+                   const int batch_stride_v, const int inc_o,
+                   const int batch_stride_o, double alpha, musaStream_t &stream,
+                   int num_mp);
\ No newline at end of file
diff --git a/vllm_musa/csrc/linear/gemv_swiglu_w4a16.mu b/vllm_musa/csrc/linear/gemv_swiglu_w4a16.mu
new file mode 100644
index 000000000..d79da34a5
--- /dev/null
+++ b/vllm_musa/csrc/linear/gemv_swiglu_w4a16.mu
@@ -0,0 +1,369 @@
+#include "common.muh"
+#include "dtype.muh"
+#include "gemv_swiglu.muh"
+#include <musa_runtime.h>
+#include <cassert>
+#include <mutex>
+
+#include "../type_convert.muh"
+#include "../dispatch_utils.h"
+
+using namespace musa::dnn;
+
+/* NOT IMPLEMENTED
+  T float16_t
+  mat weight (m, k)
+    int4b_t (store in int32_t) 4bit * 8, packed along the col direction
+    i.e.
+        weight(m, k//8)  -> intermediate(m, k//8, 8)  ->  packed(m, k//8)
+  vector input (k)
+    fp16
+
+  output (m//2)
+    fp16
+
+  scale per channel (m)
+    float32
+*/
+
+
+
+/* CURRENT IMPLEMENTATION
+  T float16_t
+  mat weight (k, m)
+    int4b_t (store in int32_t) 4bit * 8, packed along the row direction
+    i.e.
+        weight(k, m)  -> intermediate(k//8, 8, m)  ->  packed(k//8, m)
+  vector input (k)
+    fp16
+
+  output (m//2)
+    fp16
+
+  scale per channel (m)
+    float32
+*/
+template <typename T, AlignAttr ALIGN_ATTR, int BLOCK_K, int BLOCK_M, int iobit, typename PACKED_T, int coeff = 8>
+__global__ void batch_gemv_swiglu_col_continuous_w4a16(
+    T *d, const PACKED_T *mat, const T *v, const int batch, const int m, const int k,
+    const int ldm, const int64_t batch_stride_m, const int inc_x,
+    const int64_t batch_stride_x, const int inc_d, const int64_t batch_stride_d,
+    float alpha, float* scale, const int half_m) {
+  constexpr int bits_of_byte = 8;
+  constexpr int Vlen = iobit / (sizeof(PACKED_T) * bits_of_byte);
+  // PACKED_T is int32_t, coeff is 8
+  using CDtype = typename InferComputeDtype<T>::ComputeDtype;
+  using pvec = typename VecType<PACKED_T, iobit>::Ttype;
+  using tvec = typename VecType<T, coeff * sizeof(T) * bits_of_byte>::Ttype;
+  using fvec = typename VecType<float, Vlen * sizeof(float) * bits_of_byte>::Ttype;
+  using cvec = typename VecType<CDtype, 128>::Ttype;
+
+  int tid2 = threadIdx.x;
+  int bx = blockIdx.x;
+  int by = blockIdx.y;
+  if (bx * BLOCK_M >= half_m)
+    return;
+
+  // load mat using B128
+  int part_id = tid2 / (blockDim.x / 2);
+  int tid = tid2 % (blockDim.x / 2);
+  int slice_id = tid % (BLOCK_M / Vlen);
+  int lane_id = tid / (BLOCK_M / Vlen);
+  int m_id = bx * BLOCK_M + slice_id * Vlen;
+
+  extern __shared__ char shared_array[];
+  CDtype *smem = (CDtype *)shared_array;
+
+  CDtype cur_thread_sum[Vlen];
+#pragma unroll
+  for (int i = 0; i < Vlen; i++) {
+    cur_thread_sum[i] = (CDtype)0;
+  }
+
+  float s_reg[Vlen];
+  if (m_id < half_m){
+    int64_t base_m_id =((int64_t)m_id + part_id * half_m);
+    if constexpr (ALIGN_ATTR == AlignAttr::ALIGN_VEC) {
+      *(fvec *)(s_reg) = *(fvec *)(scale + base_m_id);
+    } else {
+      if (bx * BLOCK_M + BLOCK_M <= half_m) {
+        *(fvec *)(s_reg) = *(fvec *)(scale + base_m_id);
+      } else {
+        s_reg[0] = scale[base_m_id];
+  #pragma unroll
+        for (int p = 1; p < 4; p++) {
+          s_reg[p]  = (m_id + p < half_m) ? scale[base_m_id+p] : 0.f;
+        }
+      }
+    }
+  }
+
+  uint8_t tmp = 0;
+
+  for (int row_id = lane_id; row_id < k; row_id += BLOCK_K) {
+    if (m_id < half_m) {
+      PACKED_T m_reg[Vlen];
+      // 去掉mat的batch idx，mat只有input_channel, output_channel两个维度
+      int64_t base_mat_offset = row_id * (int64_t)ldm +
+                                (int64_t)m_id + part_id * half_m;
+      if constexpr (ALIGN_ATTR == AlignAttr::ALIGN_VEC) {
+        *((pvec *)m_reg) = *(pvec *)(mat + base_mat_offset);
+      } else {
+        if (bx * BLOCK_M + BLOCK_M <= half_m) {
+          *((pvec *)m_reg) = *(pvec *)(mat + base_mat_offset);
+        } else {
+          m_reg[0] = *(mat + base_mat_offset);
+#pragma unroll
+          for (int p = 1; p < Vlen; p++) {
+            m_reg[p] = (m_id + p < half_m) ? *(mat + base_mat_offset + p) : static_cast<PACKED_T>(0);
+          }
+        }
+      }
+
+      int64_t v_offset = (int64_t)row_id * coeff * inc_x + (int64_t)by * batch_stride_x;
+      T v_reg[coeff];
+       *((tvec *)v_reg) = *(tvec *)(v + v_offset);
+#pragma unroll
+      for (int p = 0; p < Vlen; p++) {
+        PACKED_T packed_val = m_reg[p];
+        CDtype val = 0;
+#pragma unroll
+        for (int i = 0; i < coeff; i++) {
+          // int32_t elem = int4b_t((packed_val >> ((coeff-1-i)*4)) & 0xF);
+          int32_t elem = ((packed_val >> (i * 4)) & 0xF);
+          elem = (elem & 0x8) ? elem - 16 : elem;
+          val += (CDtype)v_reg[i] * ((CDtype)elem * s_reg[p]);
+        }
+        cur_thread_sum[p] += val;
+      }
+    }
+  }
+
+  CDtype *smem_mat = &smem[0];
+  int sts_offset = (lane_id * BLOCK_M * 2 + part_id * BLOCK_M + slice_id * Vlen);
+#pragma unroll
+  for (int p = 0; p < Vlen/ 4; p++) {
+    *(cvec *)(&smem_mat[sts_offset + p * 4]) =
+        *(cvec *)(&cur_thread_sum[p * 4]);
+  }
+  __SYNCTHREADS_LM;
+
+#pragma unroll
+  for (int p = 0; p < Vlen / 4; p++) {
+#pragma unroll
+    for (int a = BLOCK_K / 2; a > 1; a >>= 1) {
+      if (lane_id < a) {
+        cvec w_reg = *(cvec *)(&smem_mat[sts_offset + p * 4]);
+        cvec r_reg = *(cvec *)(&smem_mat[sts_offset + p * 4 + a * BLOCK_M * 2]);
+        w_reg.x += r_reg.x;
+        w_reg.y += r_reg.y;
+        w_reg.z += r_reg.z;
+        w_reg.w += r_reg.w;
+        *(cvec *)(&smem_mat[sts_offset + p * 4]) = w_reg;
+      }
+      __SYNCTHREADS_LM;
+    }
+  }
+
+  m_id = bx * BLOCK_M + tid;
+  if (tid < BLOCK_M  && m_id < half_m) {
+    CDtype v0 = alpha * (smem_mat[tid] + smem_mat[BLOCK_M * 2 + tid]);
+    CDtype v1 = alpha * (smem_mat[tid + BLOCK_M] +
+                         smem_mat[BLOCK_M * 2  + tid + BLOCK_M]);
+     CDtype res = v0 * sigmoid(v0) * v1;
+    int64_t out_offset = by * batch_stride_d + m_id * (int64_t)inc_d;
+    d[out_offset] = res;
+  }
+}
+
+#define BATCH_GEMV_SWIGLU_COLCONTINUOUS_LAUNCH(_ALIGN_ATTR, _IOBIT)            \
+  {                                                                            \
+    int smem_bytes = 2 * block_k * block_m * sizeof(float);                    \
+    batch_gemv_swiglu_col_continuous_w4a16<T, _ALIGN_ATTR, block_k, block_m,   \
+                                            _IOBIT, int32_t, 8>                \
+        <<<grid_size, block_size, smem_bytes, stream>>>(                       \
+            o, mat, v, batch, mat_m, packed_mat_k, ldm, packed_batch_stride_m, \
+            inc_v, batch_stride_v, inc_o, batch_stride_o, alpha, scale,        \
+            half_mat_m);                                                       \
+  }
+
+#define BATCH_GEMV_SWIGLU_COLCONTINUOUS_CALL(_IOBIT)                           \
+  if (aligned_vec) {                                                           \
+    BATCH_GEMV_SWIGLU_COLCONTINUOUS_LAUNCH(AlignAttr::ALIGN_VEC, _IOBIT)       \
+  } else {                                                                     \
+    BATCH_GEMV_SWIGLU_COLCONTINUOUS_LAUNCH(AlignAttr::GENERAL, _IOBIT)         \
+  }
+
+/*
+    only accept the shape
+     - matrix (k, m), the column is continuous
+     - vector (k)
+     - output of gemv   (m)
+     - output of swiglu (m/2)
+*/
+template <typename T>
+StatusX GemvSwiGluW4A16(T *o, const int32_t *mat, const T *v, float* scale, const int batch,
+                   const int mat_m, const int mat_k, const int ldm,
+                   const int batch_stride_m, const int inc_v,
+                   const int batch_stride_v, const int inc_o,
+                   const int batch_stride_o, float alpha, musaStream_t &stream,
+                   int num_mp) {
+  // printf("%s\n", __FUNCTION__);
+
+  if (0 != mat_m % 2) {
+    printf("mat_m must be even.\n");
+    return StatusX::INVALID_PARAMETER;
+  }
+  if (0 != mat_k % 8) {
+    printf("mat_k must be must be a multiple of 8.\n");
+    return StatusX::INVALID_PARAMETER;
+  }
+  const int half_mat_m = mat_m / 2;
+  const int packed_mat_k = mat_k / 8; // 8*int4_t
+  const int packed_batch_stride_m = batch_stride_m / 8;
+  // printf("packed_mat_k %d\n",packed_mat_k);
+
+  dim3 block_size{128, 1, 1};
+  dim3 grid_size{1, static_cast<uint32_t>(batch), 1};
+
+  constexpr int iobit = 128;
+  constexpr int bits_of_byte = 8;
+  constexpr int vlen = iobit / (sizeof(int32_t) * bits_of_byte);
+  constexpr int elems_per_cacheline = 128 / sizeof(int32_t);
+
+  bool aligned_vec = half_mat_m >= vlen && half_mat_m % vlen == 0;
+  int estimated_grid_x = ceil_div(half_mat_m, elems_per_cacheline);
+
+#define BLOCK_STRATEGY_M(_BLK_M)                                         \
+{                                                                        \
+  /*printf("block_m %d\n", _BLK_M);*/                                        \
+  constexpr int block_m = _BLK_M;                                        \
+  constexpr int block_k = 512/(block_m/vlen);                            \
+  block_size.x = block_k * (block_m / vlen) * 2;                         \
+  grid_size.x = ceil_div(half_mat_m, block_m);                           \
+  BATCH_GEMV_SWIGLU_COLCONTINUOUS_CALL(iobit)                            \
+}
+
+
+  // BLOCK_STRATEGY_M(elems_per_cacheline);
+  int blk_m_list[4] = {8, 16, 32, 64};
+  int blk_scores[4] = {0, 0, 0, 0};
+  for(int i =0; i<4; i++){
+    int blks = ceil_div(half_mat_m, blk_m_list[i]);
+    if(blks < num_mp){
+      blk_scores[i] = 1;
+      break;
+    }
+    float ratio = (float)blks/num_mp;
+    blk_scores[i] =  ratio - (int)ratio;
+  }
+
+  int chose_block = 0;
+  int chose_score = 0;
+  for(int i=0; i<4; i++){
+    if(blk_scores[i]>chose_score){
+      chose_score = blk_scores[i];
+      chose_block = i;
+    }
+  }
+
+  switch(chose_block){
+    case 0:
+      BLOCK_STRATEGY_M(8);
+      break;
+    case 1:
+      BLOCK_STRATEGY_M(16);
+      break;
+    case 2:
+      BLOCK_STRATEGY_M(32);
+      break;
+    case 3:
+      BLOCK_STRATEGY_M(64);
+      break;
+    default:
+      BLOCK_STRATEGY_M(32);
+  }
+
+
+  MUSA_CHECK(musaGetLastError());
+  return StatusX::SUCCESS;
+}
+
+
+void gemm_swiglu_w4a16(torch::Tensor &output,
+                       torch::Tensor &input,
+                       torch::Tensor &weight,
+                       torch::Tensor &scale,
+                       double alpha) {
+
+
+  TORCH_CHECK(weight.scalar_type() == at::ScalarType::Int, "weight dtype must be int32.");
+  // TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16, "input dtype must be bfloat16.");
+
+  int batch = input.size(0);  //batch
+  int k = input.size(1);  // k
+  int n = weight.size(1); // int32 packint4
+  // printf("batch %d k%d n%d\n", batch, k, n);
+  int output_dim = n / 2;
+
+  int ldm = n;
+  float alpha_f = static_cast<float>(alpha);
+  int batch_stride_weight = k * n;
+  int inc_input = 1;
+  int batch_stride_input = k;
+  int inc_output = 1;
+  int batch_stride_output = output_dim;
+
+  const at::musa::OptionalMUSAGuard device_guard(device_of(input));
+  musaStream_t stream = at::musa::getCurrentMUSAStream();
+
+  // num_mp
+  int device;
+  musaGetDevice(&device);
+  musaDeviceProp device_prop;
+  musaGetDeviceProperties(&device_prop, device);
+  int num_mp = device_prop.multiProcessorCount;
+
+  // TODO: add float16
+  if (input.scalar_type() == at::ScalarType::BFloat16) {
+    GemvSwiGluW4A16<bfloat16_t>(
+      static_cast<bfloat16_t *>(output.data_ptr()),
+      reinterpret_cast<int32_t *>(weight.data_ptr()),
+      static_cast<bfloat16_t *>(input.data_ptr()),
+      static_cast<float *>(scale.data_ptr()),
+      batch,
+      n,
+      k,
+      ldm,
+      batch_stride_weight,
+      inc_input,
+      batch_stride_input,
+      inc_output,
+      batch_stride_output,
+      alpha_f,
+      stream,
+      num_mp
+    );
+  } else if (input.scalar_type() == at::ScalarType::Half) {
+    GemvSwiGluW4A16<float16_t>(
+      static_cast<float16_t *>(output.data_ptr()),
+      reinterpret_cast<int32_t *>(weight.data_ptr()),
+      static_cast<float16_t *>(input.data_ptr()),
+      static_cast<float *>(scale.data_ptr()),
+      batch,
+      n,
+      k,
+      ldm,
+      batch_stride_weight,
+      inc_input,
+      batch_stride_input,
+      inc_output,
+      batch_stride_output,
+      alpha_f,
+      stream,
+      num_mp
+    );
+  }
+
+
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/linear/gemv_w4a16.mu b/vllm_musa/csrc/linear/gemv_w4a16.mu
new file mode 100644
index 000000000..d92547ea5
--- /dev/null
+++ b/vllm_musa/csrc/linear/gemv_w4a16.mu
@@ -0,0 +1,314 @@
+#include "common.muh"
+#include "dtype.muh"
+#include "gemv_swiglu.muh"
+#include <musa_runtime.h>
+#include <cassert>
+#include <mutex>
+
+#include "../type_convert.muh"
+#include "../dispatch_utils.h"
+
+using namespace musa::dnn;
+
+template <typename T, AlignAttr ALIGN_ATTR, int BLOCK_K, int BLOCK_M, int iobit, typename PACKED_T, int coeff = 8>
+__global__ void batch_gemv_col_continuous_w4a16(
+    T *d, const PACKED_T *mat, const T *v, const int batch, const int m, const int k,
+    const int ldm, const int64_t batch_stride_m, const int inc_x,
+    const int64_t batch_stride_x, const int inc_d, const int64_t batch_stride_d,
+    float alpha, float* scale, const int half_m) {
+  constexpr int bits_of_byte = 8;
+  constexpr int Vlen = iobit / (sizeof(PACKED_T) * bits_of_byte);
+  using CDtype = typename InferComputeDtype<T>::ComputeDtype;
+  using pvec = typename VecType<PACKED_T, iobit>::Ttype;
+  using tvec = typename VecType<T, coeff * sizeof(T) * bits_of_byte>::Ttype;
+  using fvec = typename VecType<float, Vlen * sizeof(float) * bits_of_byte>::Ttype;
+  using cvec = typename VecType<CDtype, 128>::Ttype;
+
+  int tid2 = threadIdx.x;
+  int bx = blockIdx.x;
+  int by = blockIdx.y;
+  if (bx * BLOCK_M >= half_m)
+    return;
+
+  int part_id = tid2 / (blockDim.x / 2);
+  int tid = tid2 % (blockDim.x / 2);
+  int slice_id = tid % (BLOCK_M / Vlen);
+  int lane_id = tid / (BLOCK_M / Vlen);
+  int m_id = bx * BLOCK_M + slice_id * Vlen;
+
+  extern __shared__ char shared_array[];
+  CDtype *smem = (CDtype *)shared_array;
+
+  CDtype cur_thread_sum[Vlen];
+#pragma unroll
+  for (int i = 0; i < Vlen; i++) {
+    cur_thread_sum[i] = (CDtype)0;
+  }
+
+  float s_reg[Vlen];
+  if (m_id < half_m) {
+    int64_t base_m_id = m_id + part_id * half_m;
+    if constexpr (ALIGN_ATTR == AlignAttr::ALIGN_VEC) {
+      *(fvec *)(s_reg) = *(fvec *)(scale + base_m_id);
+    } else {
+      if (bx * BLOCK_M + BLOCK_M <= half_m) {
+        *(fvec *)(s_reg) = *(fvec *)(scale + base_m_id);
+      } else {
+        s_reg[0] = scale[base_m_id];
+#pragma unroll
+        for (int p = 1; p < 4; p++) {
+          s_reg[p] = (m_id + p < half_m) ? scale[base_m_id + p] : 0.f;
+        }
+      }
+    }
+  }
+
+  for (int row_id = lane_id; row_id < k; row_id += BLOCK_K) {
+    if (m_id < half_m) {
+      PACKED_T m_reg[Vlen];
+      int64_t base_mat_offset = row_id * (int64_t)ldm + m_id + part_id * half_m;
+      if constexpr (ALIGN_ATTR == AlignAttr::ALIGN_VEC) {
+        *((pvec *)m_reg) = *(pvec *)(mat + base_mat_offset);
+      } else {
+        if (bx * BLOCK_M + BLOCK_M <= half_m) {
+          *((pvec *)m_reg) = *(pvec *)(mat + base_mat_offset);
+        } else {
+          m_reg[0] = *(mat + base_mat_offset);
+#pragma unroll
+          for (int p = 1; p < Vlen; p++) {
+            m_reg[p] = (m_id + p < half_m) ? *(mat + base_mat_offset + p) : static_cast<PACKED_T>(0);
+          }
+        }
+      }
+
+      int64_t v_offset = row_id * coeff * inc_x + by * batch_stride_x;
+      T v_reg[coeff];
+      *((tvec *)v_reg) = *(tvec *)(v + v_offset);
+#pragma unroll
+      for (int p = 0; p < Vlen; p++) {
+        PACKED_T packed_val = m_reg[p];
+        CDtype val = 0;
+#pragma unroll
+        for (int i = 0; i < coeff; i++) {
+          int32_t elem = ((packed_val >> (i * 4)) & 0xF);
+          elem = (elem & 0x8) ? elem - 16 : elem;
+          val += (CDtype)v_reg[i] * ((CDtype)elem * s_reg[p]);
+        }
+        cur_thread_sum[p] += val;
+      }
+    }
+  }
+
+
+  CDtype *smem_mat = smem;
+  int sts_offset = (lane_id * BLOCK_M * 2 + part_id * BLOCK_M + slice_id * Vlen);
+#pragma unroll
+  for (int p = 0; p < Vlen / 4; p++) {
+    *(cvec *)(&smem_mat[sts_offset + p * 4]) =
+        *(cvec *)(&cur_thread_sum[p * 4]);
+  }
+  __SYNCTHREADS_LM;
+
+#pragma unroll
+  for (int p = 0; p < Vlen / 4; p++) {
+#pragma unroll
+    for (int a = BLOCK_K / 2; a > 1; a >>= 1) {
+      if (lane_id < a) {
+        cvec w_reg = *(cvec *)(&smem_mat[sts_offset + p * 4]);
+        cvec r_reg = *(cvec *)(&smem_mat[sts_offset + p * 4 + a * BLOCK_M * 2]);
+        w_reg.x += r_reg.x;
+        w_reg.y += r_reg.y;
+        w_reg.z += r_reg.z;
+        w_reg.w += r_reg.w;
+        *(cvec *)(&smem_mat[sts_offset + p * 4]) = w_reg;
+      }
+      __SYNCTHREADS_LM;
+    }
+  }
+
+  m_id = bx * BLOCK_M + tid;
+  if (tid < BLOCK_M && m_id < half_m) {
+    CDtype res0 = alpha * (smem_mat[tid] + smem_mat[BLOCK_M * 2 + tid]);
+    int64_t out_offset0 = by * batch_stride_d + m_id * (int64_t)inc_d;
+    d[out_offset0] = res0;
+  }
+  if (tid < BLOCK_M && m_id < half_m) {
+    CDtype res1 = alpha * (smem_mat[tid + BLOCK_M] + smem_mat[BLOCK_M * 2 + tid + BLOCK_M]);
+    int64_t out_offset1 = by * batch_stride_d + (m_id + half_m) * (int64_t)inc_d;
+    d[out_offset1] = res1;
+  }
+}
+
+
+#define BATCH_GEMV_COLCONTINUOUS_LAUNCH(_ALIGN_ATTR, _IOBIT)            \
+  {                                                                            \
+    int smem_bytes = 2 * block_k * block_m * sizeof(float);                    \
+    batch_gemv_col_continuous_w4a16<T, _ALIGN_ATTR, block_k, block_m,      \
+                                           _IOBIT, int32_t, 8>                  \
+        <<<grid_size, block_size, smem_bytes, stream>>>(                       \
+            o, mat, v, batch, mat_m, packed_mat_k, ldm, packed_batch_stride_m, \
+            inc_v, batch_stride_v, inc_o, batch_stride_o, alpha, scale,        \
+            half_mat_m);                                                       \
+  }
+
+#define BATCH_GEMV_COLCONTINUOUS_CALL(_IOBIT)                           \
+  if (aligned_vec) {                                                           \
+    BATCH_GEMV_COLCONTINUOUS_LAUNCH(AlignAttr::ALIGN_VEC, _IOBIT)         \
+  } else {                                                                     \
+    BATCH_GEMV_COLCONTINUOUS_LAUNCH(AlignAttr::GENERAL, _IOBIT)           \
+  }
+
+
+template <typename T>
+StatusX GemvW4A16(T *o, const int32_t *mat, const T *v, float* scale, const int batch,
+                        const int mat_m, const int mat_k, const int ldm,
+                        const int batch_stride_m, const int inc_v,
+                        const int batch_stride_v, const int inc_o,
+                        const int batch_stride_o, float alpha, musaStream_t &stream,
+                        int num_mp) {
+  if (0 != mat_m % 2) {
+    printf("mat_m must be even.\n");
+    return StatusX::INVALID_PARAMETER;
+  }
+  if (0 != mat_k % 8) {
+    printf("mat_k must be a multiple of 8.\n");
+    return StatusX::INVALID_PARAMETER;
+  }
+
+  const int half_mat_m = mat_m / 2;
+  const int packed_mat_k = mat_k / 8;
+  const int packed_batch_stride_m = batch_stride_m / 8;
+
+  dim3 block_size{128, 1, 1};
+  dim3 grid_size{1, static_cast<uint32_t>(batch), 1};
+
+  constexpr int iobit = 128;
+  constexpr int bits_of_byte = 8;
+  constexpr int vlen = iobit / (sizeof(int32_t) * bits_of_byte);
+  constexpr int elems_per_cacheline = 128 / sizeof(int32_t);
+
+  bool aligned_vec = half_mat_m >= vlen && half_mat_m % vlen == 0;
+
+
+  int blk_m_list[4] = {8, 16, 32, 64};
+  int blk_scores[4] = {0, 0, 0, 0};
+  for (int i = 0; i < 4; i++) {
+    int blks = ceil_div(half_mat_m, blk_m_list[i]);
+    if (blks < num_mp) {
+      blk_scores[i] = 1;
+      break;
+    }
+    float ratio = (float)blks / num_mp;
+    blk_scores[i] = ratio - (int)ratio;
+  }
+
+  int chose_block = 0;
+  int chose_score = 0;
+  for (int i = 0; i < 4; i++) {
+    if (blk_scores[i] > chose_score) {
+      chose_score = blk_scores[i];
+      chose_block = i;
+    }
+  }
+
+
+#define BLOCK_STRATEGY_M(_BLK_M)                                         \
+  {                                                                      \
+    constexpr int block_m = _BLK_M;                                      \
+    constexpr int block_k = 512 / (block_m / vlen);                      \
+    block_size.x = block_k * (block_m / vlen) * 2;                         \
+    grid_size.x = ceil_div(half_mat_m, block_m);                           \
+    BATCH_GEMV_COLCONTINUOUS_CALL(iobit)                            \
+  }
+
+  switch (chose_block) {
+    case 0:
+      BLOCK_STRATEGY_M(8);
+      break;
+    case 1:
+      BLOCK_STRATEGY_M(16);
+      break;
+    case 2:
+      BLOCK_STRATEGY_M(32);
+      break;
+    case 3:
+      BLOCK_STRATEGY_M(64);
+      break;
+    default:
+      BLOCK_STRATEGY_M(32);
+  }
+
+  MUSA_CHECK(musaGetLastError());
+  return StatusX::SUCCESS;
+}
+
+void gemm_w4a16(torch::Tensor &output,
+                        torch::Tensor &input,
+                        torch::Tensor &weight,
+                        torch::Tensor &scale,
+                        double alpha) {
+  TORCH_CHECK(weight.scalar_type() == at::ScalarType::Int, "weight dtype must be int32.");
+  // TORCH_CHECK(input.scalar_type() == at::ScalarType::BFloat16, "input dtype must be bfloat16.");
+
+  int batch = input.size(0);
+  int k = input.size(1);
+  int n = weight.size(1);
+  int output_dim = n;
+
+  int ldm = n;
+  float alpha_f = static_cast<float>(alpha);
+  int batch_stride_weight = k * n;
+  int inc_input = 1;
+  int batch_stride_input = k;
+  int inc_output = 1;
+  int batch_stride_output = output_dim;
+
+  const at::musa::OptionalMUSAGuard device_guard(device_of(input));
+  musaStream_t stream = at::musa::getCurrentMUSAStream();
+
+  int device;
+  musaGetDevice(&device);
+  musaDeviceProp device_prop;
+  musaGetDeviceProperties(&device_prop, device);
+  int num_mp = device_prop.multiProcessorCount;
+
+  if (input.scalar_type() == at::ScalarType::BFloat16) {
+    GemvW4A16<bfloat16_t>(
+      static_cast<bfloat16_t *>(output.data_ptr()),
+      reinterpret_cast<int32_t *>(weight.data_ptr()),
+      static_cast<bfloat16_t *>(input.data_ptr()),
+      static_cast<float *>(scale.data_ptr()),
+      batch,
+      n,
+      k,
+      ldm,
+      batch_stride_weight,
+      inc_input,
+      batch_stride_input,
+      inc_output,
+      batch_stride_output,
+      alpha_f,
+      stream,
+      num_mp
+    );
+  } else if (input.scalar_type() == at::ScalarType::Half) {
+    GemvW4A16<float16_t>(
+      static_cast<float16_t *>(output.data_ptr()),
+      reinterpret_cast<int32_t *>(weight.data_ptr()),
+      static_cast<float16_t *>(input.data_ptr()),
+      static_cast<float *>(scale.data_ptr()),
+      batch,
+      n,
+      k,
+      ldm,
+      batch_stride_weight,
+      inc_input,
+      batch_stride_input,
+      inc_output,
+      batch_stride_output,
+      alpha_f,
+      stream,
+      num_mp
+    );
+  }
+}
diff --git a/vllm_musa/csrc/linear/integer_subbyte.h b/vllm_musa/csrc/linear/integer_subbyte.h
new file mode 100644
index 000000000..b9fa54aba
--- /dev/null
+++ b/vllm_musa/csrc/linear/integer_subbyte.h
@@ -0,0 +1,193 @@
+/* Copyright @2020-2024 Moore Threads Technology Co., Ltd("Moore Threads"). All
+ * rights reserved.
+ *
+ * This software ("this software and its documentations" or "the software") is
+ * protected by Copyright and the information contained herein is confidential.
+ *
+ * The software contained herein is PROPRIETARY to Moore Threads and is being
+ * provided under the terms and conditions of a form of Moore Threads software
+ * license agreement by and between Moore Threads and Licensee ("License
+ * Agreement") or electronically accepted by Licensee. Notwithstanding any
+ * terms or conditions to the contrary in the License Agreement, copy or
+ * disclosure of the software to any third party without the express written
+ * consent of Moore Threads is prohibited.
+ *
+ * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE LICENSE
+ * AGREEMENT, MOORE THREADS MAKES NO REPRESENTATION ABOUT ANY WARRANTIES,
+ * INCLUDING BUT NOT LIMITED TO THE SUITABILITY OF THE SOFTWARE FOR ANY
+ * PURPOSE. IT IS PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF
+ * ANY KIND. MOORE THREADS DISCLAIMS ALL WARRANTIES WITH REGARD TO THE
+ * SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY,
+ * NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE.
+ * NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE
+ * LICENSE AGREEMENT, IN NO EVENT SHALL MOORE THREADS BE LIABLE FOR ANY
+ * SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY
+ * DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS,
+ * WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS
+ * ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE
+ * OF THE SOFTWARE.
+ */
+
+#ifndef INCLUDE_MUDNN_API_INTEGER_SUBBYTE_H_
+#define INCLUDE_MUDNN_API_INTEGER_SUBBYTE_H_
+
+#include <limits>
+#include <type_traits>
+
+namespace musa::dnn {
+
+// cutlass integer_subbyte class
+template <int Bits, bool Signed = true> struct integer_subbyte {
+  using Storage = uint8_t;
+
+  static_assert(Bits <= 8 * sizeof(Storage),
+                "Require a subbyte of bits in integer_subbyte");
+
+  using xint_t = typename std::conditional<Signed, int, unsigned>::type;
+
+  static constexpr Storage bits_mask_ = Storage((1 << Bits) - 1);
+
+  static constexpr Storage sign_mask_ = Storage((Signed ? 1 : 0) << (Bits - 1));
+
+  Storage storage;
+
+  __host__ __device__ constexpr integer_subbyte() {}
+
+  __host__ __device__ constexpr integer_subbyte(int value)
+      : storage(reinterpret_cast<Storage const &>(value) & bits_mask_) {}
+
+  __host__ __device__ constexpr integer_subbyte(unsigned value)
+      : storage(reinterpret_cast<Storage const &>(value) & bits_mask_) {}
+
+  __host__ __device__ constexpr operator xint_t() const {
+    if (sign_mask_ & storage) {
+      return xint_t(storage) | ~xint_t(bits_mask_);
+    } else {
+      return xint_t(storage);
+    }
+  }
+
+  __host__ __device__ constexpr bool
+  operator==(integer_subbyte const &rhs) const {
+    return storage == rhs.storage;
+  }
+
+  __host__ __device__ constexpr bool
+  operator!=(integer_subbyte const &rhs) const {
+    return storage != rhs.storage;
+  }
+
+  __host__ __device__ constexpr bool
+  operator<=(integer_subbyte const &rhs) const {
+    return storage <= rhs.storage;
+  }
+
+  __host__ __device__ constexpr bool
+  operator<(integer_subbyte const &rhs) const {
+    return storage < rhs.storage;
+  }
+
+  __host__ __device__ constexpr bool
+  operator>=(integer_subbyte const &rhs) const {
+    return storage >= rhs.storage;
+  }
+
+  __host__ __device__ constexpr bool
+  operator>(integer_subbyte const &rhs) const {
+    return storage > rhs.storage;
+  }
+};
+
+using uint1b_t = integer_subbyte<1, false>;
+using int2b_t = integer_subbyte<2, true>;
+using uint2b_t = integer_subbyte<2, false>;
+using int4b_t = integer_subbyte<4, true>;
+using uint4b_t = integer_subbyte<4, false>;
+using bin1_t = bool;
+
+} // namespace musa::dnn
+
+namespace std {
+
+template <> struct numeric_limits<musa::dnn::uint1b_t> {
+  __host__ __device__ static constexpr musa::dnn::uint1b_t const
+  lowest() noexcept {
+    return 0;
+  }
+  __host__ __device__ static constexpr musa::dnn::uint1b_t const
+  min() noexcept {
+    return 0;
+  }
+  __host__ __device__ static constexpr musa::dnn::uint1b_t const
+  max() noexcept {
+    return 1;
+  }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+template <> struct numeric_limits<musa::dnn::int2b_t> {
+  __host__ __device__ static constexpr musa::dnn::int2b_t lowest() noexcept {
+    return -2;
+  }
+  __host__ __device__ static constexpr musa::dnn::int2b_t min() noexcept {
+    return -2;
+  }
+  __host__ __device__ static constexpr musa::dnn::int2b_t max() noexcept {
+    return 1;
+  }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = true;
+};
+
+template <> struct numeric_limits<musa::dnn::uint2b_t> {
+  __host__ __device__ static constexpr musa::dnn::uint2b_t const
+  lowest() noexcept {
+    return 0;
+  }
+  __host__ __device__ static constexpr musa::dnn::uint2b_t const
+  min() noexcept {
+    return 0;
+  }
+  __host__ __device__ static constexpr musa::dnn::uint2b_t const
+  max() noexcept {
+    return 3;
+  }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+template <> struct numeric_limits<musa::dnn::int4b_t> {
+  __host__ __device__ static constexpr musa::dnn::int4b_t lowest() noexcept {
+    return -8;
+  }
+  __host__ __device__ static constexpr musa::dnn::int4b_t min() noexcept {
+    return -8;
+  }
+  __host__ __device__ static constexpr musa::dnn::int4b_t max() noexcept {
+    return 7;
+  }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = true;
+};
+
+template <> struct numeric_limits<musa::dnn::uint4b_t> {
+  __host__ __device__ static constexpr musa::dnn::uint4b_t const
+  lowest() noexcept {
+    return 0;
+  }
+  __host__ __device__ static constexpr musa::dnn::uint4b_t const
+  min() noexcept {
+    return 0;
+  }
+  __host__ __device__ static constexpr musa::dnn::uint4b_t const
+  max() noexcept {
+    return 15;
+  }
+  static constexpr bool is_integer = true;
+  static constexpr bool is_signed = false;
+};
+
+} // namespace std
+
+#endif // INCLUDE_MUDNN_API_INTEGER_SUBBYTE_H_
diff --git a/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm.mu b/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm.mu
new file mode 100644
index 000000000..f123b565a
--- /dev/null
+++ b/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm.mu
@@ -0,0 +1,160 @@
+#include "mute/tensor.hpp"
+#include <torch/torch.h>
+#include "mixed_input_gemm_kernel.hpp"
+#include "mutlass/util/device_memory.h"
+#include "mutlass/util/GPU_Clock.hpp"
+
+#include "../../type_convert.muh"
+#include "../../dispatch_utils.h"
+
+#include <vector>
+#include <cstdlib>
+
+using namespace mute;
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// GEMM kernel configurations
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+//using TileShape          = Shape<_256, _256, _32>;
+using TileShape          = Shape<_64, _64, _32>;
+
+constexpr int AlignmentA = 8;
+
+using ElementB           = mutlass::int4b_t;
+constexpr int AlignmentB = 8;
+
+
+template <typename ElementAC, typename ElementScale, int GroupSizeK_>
+void mixed_input_gemm_kernel_launcher(
+  ElementAC *A, ElementB *B, ElementAC *C, ElementScale *B_scale,
+  int M, int N, int K) {
+
+  using KernelTraits = MixedInputGemmKernelTratis<TileShape, ElementAC, AlignmentA, ElementB, AlignmentB, ElementScale, ElementAC, GroupSizeK_>;
+  auto problem_shape = make_shape(M, N, K);
+
+  typename KernelTraits::Arguments arguments {
+    A,
+    B,
+    C,
+    B_scale,
+    K,
+    K,
+    N,
+    K / KernelTraits::GroupSizeK,
+    1
+  };
+
+  typename KernelTraits::Params params = KernelTraits::to_underlying_arguments(problem_shape, arguments);
+
+  bool can_impl_check = KernelTraits::can_implement(problem_shape, arguments);
+  if (!can_impl_check) {
+    printf("Mixed Input GEMM Error: TileShape dimensions do not match the expected sizes.");
+  }
+  assert(can_impl_check && "Error: TileShape dimensions do not match the expected sizes.");
+
+  auto tile_shape = TileShape{};
+  dim3 grid_dims = {static_cast<uint32_t>(mute::ceil_div(M, get<0>(tile_shape))), static_cast<uint32_t>(mute::ceil_div(N, get<1>(tile_shape))), 1};
+
+  const musaStream_t stream = at::musa::getCurrentMUSAStream();
+  mixed_input_gemm<KernelTraits><<<grid_dims, KernelTraits::NumThreads, 0, stream>>>(params);
+}
+
+
+
+#define CALL_MIXED_INPUT_GEMM_LAUNCH(ElementAC, ElementScale, GroupSizeK_)                 \
+mixed_input_gemm_kernel_launcher<ElementAC, ElementScale, GroupSizeK_>(                    \
+    static_cast<ElementAC*>(A.data_ptr()),                                       \
+    static_cast<ElementB*>(B.data_ptr()),                                       \
+    static_cast<ElementAC*>(C.data_ptr()),                                       \
+    static_cast<ElementScale*>(B_scale.data_ptr()),                            \
+    M, N, K);                                                                   \
+
+
+
+void group_scale_gemm(
+    torch::Tensor &A,       // [bseqlen, in]
+    torch::Tensor &B,       // [out, in/8]
+    torch::Tensor &C,       // [bseqlen, out]
+    torch::Tensor &B_scale)
+{
+    TORCH_CHECK(A.dim() == 2, "A must be dim 2.")
+    TORCH_CHECK(B.dim() == 2, "B must be dim 2.")
+
+    int32_t hidden_size = A.size(1);
+    int32_t scale_k_len = B_scale.size(1);
+    int group_size = ceil_div(hidden_size, scale_k_len);
+
+    int M = A.size(0);
+    int N = B.size(0);
+    int K = hidden_size;
+
+    TORCH_CHECK(group_size == 128 || group_size == 64,
+    "group_size only support 128 or 64 for per-group w4a16");
+
+    if (A.scalar_type() == at::ScalarType::Half)
+    {
+      if (B_scale.scalar_type() == at::ScalarType::Float)
+      {
+        if (group_size == 64) {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::half_t, float, 64)
+        } else {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::half_t, float, 128)
+        }
+      }
+      else if (B_scale.scalar_type() == at::ScalarType::Half)
+      {
+        if (group_size == 64) {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::half_t, mutlass::half_t, 64)
+        } else {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::half_t, mutlass::half_t, 128)
+        }
+      }
+      else if (B_scale.scalar_type() == at::ScalarType::BFloat16)
+      {
+        if (group_size == 64) {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::half_t, mutlass::bfloat16_t, 64)
+        } else {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::half_t, mutlass::bfloat16_t, 128)
+        }
+      }
+      else
+      {
+          printf("no support on mutlass w4a16 gemm\n");
+          assert(false);
+      }
+    }
+    else if (A.scalar_type() == at::ScalarType::BFloat16)
+    {
+      if (B_scale.scalar_type() == at::ScalarType::Float)
+      {
+        if (group_size == 64) {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::bfloat16_t, float, 64)
+        } else {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::bfloat16_t, float, 128)
+        }
+      }
+      else if (B_scale.scalar_type() == at::ScalarType::Half)
+      {
+        if (group_size == 64) {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::bfloat16_t, mutlass::half_t, 64)
+        } else {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::bfloat16_t, mutlass::half_t, 128)
+        }
+      }
+      else if (B_scale.scalar_type() == at::ScalarType::BFloat16)
+      {
+        if (group_size == 64) {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::bfloat16_t, mutlass::bfloat16_t, 64)
+        } else {
+          CALL_MIXED_INPUT_GEMM_LAUNCH(mutlass::bfloat16_t, mutlass::bfloat16_t, 128)
+        }
+      }
+      else
+      {
+          printf("no support on mutlass w4a16 gemm\n");
+          assert(false);
+      }
+    }
+}
diff --git a/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm_kernel.hpp b/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm_kernel.hpp
new file mode 100644
index 000000000..40507ab88
--- /dev/null
+++ b/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm_kernel.hpp
@@ -0,0 +1,466 @@
+#pragma once
+
+#include "mixed_input_gemm_kernel_traits.hpp"
+
+template <class LayoutAccum>
+MUTE_DEVICE
+auto
+convert_layout_acc_rowcol(LayoutAccum acc_layout) {
+  if constexpr (decltype(rank<0>(acc_layout))::value == 3) { // MP31
+    // from ((2, VN, VM), MMA_M, MMA_N) to ((VM, MMA_M),(2, VN, MMA_N));
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(rank(acc_layout))::value == 3);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0,2>(l), get<1>(l)), make_layout(get<0,0>(l), get<0,1>(l), get<2>(l)));
+  } else {  // MP22
+    // from((2,4), MMA_M, MMA_N) to ((4, MMA_M), (2, MMA_N))
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(size<0, 1>(acc_layout))::value == 4);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0,1>(l), get<1>(l)), make_layout(get<0,0>(l), get<2>(l)));
+  }
+}
+
+
+template <
+	class Engine0, class Layout0,
+	class Engine1, class Layout1,
+  class Engine2, class Layout2
+>
+MUTE_DEVICE
+void
+dequant_tensor(Tensor<Engine0, Layout0>& src,
+							 Tensor<Engine1, Layout1>& dst,
+               Tensor<Engine2, Layout2>& scale) {
+	MUTE_STATIC_ASSERT_V(size(src) == size(dst));
+  using MmaType = typename Engine1::value_type;
+  using QuantType = typename Engine0::value_type;
+
+  static_assert(sizeof_bits_v<MmaType> == 16);
+  constexpr bool is_half_mma = is_same_v<MmaType, mutlass::half_t>;
+  static_assert(sizeof_bits_v<QuantType> == 4);
+
+  using v2f16_t = mute::conditional_t<is_half_mma,
+                                      _Float16 __attribute__((vector_size(4))),
+                                      mutlass::Array<mutlass::bfloat16_t, 2>>;
+
+  using v2f32_t = float __attribute__((vector_size(8)));
+  using v4f32_t = float __attribute__((vector_size(16)));
+
+  using v4u8_t = uint8_t __attribute__((vector_size(4)));
+
+#define U8_CAST 0
+
+#if U8_CAST
+  Tensor int_view = recast<uint8_t>(src);
+#else
+  Tensor int_view = recast<uint32_t>(src);
+#endif
+  Tensor v2f16_view = recast<v2f16_t>(dst);
+
+  constexpr float rescale = 1.0f / 16.0f;
+  constexpr float zp = 128.0f;
+  constexpr uint32_t mask = 0xF0F0F0F0;
+  constexpr int shift = 4;
+
+  MUTE_UNROLL
+  for (int ni = 0; ni < size<1>(int_view); ++ni) {
+    float scale_val = scale(ni) * rescale;
+    MUTE_UNROLL
+    for (int ki = 0; ki < size<2>(int_view); ++ki) {
+      MUTE_UNROLL
+      for (int vi = 0; vi < size<0>(int_view); ++vi) {
+#if U8_CAST
+        uint8_t lo = (int_view(vi, ni, ki) << 4) & 0xF0;
+        uint8_t hi = int_view(vi, ni, ki) & 0xF0;
+
+        v2f32_t value {(float(lo) - zp) * scale_val, (float(hi) - zp) * scale_val};
+
+        v2f16_view(vi, ni, ki) = __builtin_convertvector(value, v2f16_t);
+#else
+        const float scaled_zp = -zp * scale_val;
+
+        uint32_t value = int_view(vi, ni, ki);
+        uint32_t hi;
+        uint32_t lo;
+        v4f32_t f32x4_hi;
+        v4f32_t f32x4_lo;
+
+#if 1
+        asm("BIT.LUT.AND %0, %1, %2" : "=R"(hi) : "R"(value), "R"(mask));
+        asm("BIT.LSL.AND %0, %1, %2, %3" : "=R"(lo) : "R"(value), "R"(shift), "R"(mask));
+        asm("MOV.U8888_TO_F32.BST4 %0, %1" : "=R"(f32x4_hi) : "R"(hi));
+        asm("MOV.U8888_TO_F32.BST4 %0, %1" : "=R"(f32x4_lo) : "R"(lo));
+#else
+        hi = value & mask;
+        lo = (value << shift) & mask;
+
+        v4u8_t vec_hi = *reinterpret_cast<v4u8_t*>(&hi);
+        v4u8_t vec_lo = *reinterpret_cast<v4u8_t*>(&lo);
+
+        f32x4_hi = __builtin_convertvector(vec_hi, v4f32_t);
+        f32x4_lo = __builtin_convertvector(vec_lo, v4f32_t);
+#endif
+
+        MUTE_UNROLL
+        for (int i = 0; i < 4; ++i) {
+          v2f32_t value {f32x4_lo[i] * scale_val + scaled_zp, f32x4_hi[i] * scale_val + scaled_zp};
+          if constexpr (is_half_mma) {
+            v2f16_view(i, ni, ki) = __builtin_convertvector(value, v2f16_t);
+          } else {
+            float2 f32x2 = make_float2(value[0], value[1]);
+            __mt_bfloat162 res = __float22bfloat162_rn(f32x2);
+
+            v2f16_t f16x2;
+            f16x2[0] = mutlass::bfloat16_t(res.x);
+            f16x2[1] = mutlass::bfloat16_t(res.y);
+
+            v2f16_view(i, ni, ki) = f16x2;
+          }
+        }
+
+        //float f16x2;
+        //asm("FOP.FMA %0.l, %1, %2, %3" : "=R"(f16x2) : "R"(f32x4_lo[0]), "R"(scale_val), "R"(scaled_zp));
+        //asm("FOP.FMA %0.h, %1, %2, %3" : "=R"(f16x2) : "R"(f32x4_hi[0]), "R"(scale_val), "R"(scaled_zp));
+        //v2f16_view(0, ni, ki) = *reinterpret_cast<v2f16_t*>(&f16x2);
+
+        //asm("FOP.FMA %0.l, %1, %2, %3" : "=R"(f16x2) : "R"(f32x4_lo[1]), "R"(scale_val), "R"(scaled_zp));
+        //asm("FOP.FMA %0.h, %1, %2, %3" : "=R"(f16x2) : "R"(f32x4_hi[1]), "R"(scale_val), "R"(scaled_zp));
+        //v2f16_view(1, ni, ki) = *reinterpret_cast<v2f16_t*>(&f16x2);
+
+        //asm("FOP.FMA %0.l, %1, %2, %3" : "=R"(f16x2) : "R"(f32x4_lo[2]), "R"(scale_val), "R"(scaled_zp));
+        //asm("FOP.FMA %0.h, %1, %2, %3" : "=R"(f16x2) : "R"(f32x4_hi[2]), "R"(scale_val), "R"(scaled_zp));
+        //v2f16_view(2, ni, ki) = *reinterpret_cast<v2f16_t*>(&f16x2);
+
+        //asm("FOP.FMA %0.l, %1, %2, %3" : "=R"(f16x2) : "R"(f32x4_lo[3]), "R"(scale_val), "R"(scaled_zp));
+        //asm("FOP.FMA %0.h, %1, %2, %3" : "=R"(f16x2) : "R"(f32x4_hi[3]), "R"(scale_val), "R"(scaled_zp));
+        //v2f16_view(3, ni, ki) = *reinterpret_cast<v2f16_t*>(&f16x2);
+#endif
+      }
+    }
+  }
+}
+
+
+template <class KernelTraits>
+__global__ __launch_bounds__(KernelTraits::NumThreads, 1)
+void mixed_input_gemm(const typename KernelTraits::Params params) {
+
+  auto problem_shape_mnk = params.problem_shape;
+  auto M = get<0>(problem_shape_mnk);
+  auto N = get<1>(problem_shape_mnk);
+  auto K = get<2>(problem_shape_mnk);
+
+  int thread_idx = int(threadIdx.x);
+  auto blk_shape = typename KernelTraits::TileShape{};
+
+  auto [m_coord, n_coord, l_coord] = static_cast<uint3>(blockIdx);
+  auto blk_coord_mnk = make_coord(m_coord, n_coord, _);
+
+  Tensor mA_mk = make_tensor(make_gmem_ptr(params.ptr_A), make_shape(M, K), params.stride_A);
+  Tensor mB_nk = make_tensor(make_gmem_ptr<typename KernelTraits::ElementB>(params.ptr_B), make_shape(N, K), params.stride_B);
+
+
+  Tensor gA = local_tile(mA_mk, blk_shape, blk_coord_mnk, Step<_1, X, _1>{});
+  Tensor gB = local_tile(mB_nk, blk_shape, blk_coord_mnk, Step<X, _1, _1>{});
+
+  typename KernelTraits::TiledMma tiled_mma;
+  Tensor accum = partition_fragment_C(tiled_mma, take<0,2>(blk_shape)); // (MMA,MMA_M,MMA_N)
+  clear(accum);
+
+  auto k_tile_iter  = mute::make_coord_iterator(shape<2>(gA));
+  int  k_tile_count = size<2>(gA);
+
+  // Mainloop
+  __shared__ char __attribute__((aligned(128))) smem[sizeof(typename KernelTraits::SharedStorage)];
+  auto& shared_storage = *reinterpret_cast<typename KernelTraits::SharedStorage*>(smem);
+
+  Tensor sA = make_tensor(make_smem_ptr(shared_storage.smem_a.data()), typename KernelTraits::SmemLayoutA{});
+  Tensor sB = make_tensor(make_smem_ptr(shared_storage.smem_b.data()), typename KernelTraits::SmemLayoutB{});
+
+  typename KernelTraits::GmemTiledCopyA gmem_tiled_copy_a;
+  typename KernelTraits::GmemTiledCopyB gmem_tiled_copy_b;
+
+  auto copy_a_thr = gmem_tiled_copy_a.get_slice(thread_idx);
+  auto copy_b_thr = gmem_tiled_copy_b.get_slice(thread_idx);
+
+  Tensor tAgA = copy_a_thr.partition_S(gA);                                  // (ACPY,ACPY_M,ACPY_K,k)
+  Tensor tAsA = copy_a_thr.partition_D(sA);                                  // (ACPY,ACPY_M,ACPY_K)
+  Tensor tBgB = copy_b_thr.partition_S(gB);                                  // (BCPY,BCPY_N,BCPY_K,k)
+  Tensor tBsB = copy_b_thr.partition_D(sB);                                  // (BCPY,BCPY_N,BCPY_K)
+
+  // Allocate the register tiles for double buffering -- same shape as partitioned data
+  Tensor tArA = make_fragment_like(tAsA);                                    // (ACPY,ACPY_M,ACPY_K)
+  Tensor tBrB = make_fragment_like(tBsB);                                    // (BCPY,BCPY_N,BCPY_K)
+  Tensor tBrB_quant = make_fragment_like<typename KernelTraits::ElementB>(tBsB);
+
+  auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
+  Tensor tCrA  = thr_mma.partition_fragment_A(sA);                           // (MMA,MMA_M,MMA_K)
+  Tensor tCrB  = thr_mma.partition_fragment_B(sB);                           // (MMA,MMA_M,MMA_K)
+
+  MUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(accum));                     // MMA_M
+  MUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(accum));                     // MMA_N
+  MUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));                      // MMA_K
+
+  //
+  // Copy Atom retiling
+  //
+
+  auto thr_copy_A       = make_tiled_copy_A(typename KernelTraits::SmemCopyAtomA{}, tiled_mma).get_thread_slice(thread_idx);
+  Tensor tCsA           = thr_copy_A.partition_S(sA);
+  Tensor tCrA_copy_view = thr_copy_A.retile_D(tCrA);
+  MUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));            // M
+
+  auto thr_copy_B       = make_tiled_copy_B(typename KernelTraits::SmemCopyAtomB{}, tiled_mma).get_thread_slice(thread_idx);
+  Tensor tCsB           = thr_copy_B.partition_S(sB);
+  Tensor tCrB_copy_view = thr_copy_B.retile_D(tCrB);
+  MUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));            // N
+
+
+  clear(tArA);
+  if constexpr (KernelTraits::IsMixedInput) {
+    clear(tBrB_quant);
+  } else {
+    clear(tBrB);
+  }
+
+  // Scale [N, K/128]
+  Tensor tBrS = make_tensor<typename KernelTraits::ElementScale>(make_shape(size<1>(tBsB)));
+  //if (thread0()) {
+  //  print("tBsB:"); print(tBsB); print("\n");
+  //  print("tBrB_quant:"); print(tBrB_quant); print("\n");
+  //  print("tBrS:"); print(tBrS); print("\n");
+  //}
+  int offset_sn[size(tBrS)];
+  int scale_k_cnt = 0;
+  int read_scale_flag = KernelTraits::ScaleInterleave;
+
+  typename KernelTraits::ElementScale const* scale_base = params.ptr_scale + n_coord * get<1>(blk_shape) * params.stride_sn;
+
+  Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB)));
+  Tensor tBcB = copy_b_thr.partition_S(cB);
+  Tensor tBpB = make_tensor<bool>(make_shape(size<1>(tBsB), size<2>(tBsB)), Stride<_1, _0>{});
+
+  Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA)));
+  Tensor tAcA = copy_a_thr.partition_S(cA);
+  Tensor tApA = make_tensor<bool>(make_shape(size<1>(tAsA), size<2>(tAsA)), Stride<_1, _0>{});
+
+
+  int residue_m = M - m_coord * get<0>(blk_shape);
+  int residue_n = N - n_coord * get<1>(blk_shape);
+
+  MUTE_UNROLL
+  for (int m = 0; m < size<0>(tApA); ++m) {
+    tApA(m, 0) = get<0>(tAcA(0, m, 0)) < residue_m;
+  }
+
+  MUTE_UNROLL
+  for (int n = 0; n < size<0>(tBpB); ++n) {
+    tBpB(n, 0) = get<0>(tBcB(0, n, 0)) < residue_n;
+  }
+
+  //
+  // Prologue
+  //
+
+  // Copy gmem to rmem for the first k_tile
+  copy_if(gmem_tiled_copy_a.with(params.desc_A), tApA, tAgA(_,_,_,*k_tile_iter), tArA);
+  if constexpr (KernelTraits::IsMixedInput) {
+    copy_if(gmem_tiled_copy_b.with(params.desc_B), tBpB, tBgB(_,_,_,*k_tile_iter), tBrB_quant);
+  } else {
+    copy_if(gmem_tiled_copy_b.with(params.desc_B), tBpB, tBgB(_,_,_,*k_tile_iter), tBrB);
+  }
+
+  // update tBrS
+  if constexpr (KernelTraits::IsMixedInput) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size(tBrS); ++ni) {
+      offset_sn[ni] = get<0>(tBcB(0, ni, 0)) * params.stride_sn;
+      tBrS(ni) = scale_base[offset_sn[ni] + scale_k_cnt * params.stride_sk];
+    }
+    read_scale_flag--;
+    scale_k_cnt++;
+  }
+
+  if (--k_tile_count > 0) ++k_tile_iter;
+
+  copy(tArA, tAsA);
+  if constexpr (KernelTraits::IsMixedInput) {
+    dequant_tensor(tBrB_quant, tBrB, tBrS);
+  }
+  copy(tBrB, tBsB);
+
+  __syncthreads();
+
+
+  // Load A, B smem->rmem for k=0
+  copy(tCsA(_,_,0), tCrA_copy_view(_,_,0));
+  copy(tCsB(_,_,0), tCrB_copy_view(_,_,0));
+
+  // Size of the k-tiles's outer product mode (k)
+  auto K_BLOCK_MAX = size<2>(tCrA);
+
+  MUTLASS_PRAGMA_NO_UNROLL
+  while (k_tile_count > -1)
+  {
+    // Pipeline the outer products with a static for loop
+    //for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block)
+    MUTE_UNROLL
+    for (int k_block = 0; k_block < K_BLOCK_MAX; ++k_block)
+    {
+      if (k_block == K_BLOCK_MAX - 1)
+      {
+
+        if constexpr (KernelTraits::IsMixedInput) {
+          dequant_tensor(tBrB_quant, tBrB, tBrS);
+        }
+        __syncthreads();
+
+        // Copy rmem to smem
+        copy(tArA, tAsA);
+        copy(tBrB, tBsB);
+        __syncthreads();
+      }
+
+      // Load A, B smem->rmem for k+1
+      int k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX;     // static
+      copy(tCsA(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next));
+      copy(tCsB(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next));
+      if (k_block == 0)
+      {
+        // Copy gmem to rmem
+        copy_if(gmem_tiled_copy_a.with(params.desc_A), tApA, tAgA(_,_,_,*k_tile_iter), tArA);
+        if constexpr (KernelTraits::IsMixedInput) {
+          copy_if(gmem_tiled_copy_b.with(params.desc_B), tBpB, tBgB(_,_,_,*k_tile_iter), tBrB_quant);
+        } else {
+          copy_if(gmem_tiled_copy_b.with(params.desc_B), tBpB, tBgB(_,_,_,*k_tile_iter), tBrB);
+        }
+
+        // update tBrS
+        if constexpr (KernelTraits::IsMixedInput) {
+          if (read_scale_flag == 0 && k_tile_count > 0) {
+            read_scale_flag = KernelTraits::ScaleInterleave;
+            MUTE_UNROLL
+            for (int ni = 0; ni < size(tBrS); ++ni) {
+              tBrS(ni) = scale_base[offset_sn[ni] + scale_k_cnt * params.stride_sk];
+            }
+            scale_k_cnt++;
+          }
+          read_scale_flag--;
+        }
+        if (--k_tile_count > 0) ++k_tile_iter;
+      }
+
+      // Thread-level register gemm for k
+      // disambiguate gemm (shared with the namespace name)
+      mute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), accum);
+    };
+  }
+
+  // Epilogue
+  Tensor mC_mn = make_tensor(make_gmem_ptr(params.ptr_C), make_shape(M,N), params.stride_C);
+  Tensor gC_mn = local_tile(mC_mn, blk_shape, make_coord(_,_,_), Step<_1,_1,X>{});
+  Tensor gC = gC_mn(_,_,m_coord,n_coord);
+
+  Tensor tCgC = thr_mma.partition_C(gC);
+  Tensor cC = make_identity_tensor(take<0,2>(blk_shape));
+  Tensor tCcC = thr_mma.partition_C(cC);
+
+#if 1
+  constexpr bool is_half_output = is_same_v<typename KernelTraits::ElementC, mutlass::half_t>;
+  using DstVec = mute::conditional_t<is_half_output,
+                                     _Float16 __attribute__((vector_size(KernelTraits::AlignmentC * sizeof(typename KernelTraits::ElementC)))),
+                                     mutlass::Array<mutlass::bfloat16_t, KernelTraits::AlignmentC>>;
+  using SrcVec = float __attribute__((vector_size(KernelTraits::AlignmentC * sizeof(float))));
+
+  using v2f16_t = mute::conditional_t<is_half_output,
+                                      _Float16 __attribute__((vector_size(4))),
+                                      mutlass::Array<mutlass::bfloat16_t, 2>>;
+  using v2f32_t = float __attribute__((vector_size(8)));
+
+  // mn-view
+  Tensor tCcC_mn = make_tensor(tCcC.data(), convert_layout_acc_rowcol(tCcC.layout()));
+  Tensor accum_mn = make_tensor(accum.data(), convert_layout_acc_rowcol(accum.layout()));
+
+  static_assert(size<1>(accum_mn) == KernelTraits::AlignmentC);
+
+  typename KernelTraits::ElementC* ptr_C_base = params.ptr_C + m_coord * get<0>(blk_shape) * N + n_coord * get<1>(blk_shape);
+
+
+  if (residue_m >= get<0>(blk_shape) && residue_n >= get<1>(blk_shape)) {
+    MUTE_UNROLL
+    for (int mi = 0; mi < size<0>(accum_mn); ++mi) {
+      SrcVec src_vec;
+      DstVec dst_vec;
+
+      MUTE_UNROLL
+      for (int ni = 0; ni < size<1>(accum_mn) / 2; ++ni) {
+        v2f32_t f32x2_val {accum_mn(mi, ni * 2 + 0), accum_mn(mi, ni * 2 + 1)};
+
+        if constexpr (is_half_output) {
+          v2f16_t f16x2_val = __builtin_convertvector(f32x2_val, v2f16_t);
+
+          dst_vec[ni * 2 + 0] = f16x2_val[0];
+          dst_vec[ni * 2 + 1] = f16x2_val[1];
+        } else {
+          float2 f32x2 = make_float2(f32x2_val[0], f32x2_val[1]);
+          __mt_bfloat162 res = __float22bfloat162_rn(f32x2);
+          dst_vec[ni * 2 + 0] = mutlass::bfloat16_t(res.x);
+          dst_vec[ni * 2 + 1] = mutlass::bfloat16_t(res.y);
+        }
+
+      }
+
+      int offset_m = get<0>(tCcC_mn(mi,0));
+
+      *reinterpret_cast<DstVec*>(ptr_C_base + offset_m * N + get<1>(tCcC_mn(0,0))) = dst_vec;
+    }
+  } else {
+    //MUTE_UNROLL
+    //for (int mi = 0; mi < size<0>(accum_mn); ++mi) {
+    //  SrcVec src_vec;
+    //  DstVec dst_vec;
+
+    //  MUTE_UNROLL
+    //  for (int ni = 0; ni < size<1>(accum_mn) / 2; ++ni) {
+    //    v2f32_t f32x2_val {accum_mn(mi, ni * 2 + 0), accum_mn(mi, ni * 2 + 1)};
+
+    //    if constexpr (is_half_output) {
+    //      v2f16_t f16x2_val = __builtin_convertvector(f32x2_val, v2f16_t);
+
+    //      dst_vec[ni * 2 + 0] = f16x2_val[0];
+    //      dst_vec[ni * 2 + 1] = f16x2_val[1];
+    //    } else {
+    //      float2 f32x2 = make_float2(f32x2_val[0], f32x2_val[1]);
+    //      __mt_bfloat162 res = __float22bfloat162_rn(f32x2);
+    //      dst_vec[ni * 2 + 0] = mutlass::bfloat16_t(res.x);
+    //      dst_vec[ni * 2 + 1] = mutlass::bfloat16_t(res.y);
+    //    }
+    //  }
+
+    //  int offset_m = get<0>(tCcC_mn(mi,0));
+
+    //  if (offset_m < residue_m) {
+    //    *reinterpret_cast<DstVec*>(ptr_C_base + offset_m * N + get<1>(tCcC_mn(0,0))) = dst_vec;
+    //  }
+    //}
+    //
+
+
+    MUTE_UNROLL
+    for (int i = 0; i < size(accum); ++i) {
+      if (elem_less(tCcC(i), make_coord(residue_m, residue_n))) {
+        tCgC(i) = typename KernelTraits::ElementC(accum(i));
+      }
+    }
+  }
+#else
+
+  MUTE_UNROLL
+  for (int i = 0; i < size(accum); ++i) {
+    if (get<0>(tCcC(i)) < residue_m) {
+      tCgC(i) = accum(i);
+    }
+  }
+#endif
+}
diff --git a/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm_kernel_traits.hpp b/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm_kernel_traits.hpp
new file mode 100644
index 000000000..179fcf5a4
--- /dev/null
+++ b/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm_kernel_traits.hpp
@@ -0,0 +1,240 @@
+#pragma once
+
+#include <mute/tensor.hpp>
+#include <mute/atom/mma_atom.hpp>
+#include <mute/atom/copy_atom.hpp>
+
+#include <mutlass/gemm/collective/collective_builder.hpp>
+
+#include "mutlass_extension.hpp"
+
+
+using namespace mute;
+
+template <
+class ElementA,
+      class ElementB,
+      class SmemLayoutA,
+      class SmemLayoutB
+>
+struct SharedStorage {
+  mute::array_aligned<ElementA, mute::cosize_v<SmemLayoutA>, 128> smem_a;
+  mute::array_aligned<ElementB, mute::cosize_v<SmemLayoutB>, 128> smem_b;
+};
+
+template <class TileN, class PermuteTileShape = TileN>
+constexpr auto
+get_permute_tile_for_operand_B() {
+  using MmaAtomN = _16;
+  static_assert(TileN{} % PermuteTileShape{} == 0);
+  static_assert(PermuteTileShape{} % MmaAtomN{} == 0);
+
+  using RepeatInsidePermuteTile = decltype(PermuteTileShape{} / MmaAtomN{});
+  using RepeatAcrossPermuteTile = decltype(TileN{} / PermuteTileShape{});
+
+  using PermuteTileN = Layout<Shape <               MmaAtomN, RepeatInsidePermuteTile, RepeatAcrossPermuteTile>,
+                              Stride<RepeatInsidePermuteTile,                      _1,        PermuteTileShape>>;
+
+  return PermuteTileN{};
+}
+
+template <>
+constexpr auto
+get_permute_tile_for_operand_B<_256, _128>() {
+  using PermuteTileN = Layout<Shape <_16, _2,   _2, _4>,
+                              Stride< _8, _1, _128, _2>>;
+
+  return PermuteTileN{};
+}
+
+template <class TileShape_MNK, class MmaOp>
+constexpr auto
+mp22_make_tiled_mma_vectorize() {
+  constexpr int TileM = size<0>(TileShape_MNK{});
+  constexpr int TileN = size<1>(TileShape_MNK{});
+
+  constexpr int MmaOpShapeM = size<0>(typename MMA_Traits<MmaOp>::Shape_MNK{});
+  constexpr int MmaOpShapeN = size<1>(typename MMA_Traits<MmaOp>::Shape_MNK{});
+
+  static_assert(TileM % MmaOpShapeM == 0 && TileN % MmaOpShapeN == 0, "Invalid TileShape");
+
+  constexpr int MaxWarpTile = 128;
+
+  constexpr int ValueLayoutM = mute::min(MaxWarpTile, TileM) / MmaOpShapeM;
+  constexpr int ValueLayoutN = mute::min(MaxWarpTile, TileN) / MmaOpShapeN;
+
+  constexpr int AtomTileM = ValueLayoutM * MmaOpShapeM;
+  constexpr int AtomTileN = ValueLayoutN * MmaOpShapeN;
+
+  constexpr int AtomExtendM = TileM / AtomTileM;
+  constexpr int AtomExtendN = TileN / AtomTileN;
+
+  using AtomLayout = decltype(make_layout(Shape<Int<AtomExtendM>, Int<AtomExtendN>, _1>{}, LayoutLeft{}));
+
+  using PermuteLayoutK = X;
+
+  using ShapeM0 = Int<MmaOpShapeM>;
+  using ShapeM1 = Int<size<0>(AtomLayout{})>;
+  using ShapeM2 = Int<ValueLayoutM>;
+
+  using StrideM0 = _1;
+  using StrideM1 = Int<MmaOpShapeM * ValueLayoutM>;
+  using StrideM2 = Int<MmaOpShapeM>;
+
+  using PermuteLayoutM = Layout<Shape < ShapeM0,  ShapeM1,  ShapeM2>,
+                                Stride<StrideM0, StrideM1, StrideM2>>;
+
+  constexpr int MaxPermuteTile = min(128, tuple_element_t<1, TileShape_MNK>{});
+  using PermuteLayoutN = decltype(get_permute_tile_for_operand_B<tuple_element_t<1, TileShape_MNK>, Int<MaxPermuteTile>>());
+
+  using PermuteLayout = mute::Tile<PermuteLayoutM, PermuteLayoutN, PermuteLayoutK>;
+
+  return TiledMMA<MMA_Atom<MmaOp>, AtomLayout, PermuteLayout>{};
+}
+
+template <
+  class TileShape_,
+  class ElementA_,
+  int AlignmentA_,
+  class ElementB_,
+  int AlignmentB_,
+  class ElementScale_,
+  class ElementC_,
+  int GroupSizeK_ = 64
+>
+struct MixedInputGemmKernelTratis {
+  using TileShape                     = TileShape_;
+  using ElementA                      = ElementA_;
+  static constexpr int AlignmentA     = AlignmentA_;
+  using ElementB                      = ElementB_;
+  static constexpr int AlignmentB     = AlignmentB_;
+  using ElementScale                  = ElementScale_;
+  using ElementC                      = ElementC_;
+
+  static_assert(is_same_v<ElementC, ElementA>, "ElementA and ElementC should be same now.");
+
+  using IdxType = int32_t;
+
+  // Both A&B are K-major
+  using StrideA = Stride<IdxType, _1>;
+  using StrideB = Stride<IdxType, _1>;
+
+  // C is N-major
+  using StrideC = Stride<IdxType, _1>;
+
+  static_assert(sizeof_bits_v<ElementA> == 16, "Only support 16bit MMA");
+  static constexpr bool IsMixedInput = sizeof_bits_v<ElementB> < sizeof_bits_v<ElementA>;
+  static constexpr int GroupSizeK = GroupSizeK_;
+  static constexpr int ScaleInterleave = GroupSizeK / tuple_element_t<2, TileShape>{};
+
+  using MmaElementA = ElementA;
+  using MmaElementB = ElementA; // MmaElementB should be ElementA whether it is mixed-input or not
+
+  using MmaOp = decltype(mutlass::gemm::collective::detail::mp22_mma_operation_select<MmaElementA, StrideA, MmaElementB, StrideB>());
+  using TiledMma = decltype(mp22_make_tiled_mma_vectorize<TileShape, MmaOp>());
+  //using TiledMma = decltype(mutlass::gemm::collective::detail::mp22_make_tensorop_tiled_mma<TileShape, MmaOp>());
+
+  // Max WarpTile is 128
+  static constexpr int AlignmentC = min(128, tuple_element_t<1, TileShape>{}) / 16;
+
+  static constexpr int NumThreads = size(TiledMma{});
+
+  using SmemLayoutAtomA = decltype(mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<MmaElementA, StrideA>());
+  using SmemCopyAtomA = Copy_Atom<DefaultCopy, MmaElementA>;
+  using SmemLayoutA = decltype(tile_to_shape(SmemLayoutAtomA{}, select<0, 2>(TileShape{})));
+
+  using SmemLayoutAtomB = decltype(mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<MmaElementB, StrideB>());
+  using SmemCopyAtomB = Copy_Atom<DefaultCopy, MmaElementB>;
+  using SmemLayoutB = decltype(tile_to_shape(SmemLayoutAtomB{}, select<1, 2>(TileShape{})));
+
+  using SharedStorage = SharedStorage<MmaElementA, MmaElementB, SmemLayoutA, SmemLayoutB>;
+
+  using GmemTiledCopyA = decltype(mutlass::gemm::collective::detail::make_simt_tiled_copy<
+                                    NumThreads, ElementA, AlignmentA, StrideA,
+                                    get<0>(TileShape{}), get<2>(TileShape{}),
+                                    MP22_ROBUST_LOAD<uint_bit_t<AlignmentA * sizeof_bits_v<ElementA>>>>());
+
+  using GmemTiledCopyB = decltype(mutlass::gemm::collective::detail::make_simt_tiled_copy<
+                                    NumThreads, ElementB, AlignmentB, StrideB,
+                                    get<1>(TileShape{}), get<2>(TileShape{}),
+                                    MP22_ROBUST_LOAD<uint_bit_t<AlignmentB * sizeof_bits_v<ElementB>>>>());
+
+  // M, N, K
+  using ProblemShape = Shape<int, int, int>;
+
+  // Args
+  struct Arguments {
+    ElementA const* ptr_A;
+    ElementB const* ptr_B;
+    ElementC* ptr_C;
+    ElementScale const* ptr_scale;
+
+    // A [M, K], K-major
+    IdxType stride_am;
+
+    // B [N, K], K-major
+    IdxType stride_bn;
+
+    // C [M, N], N-major
+    IdxType stride_cm;
+
+    // Scale [N, K/128]
+    IdxType stride_sn;
+    IdxType stride_sk;
+  };
+
+  struct Params {
+    ProblemShape problem_shape;
+    ElementA const* ptr_A;
+    StrideA stride_A;
+    RobustDescriptor desc_A;
+
+    ElementB const* ptr_B;
+    StrideB stride_B;
+    RobustDescriptor desc_B;
+
+    ElementC* ptr_C;
+    StrideC stride_C;
+
+    ElementScale const* ptr_scale;
+    IdxType stride_sn;
+    IdxType stride_sk;
+  };
+
+  static bool
+  can_implement(ProblemShape problem_shape, Arguments args) {
+    bool implementable = true;
+
+    auto [M, N, K] = problem_shape;
+    implementable = implementable && (K % tuple_element_t<2, TileShape>{} == 0);
+    implementable = implementable && (K % GroupSizeK == 0);
+
+    return implementable;
+  }
+
+  static Params
+  to_underlying_arguments(ProblemShape problem_shape, Arguments args) {
+    StrideA stride_A {args.stride_am, _1{}};
+    StrideB stride_B {args.stride_bn, _1{}};
+    StrideC stride_C {args.stride_cm, _1{}};
+
+    auto [M, N, K] = problem_shape;
+    RobustDescriptor desc_A = make_robust_desc(args.ptr_A, M * K);
+    RobustDescriptor desc_B = make_robust_desc(args.ptr_B, N * K);
+
+    return {
+      problem_shape,
+      args.ptr_A,
+      stride_A,
+      desc_A,
+      args.ptr_B,
+      stride_B,
+      desc_B,
+      args.ptr_C,
+      stride_C,
+      args.ptr_scale,
+      args.stride_sn,
+      args.stride_sk
+    };
+  }
+};
diff --git a/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mutlass_extension.hpp b/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mutlass_extension.hpp
new file mode 100644
index 000000000..33b1348e0
--- /dev/null
+++ b/vllm_musa/csrc/linear/mp22_w4a16_groupwise_gemm/mutlass_extension.hpp
@@ -0,0 +1,192 @@
+#pragma once
+
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 220)
+#define MUTE_MP22_ROBUST_EXT
+#endif
+
+#include <mute/config.hpp>
+#include <mute/atom/copy_atom.hpp>
+#include <mute/arch/copy_mp31_desc.hpp> // RobustDescriptor
+
+
+namespace mute {
+
+//
+// Robust LD
+//
+template <class TS, class TD = TS>
+struct MP22_ROBUST_LOAD
+{
+  using SRegisters = TS[1];
+  using DRegisters = TD[1];
+
+  static_assert(sizeof(TS) == sizeof(TD),
+                "MP22_ROBUST_LOAD requires sizeof(src_value_type) == sizeof(dst_value_type)");
+
+  MUTE_HOST_DEVICE static void
+  copy(TS               const& gmem_src,
+       TD                    & rmem_dst,
+       bool                    pred,
+       RobustDescriptor const& src_desc)
+
+  {
+#if defined(MUTE_MP22_ROBUST_EXT)
+    void const* gmem_ptr = pred ? reinterpret_cast<void const*>(&gmem_src)
+                         : reinterpret_cast<void const*>(src_desc.oob_addr);
+    // meet the intrinsic requirement
+    using RobustRegV3 = uint32_t __attribute__((vector_size(12)));
+    union ExtractRobustReg {
+      RobustReg reg_v4;
+      struct {
+        RobustRegV3 reg;
+        uint32_t dummy;
+      };
+    };
+
+    ExtractRobustReg desc {src_desc.reg};
+
+    if constexpr (sizeof(TS) == 1) {
+      auto val = __musa_ld_robust_i8(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 2) {
+      auto val = __musa_ld_robust_i16(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 4) {
+      auto val = __musa_ld_robust_i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 8) {
+      auto val = __musa_ld_robust_i64(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 16) {
+      auto val = __musa_ld_robust_v4i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 32) {
+      auto val = __musa_ld_robust_v8i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 64) {
+      auto val = __musa_ld_robust_v16i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else if constexpr (sizeof(TS) == 128) {
+      auto val = __musa_ld_robust_v32i32(gmem_ptr, desc.reg);
+      rmem_dst = *reinterpret_cast<TD*>(&val);
+    } else {
+      static_assert(sizeof(TS) <= 128, "Currently we only support up-to 1024-bit load");
+    }
+#else
+    MUTE_INVALID_CONTROL_PATH(
+        "Trying to use MP22_ROBUST_LOAD without MUTE_ARCH_ROBUST_BUFFER_ACCESS_ACTIVATED.");
+#endif
+  }
+};
+
+
+//
+// Robust load traits
+//
+template <class S, class D>
+struct MP22_ROBUST_LOAD_OP : MP22_ROBUST_LOAD<S, D> {};
+
+template <class S, class D>
+struct MP22_ROBUST_PRED_LOAD_OP : MP22_ROBUST_LOAD<S, D> {};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_PRED_LOAD_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  // Predicate value that determines whether to load or return zero
+  bool pred;
+
+  // Overload copy_unpack to pass robust buffer access descriptor
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst)
+  {
+    static_assert(is_gmem<TS>::value, "Expected gmem src for MP22_ROBUST_LOAD");
+    static_assert(is_rmem<TD>::value, "Expected rmem dst for MP22_ROBUST_LOAD");
+
+    Tensor rS = recast<S>(src);
+    Tensor rD = recast<D>(dst);
+
+    MUTE_STATIC_ASSERT_V(size(rS) == Int<1>{},
+        "This src layout is incompatible with this tiled copy.");
+    MUTE_STATIC_ASSERT_V(size(rD) == Int<1>{},
+        "This dst layout is incompatible with this tiled copy.");
+
+    MP22_ROBUST_LOAD<S, D>::copy(rS[0], rD[0], traits.pred, traits.src_desc);
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_LOAD_OP<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Robust Buffer Access Descriptor
+  const RobustDescriptor src_desc;
+
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP22_ROBUST_PRED_LOAD_OP<S, D>>
+  with(bool pred) const {
+    return {src_desc, pred};
+  }
+};
+
+template <class S, class D>
+struct Copy_Traits<MP22_ROBUST_LOAD<S, D>>
+{
+  // Logical thread id to thread idx (one-thread)
+  using ThrID     = Layout<_1>;
+
+  // Map from (src-thr,src-val) to bit
+  using SrcLayout = Layout<Shape<_1, Int<sizeof_bits<S>::value>>>;
+  // Map from (dst-thr,dst-val) to bit
+  using DstLayout = Layout<Shape<_1, Int<sizeof_bits<D>::value>>>;
+
+  // Reference map from (thr,val) to bit
+  using RefLayout = SrcLayout;
+
+  // Construct an executable MP22_ROBUST_LOAD with robust buffer access descriptor
+  MUTE_HOST_DEVICE constexpr
+  Copy_Traits<MP22_ROBUST_LOAD_OP<S, D>>
+  with(RobustDescriptor desc) const {
+    return {desc};
+  }
+
+  // Don't try to execute a copy with MP22_ROBUST_LOAD before calling .with()
+  template <class TS, class SLayout,
+            class TD, class DLayout>
+  MUTE_HOST_DEVICE friend constexpr
+  void
+  copy_unpack(Copy_Traits         const& traits,
+              Tensor<TS, SLayout> const& src,
+              Tensor<TD, DLayout>      & dst) = delete;
+
+};
+
+} // namespace mute
diff --git a/vllm_musa/csrc/linear/status.md b/vllm_musa/csrc/linear/status.md
new file mode 100644
index 000000000..cf8211f8e
--- /dev/null
+++ b/vllm_musa/csrc/linear/status.md
@@ -0,0 +1,8 @@
+MUDNN_ITEM(SUCCESS)
+MUDNN_ITEM(INVALID_PARAMETER)
+MUDNN_ITEM(NOT_INITIALIZED)
+MUDNN_ITEM(ALLOC_FAILED)
+MUDNN_ITEM(NOT_SUPPORTED)
+MUDNN_ITEM(INTERNAL_ERROR)
+MUDNN_ITEM(ARCH_MISMATCH)
+MUDNN_ITEM(EXECUTION_FAILED)
diff --git a/vllm_musa/csrc/moe/deep_gemm/deep_gemm.mu b/vllm_musa/csrc/moe/deep_gemm/deep_gemm.mu
new file mode 100644
index 000000000..10836a4f7
--- /dev/null
+++ b/vllm_musa/csrc/moe/deep_gemm/deep_gemm.mu
@@ -0,0 +1,278 @@
+#include <vector>
+#include <torch/all.h>
+#include <torch/csrc/utils/pycfunction_helpers.h>
+#include <torch_musa/csrc/aten/utils/Context.h>
+#include <torch_musa/csrc/aten/utils/Utils.h>
+#include <torch_musa/csrc/core/MUSAGuard.h>
+#include <tuple>
+#include <musa_robust.h>
+#include "musa.h"
+#include "../../linear/common.muh"
+#include "musa_helper.muh"
+#define MANY_STRUCT_ASM_HELPER_S64A64 1
+using struct_S64A64 = MUtensorDescriptor;
+#include "musa_asm.h"
+
+namespace
+{
+  bool is_fp8_e4m3_dtype(const at::ScalarType t)
+  {
+    return t == at::ScalarType::Float8_e4m3fn;
+  }
+}
+
+#define LAUNCH_KERNEL(KERN_NAME)                                           \
+  status = musaLaunchKernelEx(                                             \
+      &config,                                                             \
+      KERN_NAME,                                                           \
+      output_ptr, amax_ptr, tensor_desc_lhs.tensorDesc,                    \
+      tensor_desc_rhs.tensorDesc, tensor_desc_out.tensorDesc,              \
+      robust_ptr_input_bias.get_robust_ptr(),                              \
+      robust_ptr_input_z.get_robust_ptr(),                                 \
+      robust_ptr_lhs_scale.get_robust_ptr(),                               \
+      robust_ptr_rhs_scale.get_robust_ptr(),                               \
+      robust_ptr_out_scale.get_robust_ptr(), k_part_buffer,                \
+      k_part_flag, group_idx_ptr, group_scale_idx_ptr,                     \
+      group_scale_lhs_robust_vec, group_scale_rhs_robust_vec, alpha, beta, \
+      gamma, scale_lhs_mode, scale_rhs_mode, scale_out_mode, lhs_scales_m, \
+      lhs_scales_k, rhs_scales_k, rhs_scales_n, batch_stride_scale_lhs,    \
+      batch_stride_scale_rhs, quant_tile, amax_mode, act_mode, m, n, k,    \
+      num_groups, n, n, total_tile, block_tile, batch_stride_out,          \
+      batch_stride_c, k_part, k_part_block, ldkpart,                       \
+      kpart_batch_stride_c, fast_xytile.divisor, fast_xytile.multiplier,   \
+      fast_xytile.shift_right, fast_batch.divisor,                         \
+      fast_batch.multiplier, fast_batch.shift_right, cluster_dimx,         \
+      macro_bid_max, switch_swizzle, max_tile_y,                           \
+      fast_macro_tile_x.divisor, fast_macro_tile_x.multiplier,             \
+      fast_macro_tile_x.shift_right, fast_res_tile_x.divisor,              \
+      fast_res_tile_x.multiplier, fast_res_tile_x.shift_right,             \
+      fast_macro_tile_size.divisor, fast_macro_tile_size.multiplier,       \
+      fast_macro_tile_size.shift_right, tile_a_dim0, tile_a_dim1,          \
+      tile_a_dim2, tile_b_dim0, tile_b_dim1, tile_b_dim2, tile_c_dim0,     \
+      tile_c_dim1, tile_c_dim2);
+
+#define COMMON_BLOCK_INIT()                                                    \
+  dim3 block_size(1, 1, 1);                                                    \
+  dim3 grid_size{1, 1, 1};                                                     \
+  dim3 launch_grid_size{1, 1, 1};                                              \
+  int nr_threads = block_m == 128 ? 128 : 256;                                 \
+  int block_n = 128;                                                           \
+  int block_k = 128;                                                           \
+  grid_size.x = ceil_div(n, block_n);                                          \
+  grid_size.y = ceil_div(m, block_m);                                          \
+  grid_size.z = batch;                                                         \
+  block_size.x = nr_threads;                                                   \
+                                                                               \
+  launch_grid_size = grid_size;                                                \
+  int total_tile = grid_size.x * grid_size.y * grid_size.z;                    \
+  int block_tile = 0;                                                          \
+  launch_grid_size.x = min(total_tile, num_mp);                                \
+  launch_grid_size.y = 1;                                                      \
+  launch_grid_size.z = 1;                                                      \
+  block_tile = ceil_div(total_tile, launch_grid_size.x);                       \
+  FastDivmodForAsm fast_xytile = FastDivmodForAsm(grid_size.x * grid_size.y);  \
+  int cluster_dimx = 1;                                                        \
+  bool use_swizzle = true;                                                     \
+  int switch_swizzle = default_switch_swizzle;                                 \
+  int macro_tile_x = suggest_swizzle_val;                                      \
+  int max_tile_y = grid_size.y - 1;                                            \
+  int macro_tile_size = macro_tile_x * grid_size.y;                            \
+  int res_tile_x = grid_size.x % macro_tile_x;                                 \
+  if (switch_swizzle == 1)                                                     \
+  {                                                                            \
+    if (!use_swizzle)                                                          \
+    {                                                                          \
+      macro_tile_x = grid_size.y;                                              \
+    }                                                                          \
+    max_tile_y = grid_size.x - 1;                                              \
+    macro_tile_size = macro_tile_x * grid_size.x;                              \
+    res_tile_x = grid_size.y % macro_tile_x;                                   \
+  }                                                                            \
+  int macro_bid_max = 0;                                                       \
+  FastDivmodForAsm fast_macro_tile_x;                                          \
+  FastDivmodForAsm fast_res_tile_x;                                            \
+  FastDivmodForAsm fast_macro_tile_size;                                       \
+  if (use_swizzle)                                                             \
+  {                                                                            \
+    fast_macro_tile_x = FastDivmodForAsm(macro_tile_x);                        \
+    if (res_tile_x != 0)                                                       \
+    {                                                                          \
+      fast_res_tile_x = FastDivmodForAsm(res_tile_x);                          \
+    }                                                                          \
+    fast_macro_tile_size = FastDivmodForAsm(macro_tile_size);                  \
+    if (switch_swizzle == 1)                                                   \
+    {                                                                          \
+      macro_bid_max = grid_size.y / macro_tile_x * macro_tile_x * grid_size.x; \
+    }                                                                          \
+    else                                                                       \
+    {                                                                          \
+      macro_bid_max = grid_size.x / macro_tile_x * macro_tile_x * grid_size.y; \
+    }                                                                          \
+  }
+
+template <int block_m>
+void musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous_bm(
+    const std::tuple<torch::Tensor, torch::Tensor> &lhs,
+    const std::tuple<torch::Tensor, torch::Tensor> &rhs,
+    torch::Tensor &out, const torch::Tensor &m_indices)
+{
+  torch::Tensor lhs_mat = std::get<0>(lhs);
+  torch::Tensor lhs_scales = std::get<1>(lhs);
+  torch::Tensor rhs_mat = std::get<0>(rhs);
+  torch::Tensor rhs_scales = std::get<1>(rhs);
+  TORCH_CHECK(lhs_mat.dim() == 2, "lhs matrix must be a 2D tensor.");
+  TORCH_CHECK(lhs_scales.dim() == 2, "lhs scales must be a 2D tensor.");
+  TORCH_CHECK(rhs_mat.dim() == 3, "rhs matrix must be a 3D tensor.");
+  TORCH_CHECK(rhs_scales.dim() == 3, "rhs scales must be a 3D tensor.");
+  TORCH_CHECK(out.dim() == 2, "out must be a 2D tensor.");
+
+  int m = lhs_mat.size(0);
+  int k = lhs_mat.size(1);
+  int num_groups = rhs_mat.size(0);
+  int n = rhs_mat.size(1);
+  int k_ = rhs_mat.size(2);
+  int m_ = out.size(0);
+  int n_ = out.size(1);
+  int m__ = m_indices.numel();
+  int lhs_scales_m = lhs_scales.size(0);
+  int lhs_scales_k = lhs_scales.size(1);
+  int rhs_scales_groups = rhs_scales.size(0);
+  int rhs_scales_n = rhs_scales.size(1);
+  int rhs_scales_k = rhs_scales.size(2);
+  TORCH_CHECK((m == m_) && (m_ == m__), "M dimension mismatch among lhs matrix, out and m_indices.");
+  TORCH_CHECK(n == n_, "N dimension mismatch between rhs matrix and out.");
+  TORCH_CHECK(k == k_, "K dimension mismatch between lhs and rhs matrix.");
+  TORCH_CHECK((lhs_scales_m == m) && (lhs_scales_k == ceil_div(k, 128)), "Unsupported lhs_scales shape.");
+  TORCH_CHECK((rhs_scales_groups == num_groups) && (rhs_scales_n == ceil_div(n, 128) && (rhs_scales_k == ceil_div(k, 128))), "Unsupported rhs_scales shape.");
+  TORCH_CHECK(lhs_mat.is_contiguous() && rhs_mat.is_contiguous(), "lhs and rhs matrix must be contiguous tensors.");
+  TORCH_CHECK(out.is_contiguous() && m_indices.is_contiguous(), "out and m_indices must be contiguous tensors.");
+  TORCH_CHECK(lhs_scales.is_contiguous() && rhs_scales.is_contiguous(), "lhs_scales and rhs_scales must be contiguous tensors.");
+
+  const auto lhs_dtype = lhs_mat.scalar_type();
+  const auto rhs_dtype = rhs_mat.scalar_type();
+  TORCH_CHECK(is_fp8_e4m3_dtype(lhs_dtype) && is_fp8_e4m3_dtype(rhs_dtype),
+              "lhs and rhs matrix must be float8 data types.");
+  TORCH_CHECK((lhs_scales.scalar_type() == at::ScalarType::Float) && (rhs_scales.scalar_type() == at::ScalarType::Float),
+              "lhs_scales and rhs_scales must be float32 data types.");
+  TORCH_CHECK(out.scalar_type() == at::ScalarType::BFloat16, "out must be bfloat16 data type.");
+  TORCH_CHECK(m_indices.scalar_type() == at::ScalarType::Int, "m_indices must be int dtype.");
+
+  if (m == 0)
+  {
+    return;
+  }
+
+  bfloat16_t *output_ptr = (bfloat16_t *)(out.data_ptr());
+
+  size_t lhs_dim_scale = 1;
+  size_t rhs_dim_scale = 1;
+  size_t out_dim_scale = 1;
+  size_t lhs_dim0 = k / lhs_dim_scale;
+  size_t lhs_dim1 = m;
+  size_t rhs_dim0 = k / rhs_dim_scale;
+  size_t rhs_dim1 = n;
+  constexpr int batch = 1;
+  FastDivmodForAsm fast_batch = FastDivmodForAsm(batch);
+
+  TensorDesc tensor_desc_lhs = TensorDesc({lhs_dim0, lhs_dim1, (uint64_t)batch},
+                                          (void *)(lhs_mat.data_ptr<FP8_TYPE>()),
+                                          MU_TENSOR_DESCRIPTOR_DATA_TYPE_INT8, 1, false);
+  TensorDesc tensor_desc_rhs = TensorDesc({rhs_dim0, rhs_dim1, (uint64_t)num_groups},
+                                          (void *)(rhs_mat.data_ptr<FP8_TYPE>()),
+                                          MU_TENSOR_DESCRIPTOR_DATA_TYPE_INT8, 1, false);
+  TensorDesc tensor_desc_out = TensorDesc({n / out_dim_scale, (uint64_t)m, (uint64_t)batch},
+                                          (void *)output_ptr,
+                                          MU_TENSOR_DESCRIPTOR_DATA_TYPE_BFLOAT16, 2, false);
+
+  size_t lhs_scale_buffer_size = lhs_scales.numel() * sizeof(float);
+  size_t rhs_scale_buffer_size = rhs_scales.numel() * sizeof(float);
+
+  __musa::robust_ptr<float> robust_ptr_input_bias(nullptr, 0);
+  __musa::robust_ptr<float> robust_ptr_input_z(nullptr, 0);
+  __musa::robust_ptr<float> robust_ptr_lhs_scale((float *)(lhs_scales.data_ptr<float>()), lhs_scale_buffer_size);
+  __musa::robust_ptr<float> robust_ptr_rhs_scale((float *)(rhs_scales.data_ptr<float>()), rhs_scale_buffer_size);
+  __musa::robust_ptr<float> robust_ptr_out_scale(nullptr, 0);
+
+  int warpg_m = 128;
+  int warpg_n = 128;
+  int tme_warp = 1;
+  void *amax_ptr = nullptr;
+  void *k_part_buffer = nullptr;
+  int *k_part_flag = nullptr;
+  int *group_idx_ptr = m_indices.data_ptr<int>();
+  int *group_scale_idx_ptr = nullptr;
+  int *group_scale_lhs_robust_vec = nullptr;
+  int *group_scale_rhs_robust_vec = nullptr;
+  float alpha = 1.f;
+  float beta = 0.f;
+  float gamma = 0.f;
+  int scale_lhs_mode = 4; // lhs matrix per-group quant
+  int scale_rhs_mode = 3; // rhs matrix per-block quant
+  int scale_out_mode = 0;
+  int batch_stride_scale_lhs = lhs_scales_m * lhs_scales_k;
+  int batch_stride_scale_rhs = rhs_scales_n * rhs_scales_k;
+  int quant_tile = 128; // quantization tile size
+  int amax_mode = 0;    // no amax
+  int act_mode = 0;     // no activation
+  size_t batch_stride_out = m * n;
+  size_t batch_stride_c = batch_stride_out;
+  int k_part = 1;
+  int k_part_block = k;
+  int ldkpart = ceil_div(n, 128) * 128;
+  size_t kpart_batch_stride_c = 0;
+  int default_switch_swizzle = 0;
+  int device;
+  MUSART_CHECK(musaGetDevice(&device));
+  musaDeviceProp device_prop;
+  MUSART_CHECK(musaGetDeviceProperties(&device_prop, device));
+  int num_mp = device_prop.multiProcessorCount;
+  int suggest_swizzle_val = num_mp % 4 == 0 ? 4 : 2;
+  COMMON_BLOCK_INIT();
+  int tile_a_dim0 = block_k;
+  int tile_a_dim1 = warpg_m / tme_warp;
+  int tile_a_dim2 = 1;
+  int tile_b_dim0 = block_k;
+  int tile_b_dim1 = warpg_n / tme_warp;
+  int tile_b_dim2 = 1;
+  int tile_c_dim0 = warpg_n;
+  int tile_c_dim1 = 4;
+  int tile_c_dim2 = 1;
+
+  const musaStream_t stream = at::musa::getCurrentMUSAStream();
+  musaLaunchConfig_t config = {0};
+  config.gridDim = launch_grid_size;
+  config.blockDim = block_size;
+  config.dynamicSmemBytes = 0;
+  config.stream = stream;
+
+  musaError_t status = musaSuccess;
+  if constexpr (block_m == 128)
+  {
+    LAUNCH_KERNEL(e4m3e4m3bf16bf16gemm_gm3_nt_tce_128_128x128B128_epilogue_group_block_128);
+  }
+  else
+  {
+    LAUNCH_KERNEL(e4m3e4m3bf16bf16gemm_gm3_nt_tce_256_256x128B128_epilogue_group_block_128);
+  }
+  MUSART_CHECK(status);
+  MUSART_CHECK(musaGetLastError());
+}
+
+void musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
+    const std::tuple<torch::Tensor, torch::Tensor> &lhs,
+    const std::tuple<torch::Tensor, torch::Tensor> &rhs,
+    torch::Tensor &out, const torch::Tensor &m_indices, const int block_m)
+{
+  if (block_m == 128)
+  {
+    musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous_bm<128>(lhs, rhs, out, m_indices);
+  }
+  else if (block_m == 256)
+  {
+    musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous_bm<256>(lhs, rhs, out, m_indices);
+  }
+  else
+  {
+    TORCH_CHECK(false, "Unsupported block m size: ", block_m);
+  }
+}
diff --git a/vllm_musa/csrc/moe/deep_gemm/libmusa_asm_lib_gen.a b/vllm_musa/csrc/moe/deep_gemm/libmusa_asm_lib_gen.a
new file mode 100644
index 0000000000000000000000000000000000000000..3af7fd15a4cf5791f7ca014f4ece2cd41ae5484e
GIT binary patch
diff --git a/vllm_musa/csrc/moe/deep_gemm/musa_asm.h b/vllm_musa/csrc/moe/deep_gemm/musa_asm.h
new file mode 100644
index 000000000..e0e7f6e09
--- /dev/null
+++ b/vllm_musa/csrc/moe/deep_gemm/musa_asm.h
@@ -0,0 +1,2 @@
+#include "musa_tce_hgemm_multi_stage_musaasm_trans_a_False_trans_b_True_thread_per_block_128_a_dtype_fp8_e4m3_b_dtype_fp8_e4m3_bias_dtype_bfloat16_c_dtype_bfloat16_tile_m_128_tile_n_128_tile_k_bytes_128_nr_buffer_3_aqmode_4_bqmode_3_qtile_128_gmode_3_fsl_False.h"
+#include "musa_tce_hgemm_multi_stage_musaasm_trans_a_False_trans_b_True_thread_per_block_256_a_dtype_fp8_e4m3_b_dtype_fp8_e4m3_bias_dtype_bfloat16_c_dtype_bfloat16_tile_m_256_tile_n_128_tile_k_bytes_128_nr_buffer_4_aqmode_4_bqmode_3_qtile_128_gmode_3_fsl_False.h"
diff --git a/vllm_musa/csrc/moe/deep_gemm/musa_helper.muh b/vllm_musa/csrc/moe/deep_gemm/musa_helper.muh
new file mode 100644
index 000000000..1d6d2acb5
--- /dev/null
+++ b/vllm_musa/csrc/moe/deep_gemm/musa_helper.muh
@@ -0,0 +1,151 @@
+#include <vector>
+#include <numeric>
+#include "musa.h"
+#include "musa_runtime.h"
+#include "../../linear/common.muh"
+#include "../../quantization/fp8/common.muh"
+
+#ifndef MU_CHECK
+#define MU_CHECK(expression)                                                   \
+    {                                                                          \
+        MUresult status = (expression);                                        \
+        if (status != MUSA_SUCCESS)                                            \
+        {                                                                      \
+            const char *err_str;                                               \
+            muGetErrorString(status, &err_str);                                \
+            std::cerr << "musa driver failed  " << __FILE__<<":"<< __LINE__    \
+            << " : " << err_str << std::endl;                                  \
+            std::exit(EXIT_FAILURE);                                           \
+        }                                                                      \
+    }
+#endif
+
+#ifndef MUSART_CHECK
+#define MUSART_CHECK(expr)                                                            \
+    do                                                                                \
+    {                                                                                 \
+        musaError_t _err = (expr);                                                    \
+        if (_err != musaSuccess)                                                      \
+        {                                                                             \
+            fprintf(stderr, "musa runtime failed " #expr ": err=%d(%s) line=%s:%d\n", \
+                    (int)_err, musaGetErrorString(_err),__FILE__,  __LINE__);         \
+            abort();                                                                  \
+        }                                                                             \
+    } while (0)
+#endif
+
+struct TensorDesc {
+  TensorDesc(const std::vector<uint64_t> &dims, void* ptr, MUtensorDescriptorDataType tensor_dtype, int dtype_size, bool h2d = true) {
+    d_size = dtype_size;
+    tensorRank = dims.size();
+    uint32_t last_stride = 1;
+    for (size_t i = 0; i < tensorRank; ++i) {
+        globalDim[i] = dims[i];
+        if (dims[i] == 0) {
+            is_empty = true;
+        }
+        last_stride = last_stride * dims[i];
+        if (i != tensorRank - 1) {
+            globalStrides[i] = last_stride * dtype_size;
+        }
+    }
+    tensorDataType = tensor_dtype;
+    interleave = MU_TENSOR_DESCRIPTOR_INTERLEAVE_NONE;
+    MU_CHECK(muTensorDescriptorEncode(&tensorDesc, tensorDataType, tensorRank, ptr, globalDim, globalStrides, interleave, oobConstantFill));
+    if (h2d) {
+        MUSART_CHECK(musaMalloc(&desc_addr, 64));
+        MUSART_CHECK(musaMemcpy(desc_addr, (void *)&tensorDesc, 64, musaMemcpyHostToDevice));
+    }
+  }
+
+  MUtensorDescriptor tensorDesc;
+  MUtensorDescriptorDataType tensorDataType;
+  uint32_t tensorRank;
+  uint64_t globalDim[5];
+  uint64_t globalStrides[4]; // 8 * 4 = 32 , 0x20
+  MUtensorDescriptorInterleave interleave;
+  uint64_t oobConstantFill{0};
+  void *desc_addr;
+  int d_size;
+  bool is_empty{false};
+};
+
+// FastDivmod functions
+/**
+ * log2 computation, what's the
+ * difference between the below codes and
+ * log2_up/down codes?
+ */
+template <typename value_t>
+__device__ __host__ __forceinline__ value_t clz(value_t x) {
+  for (int i = 31; i >= 0; --i) {
+    if ((1 << i) & x)
+      return 31 - i;
+  }
+  return 32;
+}
+
+template <typename value_t>
+__device__ __host__ __forceinline__ value_t find_log2(value_t x) {
+  int a = int(31 - clz(x));
+  a += (x & (x - 1)) != 0;  // Round up, add 1 if not a power of 2.
+  return a;
+}
+
+/**
+ * Find divisor, using find_log2
+ */
+__device__ __host__ __forceinline__ void find_divisor(uint32_t& mul,
+                                                      uint32_t& shr,
+                                                      uint32_t denom) {
+  if ((denom & (denom - 1)) == 0) {
+    mul = 1;
+    shr = find_log2(denom);
+  } else {
+    uint32_t p = 31 + find_log2(denom);
+    unsigned m =
+        unsigned(((1ull << p) + unsigned(denom) - 1) / unsigned(denom));
+
+    mul = m;
+    shr = p;
+  }
+}
+
+/**
+ * Find quotient and remainder using device-side intrinsics
+ */
+__device__ __host__ __forceinline__ void fast_divmod(uint32_t& quo,
+                                                     uint32_t& rem,
+                                                     uint32_t src, uint32_t div,
+                                                     uint32_t mul,
+                                                     uint32_t shr) {
+  quo = (uint32_t)(((uint64_t)src * mul) >> shr);
+
+  // The remainder.
+  rem = src - (quo * div);
+}
+
+struct FastDivmodForAsm {
+  uint32_t divisor;
+  uint32_t multiplier;
+  uint32_t shift_right;
+
+  FastDivmodForAsm() : divisor(0), multiplier(0), shift_right(0) {}
+  explicit FastDivmodForAsm(uint32_t divisor_) : divisor(divisor_) {
+    find_divisor_for_asm(multiplier, shift_right, divisor);
+  }
+
+  void find_divisor_for_asm(uint32_t& mul, uint32_t& shr, uint32_t denom) {
+    if (denom == 1) {
+      mul = 0;
+      shr = 0;
+    } else {
+      uint32_t p = 31 + find_log2(denom);
+      unsigned m =
+          unsigned(((1ull << p) + unsigned(denom) - 1) / unsigned(denom));
+
+      mul = m;
+      shr = p - 32;
+    }
+  }
+};
diff --git a/vllm_musa/csrc/moe/deep_gemm/musa_tce_hgemm_multi_stage_musaasm_trans_a_False_trans_b_True_thread_per_block_128_a_dtype_fp8_e4m3_b_dtype_fp8_e4m3_bias_dtype_bfloat16_c_dtype_bfloat16_tile_m_128_tile_n_128_tile_k_bytes_128_nr_buffer_3_aqmode_4_bqmode_3_qtile_128_gmode_3_fsl_False.h b/vllm_musa/csrc/moe/deep_gemm/musa_tce_hgemm_multi_stage_musaasm_trans_a_False_trans_b_True_thread_per_block_128_a_dtype_fp8_e4m3_b_dtype_fp8_e4m3_bias_dtype_bfloat16_c_dtype_bfloat16_tile_m_128_tile_n_128_tile_k_bytes_128_nr_buffer_3_aqmode_4_bqmode_3_qtile_128_gmode_3_fsl_False.h
new file mode 100644
index 000000000..6d0f9704a
--- /dev/null
+++ b/vllm_musa/csrc/moe/deep_gemm/musa_tce_hgemm_multi_stage_musaasm_trans_a_False_trans_b_True_thread_per_block_128_a_dtype_fp8_e4m3_b_dtype_fp8_e4m3_bias_dtype_bfloat16_c_dtype_bfloat16_tile_m_128_tile_n_128_tile_k_bytes_128_nr_buffer_3_aqmode_4_bqmode_3_qtile_128_gmode_3_fsl_False.h
@@ -0,0 +1,30 @@
+
+            #ifndef MANY_INT_ASM_HELPER_3
+                #define MANY_INT_ASM_HELPER_3 1
+                struct ManyIntAsm3{int x[3];};
+                #define EXPAND_MANY_INT_HELPER_3(__MANY_INT_OBJ__) __MANY_INT_OBJ__.x[0],__MANY_INT_OBJ__.x[1],__MANY_INT_OBJ__.x[2]
+            #endif
+
+
+            #ifndef MANY_INT_ASM_HELPER_3
+                #define MANY_INT_ASM_HELPER_3 1
+                struct ManyIntAsm3{int x[3];};
+                #define EXPAND_MANY_INT_HELPER_3(__MANY_INT_OBJ__) __MANY_INT_OBJ__.x[0],__MANY_INT_OBJ__.x[1],__MANY_INT_OBJ__.x[2]
+            #endif
+
+
+            #ifndef MANY_INT_ASM_HELPER_3
+                #define MANY_INT_ASM_HELPER_3 1
+                struct ManyIntAsm3{int x[3];};
+                #define EXPAND_MANY_INT_HELPER_3(__MANY_INT_OBJ__) __MANY_INT_OBJ__.x[0],__MANY_INT_OBJ__.x[1],__MANY_INT_OBJ__.x[2]
+            #endif
+
+
+            #ifndef MANY_STRUCT_ASM_HELPER_S64A64
+                #define MANY_STRUCT_ASM_HELPER_S64A64 1
+                struct alignas(64) struct_S64A64 {
+                    int8_t x[64];
+                };
+            #endif
+
+extern "C"{__global__ void e4m3e4m3bf16bf16gemm_gm3_nt_tce_128_128x128B128_epilogue_group_block_128 (void* output_ptr,void* amax_ptr,struct_S64A64 a_desc_addr,struct_S64A64 b_desc_addr,struct_S64A64 c_desc_addr,__musa::robust_data_t input_bias,__musa::robust_data_t input_z,__musa::robust_data_t input_a_scale,__musa::robust_data_t input_b_scale,__musa::robust_data_t input_d_scale,void* k_part_buffer,int* k_part_flag_ptr,int* group_idx_ptr,int* group_scale_idx_ptr,int* group_scale_a_robust_vec,int* group_scale_b_robust_vec,float alpha,float beta,float gamma,int scale_a_mode,int scale_b_mode,int scale_d_mode,int scale_a_m,int scale_a_k,int scale_b_k,int scale_b_n,int batch_stride_scale_a,int batch_stride_scale_b,int quant_tile,int amax_mode,int act_mode,int nr_m,int nr_n,int nr_k,int nr_group,int ldc,int ldz,int total_tile,int block_tile,long batch_stride_c,long batch_stride_z,int k_part,int k_part_block,int ldkpart,long kpart_batch_stride_c,int fast_xytile_ori,int fast_xytile_mul,int fast_xytile_sft,int fast_batch_ori,int fast_batch_mul,int fast_batch_sft,int cluster_dimx,int macro_bid_max,int switch_swizzle,int max_tile_y,int macro_tile_x_ori,uint macro_tile_x_mul,uint macro_tile_x_sft,int res_tile_x_ori,uint res_tile_x_mul,uint res_tile_x_sft,int macro_tile_size_ori,uint macro_tile_size_mul,uint macro_tile_size_sft,int tile_a_dims_0,int tile_a_dims_1,int tile_a_dims_2,int tile_b_dims_0,int tile_b_dims_1,int tile_b_dims_2,int tile_c_dims_0,int tile_c_dims_1,int tile_c_dims_2);}
diff --git a/vllm_musa/csrc/moe/deep_gemm/musa_tce_hgemm_multi_stage_musaasm_trans_a_False_trans_b_True_thread_per_block_256_a_dtype_fp8_e4m3_b_dtype_fp8_e4m3_bias_dtype_bfloat16_c_dtype_bfloat16_tile_m_256_tile_n_128_tile_k_bytes_128_nr_buffer_4_aqmode_4_bqmode_3_qtile_128_gmode_3_fsl_False.h b/vllm_musa/csrc/moe/deep_gemm/musa_tce_hgemm_multi_stage_musaasm_trans_a_False_trans_b_True_thread_per_block_256_a_dtype_fp8_e4m3_b_dtype_fp8_e4m3_bias_dtype_bfloat16_c_dtype_bfloat16_tile_m_256_tile_n_128_tile_k_bytes_128_nr_buffer_4_aqmode_4_bqmode_3_qtile_128_gmode_3_fsl_False.h
new file mode 100644
index 000000000..35b874106
--- /dev/null
+++ b/vllm_musa/csrc/moe/deep_gemm/musa_tce_hgemm_multi_stage_musaasm_trans_a_False_trans_b_True_thread_per_block_256_a_dtype_fp8_e4m3_b_dtype_fp8_e4m3_bias_dtype_bfloat16_c_dtype_bfloat16_tile_m_256_tile_n_128_tile_k_bytes_128_nr_buffer_4_aqmode_4_bqmode_3_qtile_128_gmode_3_fsl_False.h
@@ -0,0 +1,30 @@
+
+            #ifndef MANY_INT_ASM_HELPER_3
+                #define MANY_INT_ASM_HELPER_3 1
+                struct ManyIntAsm3{int x[3];};
+                #define EXPAND_MANY_INT_HELPER_3(__MANY_INT_OBJ__) __MANY_INT_OBJ__.x[0],__MANY_INT_OBJ__.x[1],__MANY_INT_OBJ__.x[2]
+            #endif
+
+
+            #ifndef MANY_INT_ASM_HELPER_3
+                #define MANY_INT_ASM_HELPER_3 1
+                struct ManyIntAsm3{int x[3];};
+                #define EXPAND_MANY_INT_HELPER_3(__MANY_INT_OBJ__) __MANY_INT_OBJ__.x[0],__MANY_INT_OBJ__.x[1],__MANY_INT_OBJ__.x[2]
+            #endif
+
+
+            #ifndef MANY_INT_ASM_HELPER_3
+                #define MANY_INT_ASM_HELPER_3 1
+                struct ManyIntAsm3{int x[3];};
+                #define EXPAND_MANY_INT_HELPER_3(__MANY_INT_OBJ__) __MANY_INT_OBJ__.x[0],__MANY_INT_OBJ__.x[1],__MANY_INT_OBJ__.x[2]
+            #endif
+
+
+            #ifndef MANY_STRUCT_ASM_HELPER_S64A64
+                #define MANY_STRUCT_ASM_HELPER_S64A64 1
+                struct alignas(64) struct_S64A64 {
+                    int8_t x[64];
+                };
+            #endif
+
+extern "C"{__global__ void e4m3e4m3bf16bf16gemm_gm3_nt_tce_256_256x128B128_epilogue_group_block_128 (void* output_ptr,void* amax_ptr,struct_S64A64 a_desc_addr,struct_S64A64 b_desc_addr,struct_S64A64 c_desc_addr,__musa::robust_data_t input_bias,__musa::robust_data_t input_z,__musa::robust_data_t input_a_scale,__musa::robust_data_t input_b_scale,__musa::robust_data_t input_d_scale,void* k_part_buffer,int* k_part_flag_ptr,int* group_idx_ptr,int* group_scale_idx_ptr,int* group_scale_a_robust_vec,int* group_scale_b_robust_vec,float alpha,float beta,float gamma,int scale_a_mode,int scale_b_mode,int scale_d_mode,int scale_a_m,int scale_a_k,int scale_b_k,int scale_b_n,int batch_stride_scale_a,int batch_stride_scale_b,int quant_tile,int amax_mode,int act_mode,int nr_m,int nr_n,int nr_k,int nr_group,int ldc,int ldz,int total_tile,int block_tile,long batch_stride_c,long batch_stride_z,int k_part,int k_part_block,int ldkpart,long kpart_batch_stride_c,int fast_xytile_ori,int fast_xytile_mul,int fast_xytile_sft,int fast_batch_ori,int fast_batch_mul,int fast_batch_sft,int cluster_dimx,int macro_bid_max,int switch_swizzle,int max_tile_y,int macro_tile_x_ori,uint macro_tile_x_mul,uint macro_tile_x_sft,int res_tile_x_ori,uint res_tile_x_mul,uint res_tile_x_sft,int macro_tile_size_ori,uint macro_tile_size_mul,uint macro_tile_size_sft,int tile_a_dims_0,int tile_a_dims_1,int tile_a_dims_2,int tile_b_dims_0,int tile_b_dims_1,int tile_b_dims_2,int tile_c_dims_0,int tile_c_dims_1,int tile_c_dims_2);}
diff --git a/vllm_musa/csrc/moe/moe_align_sum_kernels.mu b/vllm_musa/csrc/moe/moe_align_sum_kernels.mu
new file mode 100644
index 000000000..d83c67577
--- /dev/null
+++ b/vllm_musa/csrc/moe/moe_align_sum_kernels.mu
@@ -0,0 +1,440 @@
+#include <torch/all.h>
+#include "../musa.h"
+
+#include <ATen/ATen.h>
+#include <THC/THCAtomics.muh> // copy MUSAMarcos.muh to torch_musa package
+
+#include "../cuda_compat.h"
+#include "../dispatch_utils.h"
+
+#define CEILDIV(x, y) (((x) + (y) - 1) / (y))
+
+namespace vllm_musa {
+namespace moe {
+
+namespace {
+__device__ __forceinline__ int32_t index(int32_t total_col, int32_t row,
+                                         int32_t col) {
+  // don't worry about overflow because num_experts is relatively small
+  return row * total_col + col;
+}
+}  // namespace
+
+template <typename scalar_t, typename token_cnts_t>
+__global__ void moe_align_block_size_kernel(scalar_t* __restrict__ topk_ids,
+                                            int32_t* sorted_token_ids,
+                                            int32_t* expert_ids,
+                                            int32_t* total_tokens_post_pad,
+                                            int32_t num_experts,
+                                            int32_t block_size, size_t numel) {
+  const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
+  const size_t start_idx = threadIdx.x * tokens_per_thread;
+
+  extern __shared__ int32_t shared_mem[];
+  int32_t* cumsum = shared_mem;  // 1d tensor with shape (num_experts + 1)
+  token_cnts_t* tokens_cnts =
+      (token_cnts_t*)(shared_mem + num_experts +
+                      1);  // 2d tensor with shape (blockDim.x + 1, num_experts)
+
+  for (int i = 0; i < num_experts; ++i) {
+    tokens_cnts[index(num_experts, threadIdx.x + 1, i)] = 0;
+  }
+
+  /**
+   * In the first step we compute token_cnts[thread_index + 1][expert_index],
+   * which counts how many tokens in the token shard of thread_index are
+   * assigned to expert expert_index.
+   */
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    ++tokens_cnts[index(num_experts, threadIdx.x + 1, topk_ids[i])];
+  }
+
+  __syncthreads();
+
+  // For each expert we accumulate the token counts from the different threads.
+  if (threadIdx.x < num_experts) {
+    tokens_cnts[index(num_experts, 0, threadIdx.x)] = 0;
+    for (int i = 1; i <= blockDim.x; ++i) {
+      tokens_cnts[index(num_experts, i, threadIdx.x)] +=
+          tokens_cnts[index(num_experts, i - 1, threadIdx.x)];
+    }
+  }
+
+  __syncthreads();
+
+  // We accumulate the token counts of all experts in thread 0.
+  if (threadIdx.x == 0) {
+    cumsum[0] = 0;
+    for (int i = 1; i <= num_experts; ++i) {
+      cumsum[i] = cumsum[i - 1] +
+                  CEILDIV(tokens_cnts[index(num_experts, blockDim.x, i - 1)],
+                          block_size) *
+                      block_size;
+    }
+    *total_tokens_post_pad = static_cast<int32_t>(cumsum[num_experts]);
+  }
+
+  __syncthreads();
+
+  /**
+   * For each expert, each thread processes the tokens of the corresponding
+   * blocks and stores the corresponding expert_id for each block.
+   */
+  if (threadIdx.x < num_experts) {
+    for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
+         i += block_size) {
+      expert_ids[i / block_size] = threadIdx.x;
+    }
+  }
+
+  /**
+   * Each thread processes a token shard, calculating the index of each token
+   * after sorting by expert number. Given the example topk_ids =
+   * [0,1,2,1,2,3,0,3,4] and block_size = 4, then the output would be [0, 6, *,
+   * *, 1, 3, *, *, 2, 4, *, *, 5, 7, *, *, 8, *, *, *], where * represents a
+   * padding value(preset in python).
+   */
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int32_t expert_id = topk_ids[i];
+    /** The cumsum[expert_id] stores the starting index of the tokens that the
+     * expert with expert_id needs to process, and
+     * tokens_cnts[threadIdx.x][expert_id] stores the indices of the tokens
+     * processed by the expert with expert_id within the current thread's token
+     * shard.
+     */
+    int32_t rank_post_pad =
+        tokens_cnts[index(num_experts, threadIdx.x, expert_id)] +
+        cumsum[expert_id];
+    sorted_token_ids[rank_post_pad] = i;
+    ++tokens_cnts[index(num_experts, threadIdx.x, expert_id)];
+  }
+}
+
+// TODO(simon): this is temporarily adapted from
+// https://github.com/sgl-project/sglang/commit/31548116a8dc8c6df7e146e0587335a59fc5b9d7
+// we did this to unblock Deepseek V3 but there should be a better
+// implementation to manage shared memory.
+template <typename scalar_t>
+__global__ void moe_align_block_size_global_mem_kernel(
+    scalar_t* __restrict__ topk_ids, int32_t* sorted_token_ids,
+    int32_t* expert_ids, int32_t* total_tokens_post_pad, int32_t num_experts,
+    int32_t block_size, size_t numel, int32_t* tokens_cnts, int32_t* cumsum) {
+  const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
+  const size_t start_idx = threadIdx.x * tokens_per_thread;
+
+  for (int i = 0; i < num_experts; ++i) {
+    tokens_cnts[index(num_experts, threadIdx.x + 1, i)] = 0;
+  }
+
+  /**
+   * In the first step we compute token_cnts[thread_index + 1][expert_index],
+   * which counts how many tokens in the token shard of thread_index are
+   * assigned to expert expert_index.
+   */
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    ++tokens_cnts[index(num_experts, threadIdx.x + 1, topk_ids[i])];
+  }
+
+  __syncthreads();
+
+  // For each expert we accumulate the token counts from the different threads.
+  if (threadIdx.x < num_experts) {
+    tokens_cnts[index(num_experts, 0, threadIdx.x)] = 0;
+    for (int i = 1; i <= blockDim.x; ++i) {
+      tokens_cnts[index(num_experts, i, threadIdx.x)] +=
+          tokens_cnts[index(num_experts, i - 1, threadIdx.x)];
+    }
+  }
+
+  __syncthreads();
+
+  // We accumulate the token counts of all experts in thread 0.
+  if (threadIdx.x == 0) {
+    cumsum[0] = 0;
+    for (int i = 1; i <= num_experts; ++i) {
+      cumsum[i] = cumsum[i - 1] +
+                  CEILDIV(tokens_cnts[index(num_experts, blockDim.x, i - 1)],
+                          block_size) *
+                      block_size;
+    }
+    *total_tokens_post_pad = cumsum[num_experts];
+  }
+
+  __syncthreads();
+
+  /**
+   * For each expert, each thread processes the tokens of the corresponding
+   * blocks and stores the corresponding expert_id for each block.
+   */
+  if (threadIdx.x < num_experts) {
+    for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
+         i += block_size) {
+      expert_ids[i / block_size] = threadIdx.x;
+    }
+  }
+
+  /**
+   * Each thread processes a token shard, calculating the index of each token
+   * after sorting by expert number. Given the example topk_ids =
+   * [0,1,2,1,2,3,0,3,4] and block_size = 4, then the output would be [0, 6, *,
+   * *, 1, 3, *, *, 2, 4, *, *, 5, 7, *, *, 8, *, *, *], where * represents a
+   * padding value(preset in python).
+   */
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int32_t expert_id = topk_ids[i];
+    /** The cumsum[expert_id] stores the starting index of the tokens that the
+     * expert with expert_id needs to process, and
+     * tokens_cnts[threadIdx.x][expert_id] stores the indices of the tokens
+     * processed by the expert with expert_id within the current thread's token
+     * shard.
+     */
+    int32_t rank_post_pad =
+        tokens_cnts[index(num_experts, threadIdx.x, expert_id)] +
+        cumsum[expert_id];
+    sorted_token_ids[rank_post_pad] = i;
+    ++tokens_cnts[index(num_experts, threadIdx.x, expert_id)];
+  }
+}
+
+// taken from
+// https://github.com/sgl-project/sglang/commit/ded9fcd09a43d5e7d5bb31a2bc3e9fc21bf65d2a
+template <typename scalar_t>
+__global__ void sgl_moe_align_block_size_kernel(
+    scalar_t* __restrict__ topk_ids, int32_t* sorted_token_ids,
+    int32_t* expert_ids, int32_t* total_tokens_post_pad, int32_t num_experts,
+    int32_t block_size, size_t numel, int32_t* cumsum) {
+  __shared__ int32_t shared_counts[32][8];
+  __shared__ int32_t local_offsets[256];
+
+  const int warp_id = threadIdx.x / 32;
+  const int lane_id = threadIdx.x % 32;
+  const int experts_per_warp = 8;
+  const int my_expert_start = warp_id * experts_per_warp;
+
+  for (int i = 0; i < experts_per_warp; ++i) {
+    if (my_expert_start + i < num_experts) {
+      shared_counts[warp_id][i] = 0;
+    }
+  }
+
+  const size_t tokens_per_thread = CEILDIV(numel, blockDim.x);
+  const size_t start_idx = threadIdx.x * tokens_per_thread;
+
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int expert_id = topk_ids[i];
+    int warp_idx = expert_id / experts_per_warp;
+    int expert_offset = expert_id % experts_per_warp;
+    atomicAdd(&shared_counts[warp_idx][expert_offset], 1);
+  }
+
+  __syncthreads();
+
+  if (threadIdx.x == 0) {
+    cumsum[0] = 0;
+    for (int i = 1; i <= num_experts; ++i) {
+      int expert_count = 0;
+      int warp_idx = (i - 1) / experts_per_warp;
+      int expert_offset = (i - 1) % experts_per_warp;
+      expert_count = shared_counts[warp_idx][expert_offset];
+
+      cumsum[i] =
+          cumsum[i - 1] + CEILDIV(expert_count, block_size) * block_size;
+    }
+    *total_tokens_post_pad = cumsum[num_experts];
+  }
+
+  __syncthreads();
+
+  if (threadIdx.x < num_experts) {
+    for (int i = cumsum[threadIdx.x]; i < cumsum[threadIdx.x + 1];
+         i += block_size) {
+      expert_ids[i / block_size] = threadIdx.x;
+    }
+    local_offsets[threadIdx.x] = cumsum[threadIdx.x];
+  }
+
+  __syncthreads();
+
+  for (int i = start_idx; i < numel && i < start_idx + tokens_per_thread; ++i) {
+    int32_t expert_id = topk_ids[i];
+    int32_t rank_post_pad = atomicAdd(&local_offsets[expert_id], 1);
+    sorted_token_ids[rank_post_pad] = i;
+  }
+}
+
+template <typename scalar_t, int TOPK>
+__global__ void moe_sum_kernel(
+    scalar_t* __restrict__ out,          // [..., d]
+    const scalar_t* __restrict__ input,  // [..., topk, d]
+    const int d) {
+  const int64_t token_idx = blockIdx.x;
+  for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
+    scalar_t x = 0.0;
+#pragma unroll
+    for (int k = 0; k < TOPK; ++k) {
+      x += VLLM_LDG(&input[token_idx * TOPK * d + k * d + idx]);
+    }
+    out[token_idx * d + idx] = x;
+  }
+}
+
+}  // namespace moe
+}  // namespace vllm_musa
+
+void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
+                          int64_t block_size, torch::Tensor sorted_token_ids,
+                          torch::Tensor experts_ids,
+                          torch::Tensor num_tokens_post_pad) {
+  const musaStream_t stream = at::musa::getCurrentCUDAStream();
+
+  int device_max_shared_mem;
+  auto dev = topk_ids.get_device();
+  musaDeviceGetAttribute(&device_max_shared_mem,
+                         musaDevAttrMaxSharedMemoryPerBlockOptin, dev);
+
+  const int32_t num_thread = max((int32_t)num_experts, WARP_SIZE);
+  const int32_t shared_mem_i32 =
+      ((num_thread + 1) * num_experts + (num_experts + 1)) * sizeof(int32_t);
+  const int32_t shared_mem_i16 =
+      ((num_thread + 1) * num_experts) * sizeof(uint16_t) +
+      (num_experts + 1) * sizeof(int32_t);
+
+  bool use_global_memory = false;
+  bool use_i16 = false;  // Use uint16_t for shared memory token counts
+  if (shared_mem_i32 < device_max_shared_mem) {
+    // Do nothing in this case. We're all set to use int32_t token counts
+  } else if (shared_mem_i16 < device_max_shared_mem &&
+             topk_ids.numel() <= 65535) {
+    // when nelements of topk_ids is smaller than 65535 (max value of uint16),
+    // element value of token_cnts would also smaller than 65535,
+    // so we can use uint16 as dtype of token_cnts
+    use_i16 = true;
+  } else {
+    use_global_memory = true;
+  }
+
+  if (use_global_memory) {
+    VLLM_DISPATCH_INTEGRAL_TYPES(
+        topk_ids.scalar_type(), "moe_align_block_size_global_mem_kernel", [&] {
+          // calc needed amount of shared mem for `tokens_cnts` and `cumsum`
+          // tensors
+          const int32_t num_thread = max((int32_t)num_experts, WARP_SIZE);
+
+          auto options_int = torch::TensorOptions()
+                                 .dtype(torch::kInt)
+                                 .device(topk_ids.device());
+          torch::Tensor token_cnts_buffer =
+              torch::empty({(num_experts + 1) * num_experts}, options_int);
+          torch::Tensor cumsum_buffer =
+              torch::empty({num_experts + 1}, options_int);
+
+          auto kernel =
+              vllm_musa::moe::moe_align_block_size_global_mem_kernel<scalar_t>;
+          kernel<<<1, num_thread, 0, stream>>>(
+              topk_ids.data_ptr<scalar_t>(),
+              sorted_token_ids.data_ptr<int32_t>(),
+              experts_ids.data_ptr<int32_t>(),
+              num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
+              topk_ids.numel(), token_cnts_buffer.data_ptr<int32_t>(),
+              cumsum_buffer.data_ptr<int32_t>());
+        });
+  } else if (use_i16) {
+    VLLM_DISPATCH_INTEGRAL_TYPES(
+        topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
+          // set dynamic shared mem
+          auto kernel =
+              vllm_musa::moe::moe_align_block_size_kernel<scalar_t, uint16_t>;
+          AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
+              (void*)kernel, shared_mem_i16));
+          kernel<<<1, num_thread, shared_mem_i16, stream>>>(
+              topk_ids.data_ptr<scalar_t>(),
+              sorted_token_ids.data_ptr<int32_t>(),
+              experts_ids.data_ptr<int32_t>(),
+              num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
+              topk_ids.numel());
+        });
+  } else {
+    VLLM_DISPATCH_INTEGRAL_TYPES(
+        topk_ids.scalar_type(), "moe_align_block_size_kernel", [&] {
+          auto kernel =
+              vllm_musa::moe::moe_align_block_size_kernel<scalar_t, int32_t>;
+          AT_CUDA_CHECK(VLLM_DevFuncAttribute_SET_MaxDynamicSharedMemorySize(
+              (void*)kernel, shared_mem_i32));
+          kernel<<<1, num_thread, shared_mem_i32, stream>>>(
+              topk_ids.data_ptr<scalar_t>(),
+              sorted_token_ids.data_ptr<int32_t>(),
+              experts_ids.data_ptr<int32_t>(),
+              num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
+              topk_ids.numel());
+        });
+  }
+}
+
+void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
+                              int64_t block_size,
+                              torch::Tensor sorted_token_ids,
+                              torch::Tensor experts_ids,
+                              torch::Tensor num_tokens_post_pad) {
+  const musaStream_t stream = at::musa::getCurrentCUDAStream();
+  VLLM_DISPATCH_INTEGRAL_TYPES(
+      topk_ids.scalar_type(), "sgl_moe_align_block_size_kernel", [&] {
+        // calc needed amount of shared mem for `tokens_cnts` and `cumsum`
+        // tensors
+        auto options_int =
+            torch::TensorOptions().dtype(torch::kInt).device(topk_ids.device());
+        // torch::Tensor token_cnts_buffer =
+        //     torch::empty({(num_experts + 1) * num_experts}, options_int);
+        torch::Tensor cumsum_buffer =
+            torch::empty({num_experts + 1}, options_int);
+
+        auto kernel = vllm_musa::moe::sgl_moe_align_block_size_kernel<scalar_t>;
+        kernel<<<1, 1024, 0, stream>>>(
+            topk_ids.data_ptr<scalar_t>(), sorted_token_ids.data_ptr<int32_t>(),
+            experts_ids.data_ptr<int32_t>(),
+            num_tokens_post_pad.data_ptr<int32_t>(), num_experts, block_size,
+            topk_ids.numel(), cumsum_buffer.data_ptr<int32_t>());
+      });
+}
+
+void moe_sum(torch::Tensor& input,   // [num_tokens, topk, hidden_size]
+             torch::Tensor& output)  // [num_tokens, hidden_size]
+{
+  const int hidden_size = input.size(-1);
+  const int num_tokens = output.numel() / hidden_size;
+  const int topk = input.size(1);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min(hidden_size, 1024));
+  const at::musa::OptionalCUDAGuard device_guard(device_of(output));
+  const musaStream_t stream = at::musa::getCurrentCUDAStream();
+
+  switch (topk) {
+    case 2:
+      VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
+        vllm_musa::moe::moe_sum_kernel<scalar_t, 2><<<grid, block, 0, stream>>>(
+            output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
+            hidden_size);
+      });
+      break;
+
+    case 3:
+      VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
+        vllm_musa::moe::moe_sum_kernel<scalar_t, 3><<<grid, block, 0, stream>>>(
+            output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
+            hidden_size);
+      });
+      break;
+
+    case 4:
+      VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "moe_sum_kernel", [&] {
+        vllm_musa::moe::moe_sum_kernel<scalar_t, 4><<<grid, block, 0, stream>>>(
+            output.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
+            hidden_size);
+      });
+      break;
+
+    default:
+      at::sum_out(output, input, 1);
+      break;
+  }
+}
diff --git a/vllm_musa/csrc/moe/moe_gemv_swiglu.mu b/vllm_musa/csrc/moe/moe_gemv_swiglu.mu
new file mode 100644
index 000000000..e18088474
--- /dev/null
+++ b/vllm_musa/csrc/moe/moe_gemv_swiglu.mu
@@ -0,0 +1,807 @@
+#include "../linear/common.muh"
+#include "../linear/dtype.muh"
+#include "../linear/gemv_swiglu.muh"
+#include <musa_runtime.h>
+#include <cassert>
+#include <mutex>
+#include "../flash_attn/impl/helper.h"
+#include "../type_convert.muh"
+#include "../dispatch_utils.h"
+
+using namespace musa::dnn;
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ == 310
+#define ThreadNumPerWarp 32
+#else
+#define ThreadNumPerWarp 128
+#endif
+
+#define SYNC_IF_NEEDED() \
+    if constexpr (BLOCK_N * BLOCK_K > ThreadNumPerWarp) { \
+        __SYNCTHREADS_LM; \
+    }
+
+// Optimized macro definitions
+#define CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, _IS_FP8, _IS_RMSNORM) \
+    if (scale_k_group_tile == 128) { \
+        moe_gemv_kernel<_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, false, false, _IS_FP8, 128, _IS_RMSNORM> \
+            <<<grid_size, block_size, shmem_size, stream>>>( \
+                static_cast<_CDTYPE*>(C.data_ptr()), \
+                static_cast<_ADTYPE*>(A.data_ptr()), \
+                static_cast<_BDTYPE*>(B.data_ptr()), \
+                static_cast<int*>(topk_ids_ptr), \
+                static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \
+                static_cast<_SCALE_DTYPE*>(a_scale_ptr), \
+                static_cast<_SCALE_DTYPE*>(b_scale_ptr), \
+                topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \
+                static_cast<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(rms_count_ptr), eps); \
+    } else { \
+        moe_gemv_kernel<_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, false, false, _IS_FP8, 64, _IS_RMSNORM> \
+            <<<grid_size, block_size, shmem_size, stream>>>( \
+                static_cast<_CDTYPE*>(C.data_ptr()), \
+                static_cast<_ADTYPE*>(A.data_ptr()), \
+                static_cast<_BDTYPE*>(B.data_ptr()), \
+                static_cast<int*>(topk_ids_ptr), \
+                static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \
+                static_cast<_SCALE_DTYPE*>(a_scale_ptr), \
+                static_cast<_SCALE_DTYPE*>(b_scale_ptr), \
+                topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \
+                static_cast<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(rms_count_ptr), eps); \
+    } \
+    return;
+
+#define RUN_SCALE_ROUTE_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _IS_FP8) \
+    if (mul_routed_weight) { \
+        if (use_swigelu) { \
+            CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, true, _IS_FP8, false) \
+        } else if(use_rms_norm) { \
+            CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, false, _IS_FP8, true) \
+        } else { \
+            CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, false, _IS_FP8, false) \
+        } \
+    } else { \
+        if (use_swigelu) { \
+            CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, true, _IS_FP8, false) \
+        } else if (use_rms_norm) { \
+            CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, false, _IS_FP8, true) \
+        } else { \
+            CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, false, _IS_FP8, false) \
+        } \
+    }
+
+#define CAL_MOE_GEMV_W4A16(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, _IS_RMS_NROM) \
+    if (is_pergroup_scale) { \
+        if (scale_k_group_tile == 128) { \
+            moe_gemv_kernel<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, true, true, false, 128, _IS_RMS_NROM> \
+                <<<grid_size, block_size, shmem_size, stream>>>( \
+                    static_cast<_ADTYPE*>(C.data_ptr()), \
+                    static_cast<_ADTYPE*>(A.data_ptr()), \
+                    static_cast<_BDTYPE*>(B.data_ptr()), \
+                    static_cast<int*>(topk_ids_ptr), \
+                    static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \
+                    static_cast<_SCALE_DTYPE*>(a_scale_ptr), \
+                    static_cast<_SCALE_DTYPE*>(b_scale_ptr), \
+                    topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \
+                    static_cast<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(rms_count_ptr), eps); \
+            return; \
+        } else { \
+            moe_gemv_kernel<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, true, true, false, 64, _IS_RMS_NROM> \
+                <<<grid_size, block_size, shmem_size, stream>>>( \
+                    static_cast<_ADTYPE*>(C.data_ptr()), \
+                    static_cast<_ADTYPE*>(A.data_ptr()), \
+                    static_cast<_BDTYPE*>(B.data_ptr()), \
+                    static_cast<int*>(topk_ids_ptr), \
+                    static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \
+                    static_cast<_SCALE_DTYPE*>(a_scale_ptr), \
+                    static_cast<_SCALE_DTYPE*>(b_scale_ptr), \
+                    topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \
+                    static_cast<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(rms_count_ptr), eps); \
+            return; \
+        } \
+    } else { \
+        moe_gemv_kernel<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, true, false, false, 1, _IS_RMS_NROM> \
+            <<<grid_size, block_size, shmem_size, stream>>>( \
+                static_cast<_ADTYPE*>(C.data_ptr()), \
+                static_cast<_ADTYPE*>(A.data_ptr()), \
+                static_cast<_BDTYPE*>(B.data_ptr()), \
+                static_cast<int*>(topk_ids_ptr), \
+                static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \
+                static_cast<_SCALE_DTYPE*>(a_scale_ptr), \
+                static_cast<_SCALE_DTYPE*>(b_scale_ptr), \
+                topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \
+                static_cast<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(rms_count_ptr), eps); \
+        return; \
+    }
+
+#define CAL_MOE_GEMV(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, _IS_RMS_NROM) \
+    moe_gemv_kernel<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, false, false, false, 1, _IS_RMS_NROM> \
+        <<<grid_size, block_size, shmem_size, stream>>>( \
+            static_cast<_ADTYPE*>(C.data_ptr()), \
+            static_cast<_ADTYPE*>(A.data_ptr()), \
+            static_cast<_BDTYPE*>(B.data_ptr()), \
+            static_cast<int*>(topk_ids_ptr), \
+            static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \
+            nullptr, \
+            nullptr, \
+            topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \
+            static_cast<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(rms_count_ptr), eps); \
+    return;
+
+#define RUN_SCALE_ROUTE(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _CAL_FUNC) \
+    if (mul_routed_weight) { \
+        if (use_swigelu) { \
+            _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, true, false) \
+        } else if (use_rms_norm) { \
+            _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, false, true) \
+        } else { \
+            _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, false, false) \
+        } \
+    } else { \
+        if (use_swigelu) { \
+            _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, true, false) \
+        } else if (use_rms_norm) { \
+            _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, false, true) \
+        } else { \
+            _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, false, false) \
+        } \
+    }
+
+#define RUN_ROUNTE_WEIGHT(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _CAL_FUNC) \
+    if (!B_scale.has_value() || B_scale->scalar_type() == at::ScalarType::Float) { \
+        RUN_SCALE_ROUTE(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, float, _CAL_FUNC) \
+    } else if (B_scale.has_value() && B_scale->scalar_type() == at::ScalarType::BFloat16) { \
+        RUN_SCALE_ROUTE(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, bfloat16_t, _CAL_FUNC) \
+    } else if (B_scale.has_value() && B_scale->scalar_type() == at::ScalarType::Half) { \
+        RUN_SCALE_ROUTE(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, float16_t, _CAL_FUNC) \
+    }
+
+#define GEN_LAUNCH_KERN_GEMV(_BLK_N, _BLK_K) \
+    { \
+        launch_kernel = [&]() { \
+            constexpr int block_n = _BLK_N; \
+            constexpr int block_k = _BLK_K; \
+            TORCH_CHECK(nr_n % block_n == 0, "gemv n need align"); \
+            TORCH_CHECK(hidden_size % block_k == 0, "gemv k need align"); \
+            dim3 block_size{block_n * block_k, 1, 1}; \
+            dim3 grid_size{(uint32_t)ceil_div(reduce_size, block_n), (uint32_t)topk, (uint32_t)bseqlen}; \
+            int shmem_size = block_n * sizeof(float) * block_k; \
+            if (use_int4_w4a16) { \
+                if (A.scalar_type() == at::ScalarType::BFloat16) { \
+                    RUN_ROUNTE_WEIGHT(bfloat16_t, int8_t, float, CAL_MOE_GEMV_W4A16) \
+                } else if (A.scalar_type() == at::ScalarType::Half) { \
+                    RUN_ROUNTE_WEIGHT(float16_t, int8_t, float, CAL_MOE_GEMV_W4A16) \
+                } \
+            } else if (is_fp8) { \
+                if (A.dtype() == at::ScalarType::BFloat16) { \
+                    RUN_SCALE_ROUTE_FP8(bfloat16_t, __mt_fp8_e4m3, bfloat16_t, float, float, true) \
+                } else { \
+                    RUN_SCALE_ROUTE_FP8(__mt_fp8_e4m3, __mt_fp8_e4m3, bfloat16_t, float, float, true) \
+                } \
+            } else { \
+                if (A.scalar_type() == at::ScalarType::BFloat16) { \
+                    RUN_ROUNTE_WEIGHT(bfloat16_t, bfloat16_t, float, CAL_MOE_GEMV) \
+                } else if (A.scalar_type() == at::ScalarType::Half) { \
+                    RUN_ROUNTE_WEIGHT(float16_t, float16_t, float, CAL_MOE_GEMV) \
+                } \
+            } \
+            printf("no support on moe gemv\n"); \
+            assert(false); \
+        }; \
+    }
+
+#define GEN_LAUNCH_KERN(_BLK_N, _BLK_K) \
+    { \
+        launch_kernel = [&]() { \
+            constexpr int block_n = _BLK_N; \
+            constexpr int block_k = _BLK_K; \
+            TORCH_CHECK(nr_n % block_n == 0, "gemv n need align"); \
+            TORCH_CHECK(hidden_size % block_k == 0, "gemv k need align"); \
+            dim3 block_size{block_n * block_k, 1, 1}; \
+            dim3 grid_size{(uint32_t)ceil_div(reduce_size, block_n), (uint32_t)topk, (uint32_t)bseqlen}; \
+            int shmem_size = block_n * sizeof(float) * block_k; \
+            if (use_int4_w4a16) { \
+                if (A.scalar_type() == at::ScalarType::BFloat16) { \
+                    if (topk_weights.scalar_type() == at::ScalarType::Float) { \
+                        RUN_ROUNTE_WEIGHT(bfloat16_t, int8_t, float, CAL_MOE_GEMV_W4A16) \
+                    } else if (topk_weights.scalar_type() == at::ScalarType::BFloat16) { \
+                        RUN_ROUNTE_WEIGHT(bfloat16_t, int8_t, bfloat16_t, CAL_MOE_GEMV_W4A16) \
+                    } \
+                } else if (A.scalar_type() == at::ScalarType::Half) { \
+                    if (topk_weights.scalar_type() == at::ScalarType::Float) { \
+                        RUN_ROUNTE_WEIGHT(float16_t, int8_t, float, CAL_MOE_GEMV_W4A16) \
+                    } else if (topk_weights.scalar_type() == at::ScalarType::Half) { \
+                        RUN_ROUNTE_WEIGHT(float16_t, int8_t, float16_t, CAL_MOE_GEMV_W4A16) \
+                    } \
+                } \
+            } else if (is_fp8) { \
+                if (A.dtype() == at::ScalarType::BFloat16) { \
+                    RUN_SCALE_ROUTE_FP8(bfloat16_t, __mt_fp8_e4m3, bfloat16_t, float, float, true) \
+                } else { \
+                    RUN_SCALE_ROUTE_FP8(__mt_fp8_e4m3, __mt_fp8_e4m3, bfloat16_t, float, float, true) \
+                } \
+            } else { \
+                if (A.scalar_type() == at::ScalarType::BFloat16) { \
+                    if (topk_weights.scalar_type() == at::ScalarType::Float) { \
+                        RUN_ROUNTE_WEIGHT(bfloat16_t, bfloat16_t, float, CAL_MOE_GEMV) \
+                    } else if (topk_weights.scalar_type() == at::ScalarType::BFloat16) { \
+                        RUN_ROUNTE_WEIGHT(bfloat16_t, bfloat16_t, bfloat16_t, CAL_MOE_GEMV) \
+                    } \
+                } else if (A.scalar_type() == at::ScalarType::Half) { \
+                    if (topk_weights.scalar_type() == at::ScalarType::Float) { \
+                        RUN_ROUNTE_WEIGHT(float16_t, float16_t, float, CAL_MOE_GEMV) \
+                    } else if (topk_weights.scalar_type() == at::ScalarType::Half) { \
+                        RUN_ROUNTE_WEIGHT(float16_t, float16_t, float16_t, CAL_MOE_GEMV) \
+                    } \
+                } \
+            } \
+            printf("no support on moe gemv\n"); \
+            assert(false); \
+        }; \
+    }
+
+// Main kernel implementation
+template <typename AType, typename BType, typename CType, typename ScoreType, typename ScaleType,
+          int BLOCK_N, int BLOCK_K, int iobit, bool mul_routed_weight, bool is_swigelu,
+          bool is_w4a16, bool is_per_group_scale, bool is_fp8, int scale_block, bool use_rms_norm>
+__global__ void moe_gemv_kernel(
+    CType *c_ptr,
+    const AType *a_ptr,
+    const BType *b_ptr,
+    int *expert_idx_table,
+    ScoreType *score_ptr,
+    ScaleType *scale_a,
+    ScaleType *scale_b,
+    int topk,
+    int expert_offset_stride,
+    int n, int k,
+    int nr_expert,
+    int half_n_idx,
+    int scale_k_len,
+    bfloat16_t *gamma,
+    float* sum_out,
+    volatile int *count,
+    float eps) {
+
+    constexpr int bits_of_byte = 8;
+    constexpr int Vlen = is_w4a16 ? (iobit / (sizeof(AType) * bits_of_byte)) : (iobit / (sizeof(BType) * bits_of_byte));
+    constexpr int half_blockn = BLOCK_N / 2;
+    constexpr int b_vec_bits = is_w4a16 ? 32 : 128; // TODO: There might be some room for performance improvement here
+    constexpr int w4a16_shift = is_w4a16 ? 2 : 1;
+    constexpr int scale_k_load_cntdown_init = ceil_div(scale_block, (BLOCK_K * Vlen));
+    constexpr bool fuse_castfp8 = (is_fp8 && !std::is_same_v<AType, __mt_fp8_e4m3>);
+
+    using AVecSType = std::conditional_t<std::is_same_v<AType, __mt_fp8_e4m3>, uint8_t, AType>;
+    using BVecSType = std::conditional_t<is_fp8, uint8_t, BType>;
+    using v8f32_t = float __attribute__((vector_size(32)));
+    using AVec = std::conditional_t<fuse_castfp8, v8f32_t, typename VecType<AVecSType, 128>::Ttype>;
+    using BVec = typename VecType<BVecSType, b_vec_bits>::Ttype;
+
+    int token_idx = blockIdx.z;
+    int expert_idx = blockIdx.y;
+    int real_expert_idx = 0;
+    int t_n_idx = threadIdx.x / BLOCK_K;
+    int t_k_idx = threadIdx.x % BLOCK_K;
+    int n_idx = blockIdx.x * BLOCK_N + t_n_idx;
+
+    if (expert_idx_table != nullptr) {
+        real_expert_idx = expert_idx_table[token_idx * topk + expert_idx];
+        if (real_expert_idx < 0 || real_expert_idx >= nr_expert) {
+              if constexpr (is_swigelu) {
+                if (n_idx < half_n_idx) {
+                  int offsets = (token_idx * topk + expert_idx) * half_n_idx + n_idx;
+                  c_ptr[offsets] = 0;
+                }
+              } else {
+                int offsets = (token_idx * topk + expert_idx) * half_n_idx * 2 + n_idx;
+                c_ptr[offsets] = 0;
+              }
+            return;
+        }
+    }
+
+    constexpr int thread_sum_len = is_fp8 ? Vlen / 4 : Vlen;
+    float cur_thread_sum[thread_sum_len];
+    int scale_k_load_cntdown = scale_k_load_cntdown_init;
+
+    extern __shared__ float shared_array[];
+
+    if constexpr (is_swigelu) {
+        if (t_n_idx < half_blockn) {
+            n_idx = blockIdx.x * half_blockn + t_n_idx;
+        } else {
+            n_idx = blockIdx.x * half_blockn + t_n_idx - half_blockn + half_n_idx;
+        }
+    }
+
+
+    #pragma unroll
+    for (int i = 0; i < thread_sum_len; i++) {
+        cur_thread_sum[i] = 0.0f;
+    }
+
+    float scale_a_val = 1.0f;
+    float scale_b_val = 1.0f;
+    int scale_a_offset = 0;
+    int scale_b_offset = 0;
+
+    if constexpr (is_w4a16) {
+        scale_b_offset = (real_expert_idx * n + n_idx) * scale_k_len + t_k_idx * Vlen / scale_block;
+        if constexpr (is_swigelu) {
+            scale_b_offset = (real_expert_idx * 2 * n + n_idx) * scale_k_len;
+        }
+        scale_b_val = scale_b[scale_b_offset];
+        scale_k_load_cntdown -= 1;
+    } else if (is_fp8) {
+        scale_a_offset = token_idx * scale_k_len + t_k_idx * Vlen / scale_block;
+        scale_b_offset = (real_expert_idx * n + n_idx) / scale_block * scale_k_len + t_k_idx * Vlen / scale_block;
+        if constexpr (is_swigelu) {
+            scale_b_offset = (real_expert_idx * 2 * n + n_idx) / scale_block * scale_k_len + t_k_idx * Vlen / scale_block;
+        }
+        if constexpr (!fuse_castfp8) {
+            scale_a_val = scale_a[scale_a_offset];
+        }
+        scale_b_val = scale_b[scale_b_offset];
+        scale_k_load_cntdown -= 1;
+    }
+
+    const BType *b_base_ptr = b_ptr + ((size_t)real_expert_idx * expert_offset_stride + n_idx * k + t_k_idx * Vlen) / w4a16_shift;
+
+    for (int k_idx = 0; k_idx < k; k_idx += Vlen * BLOCK_K) {
+        AType a_reg[Vlen];
+        BType b_reg[Vlen / w4a16_shift];
+        *(AVec *)(a_reg) = *(AVec *)(a_ptr + token_idx * k + t_k_idx * Vlen + k_idx);
+        *(BVec *)(b_reg) = *(BVec *)(b_base_ptr + k_idx / w4a16_shift);
+
+        if constexpr (is_w4a16 && !is_fp8) {
+            float b_reg_float[Vlen];
+            #pragma unroll
+            for (int i = 0; i < Vlen / 2; i++) {
+                if constexpr (is_per_group_scale) {
+                    uint8_t read_u8 = b_reg[i];
+                    b_reg_float[i * 2 + 0] = scale_b_val * ((float)(read_u8 & 0xF) - 8.f);
+                    b_reg_float[i * 2 + 1] = scale_b_val * ((float)(read_u8 >> 4) - 8.f);
+                } else {
+                    int8_t read_s8 = b_reg[i];
+                    b_reg_float[i * 2 + 0] = scale_b_val * (float)((int8_t)(read_s8 << 4));
+                    b_reg_float[i * 2 + 1] = scale_b_val * (float)((int8_t)(read_s8 & 0xF0));
+                }
+            }
+            if constexpr (is_per_group_scale) {
+                if (scale_k_load_cntdown == 0) {
+                    scale_b_offset += 1;
+                    scale_b_val = scale_b[scale_b_offset];
+                    scale_k_load_cntdown = scale_k_load_cntdown_init;
+                }
+                scale_k_load_cntdown -= ceil_div(BLOCK_K * Vlen, scale_block);
+            }
+            #pragma unroll
+            for (int i = 0; i < thread_sum_len; i++) {
+                cur_thread_sum[i] += b_reg_float[i] * (float)a_reg[i];
+            }
+        } else if constexpr (is_fp8) {
+            float scale_val = scale_a_val * scale_b_val;
+            if (scale_k_load_cntdown == 0) {
+                scale_a_offset += ceil_div(BLOCK_K * Vlen, scale_block);
+                scale_b_offset += ceil_div(BLOCK_K * Vlen, scale_block);
+                if constexpr (!fuse_castfp8) {
+                    scale_a_val = scale_a[scale_a_offset];
+                }
+                scale_b_val = scale_b[scale_b_offset];
+                scale_k_load_cntdown = scale_k_load_cntdown_init;
+            }
+            scale_k_load_cntdown -= 1;
+            for (int i = 0; i < thread_sum_len; i++) {
+                typedef _Float16 _half_v4 __attribute__((ext_vector_type(4)));
+                typedef _Float32 _float_v4 __attribute__((ext_vector_type(4)));
+                _half_v4 a_halfv4;
+                _half_v4 b_halfv4;
+                _float_v4 a_float4;
+                _float_v4 b_float4;
+                if constexpr (fuse_castfp8) {
+                    asm("CVT.RTE.DENORM.BST4 %0.f16, %1.e4m3;" : "=R"(b_halfv4):"R"(reinterpret_cast<float*>(b_reg)[i]));
+                    #pragma unroll
+                    for (int j = 0; j < 4; j++) {
+                        cur_thread_sum[i] += scale_val * float(a_reg[i * 4 + j]) * (b_halfv4[j]);
+                    }
+                } else {
+                    asm("CVT.RTE.DENORM.BST4 %0.f16, %1.e4m3;" : "=R"(a_halfv4):"R"(reinterpret_cast<float*>(a_reg)[i]));
+                    asm("CVT.RTE.DENORM.BST4 %0.f16, %1.e4m3;" : "=R"(b_halfv4):"R"(reinterpret_cast<float*>(b_reg)[i]));
+                    #pragma unroll
+                    for (int j = 0; j < 4; j++) {
+                        cur_thread_sum[i] += scale_val * (a_halfv4[j]) * (b_halfv4[j]);
+                    }
+                }
+            }
+        } else {
+            #pragma unroll
+            for (int i = 0; i < thread_sum_len; i++) {
+                cur_thread_sum[i] += (float)b_reg[i] * (float)a_reg[i];
+            }
+        }
+    }
+
+    float rst = 0;
+    #pragma unroll
+    for (int i = 0; i < thread_sum_len; i++) {
+        rst += cur_thread_sum[i];
+    }
+
+    if constexpr (is_w4a16 && !is_per_group_scale) {
+        rst = rst / 16.f;
+    }
+
+    if constexpr (BLOCK_K > 1) {
+        shared_array[threadIdx.x] = rst;
+        SYNC_IF_NEEDED()
+        if (threadIdx.x < BLOCK_N) {
+            rst = 0;
+            #pragma unroll
+            for (int i = 0; i < BLOCK_K; i++) {
+                rst += shared_array[threadIdx.x * BLOCK_K + i];
+            }
+        }
+        if constexpr (is_swigelu) {
+            SYNC_IF_NEEDED()
+        }
+    }
+
+    static_assert(BLOCK_N <= ThreadNumPerWarp, "BLOCK_N must less than ThreadNumPerWarp");
+    if (threadIdx.x < BLOCK_N) {
+        int dst_n_idx = blockIdx.x * BLOCK_N + threadIdx.x;
+        if constexpr (is_swigelu) {
+            dst_n_idx = blockIdx.x * half_blockn + threadIdx.x;
+        }
+
+        if constexpr (mul_routed_weight) {
+            float score = (float)score_ptr[token_idx * topk + expert_idx];
+            rst = rst * score;
+        }
+
+        if constexpr (is_swigelu) {
+            shared_array[threadIdx.x] = rst;
+            if (threadIdx.x < half_blockn) {
+                float b = shared_array[threadIdx.x + half_blockn];
+                rst = rst * sigmoid(rst) * b;
+                c_ptr[token_idx * topk * n + expert_idx * n + dst_n_idx] = rst;
+            }
+        } else if constexpr (use_rms_norm) {
+            float rms = rst * rst;
+            int count_val = 0;
+
+            for (int offset = 1; offset < BLOCK_N; offset *= 2) {
+                float peer = __shfl_xor_sync(BLOCK_N, rms, offset);
+                rms += peer;
+            }
+
+            if (threadIdx.x == 0) {
+                atomicAdd(sum_out, rms);
+                __threadfence_block();
+                atomicAdd((int*)(count), 1);
+            }
+
+            while (count_val < gridDim.x) {
+                count_val = count[0];
+            }
+
+            rms = sum_out[0];
+            rst = rst * rsqrtf(rms / n + eps) * float(gamma[dst_n_idx]);
+            c_ptr[token_idx * topk * n + expert_idx * n + dst_n_idx] = rst;
+        } else {
+            c_ptr[token_idx * topk * n + expert_idx * n + dst_n_idx] = rst;
+        }
+    }
+}
+
+struct BlockConfig {
+    int block_n;
+    int block_k;
+    float score;
+    bool valid;
+};
+
+// Entry functions
+void fused_gemv(
+    torch::Tensor &A,
+    torch::Tensor &B,
+    torch::Tensor &C,
+    const c10::optional<torch::Tensor> &A_scale,
+    const c10::optional<torch::Tensor> &B_scale,
+    bool use_int4_w4a16,
+    bool use_swigelu,
+    bool use_rms_norm,
+    const c10::optional<torch::Tensor> &gamma,
+    float eps) {
+
+    TORCH_CHECK(A.dim() == 2, "A must be dim 2.")
+    TORCH_CHECK(B.dim() == 2, "B must be dim 2.")
+
+    bool mul_routed_weight = false;
+    int topk = 1;
+    int32_t bseqlen = A.size(0);
+    int32_t hidden_size = A.size(1);
+    int32_t num_experts = 1;
+    int32_t reduce_size = B.size(0);
+    bool is_fp8 = false;
+
+    if (B.dtype() == torch::kFloat8_e4m3fn) {
+        is_fp8 = true;
+    }
+
+    int current_arch = at::musa::getMUSAArch();
+    if (current_arch < 300) {
+        if (is_fp8) {
+            TORCH_CHECK(false, "gemv moe not support Float8_e4m3fn on MUSA arch ", current_arch);
+        }
+    }
+
+    const at::musa::OptionalMUSAGuard device_guard(device_of(A));
+    musaStream_t stream = at::musa::getCurrentMUSAStream();
+
+    void *topk_ids_ptr = nullptr;
+    void *topk_weights_ptr = nullptr;
+    void *a_scale_ptr = nullptr;
+    void *b_scale_ptr = nullptr;
+
+    if (A_scale.has_value()) {
+        a_scale_ptr = A_scale.value().data_ptr();
+    }
+    if (B_scale.has_value()) {
+        b_scale_ptr = B_scale.value().data_ptr();
+    }
+
+    void *rms_gamma_ptr = nullptr;
+    void *rms_sum_out_ptr = nullptr;
+    void *rms_count_ptr = nullptr;
+
+    if (use_rms_norm && gamma.has_value()) {
+        torch::Tensor sum_out = torch::zeros({1}, A.options().dtype(torch::kFloat));
+        torch::Tensor count = torch::zeros({1}, A.options().dtype(torch::kInt));
+        rms_gamma_ptr = gamma.value().data_ptr();
+        rms_sum_out_ptr = sum_out.data_ptr();
+        rms_count_ptr = count.data_ptr();
+    }
+
+    int device;
+    musaGetDevice(&device);
+    musaDeviceProp device_prop;
+    musaGetDeviceProperties(&device_prop, device);
+    int num_mp = device_prop.multiProcessorCount;
+    int expert_offset_stride = reduce_size * hidden_size;
+    int half_n_idx = reduce_size / 2;
+    int scale_k_len = 1;
+    int scale_k_group_tile = 128;
+
+    if (use_int4_w4a16 || is_fp8) {
+        scale_k_len = B_scale->size(1);
+        if (scale_k_len != 1) {
+            scale_k_group_tile = ceil_div(hidden_size, scale_k_len);
+            TORCH_CHECK(scale_k_group_tile == 128 || scale_k_group_tile == 64, "scale_k_group_tile only support 128 or 64");
+        }
+    }
+
+    bool is_pergroup_scale = scale_k_len != 1;
+    int nr_n = use_swigelu ? reduce_size / 2 : reduce_size;
+
+    std::function<void()> launch_kernel;
+
+    BlockConfig configs[] = {
+        {8, 16, 0.f, false},
+        {16, 8, 0.f, false},
+        {32, 4, 0.f, false},
+        {4, 32, 0.f, false},
+    };
+
+    constexpr int iobit = 128;
+    const int bits_of_byte = 8;
+    const int vlen =  use_int4_w4a16 ?
+                      (iobit / (torch::elementSize(A.scalar_type()) * bits_of_byte)):
+                      (iobit / (torch::elementSize(B.scalar_type()) * bits_of_byte));
+
+    float target_ratio = static_cast<float>(reduce_size) / hidden_size;
+
+    for (auto& config : configs) {
+        int load_size = config.block_k * vlen;
+        config.valid = (load_size <= hidden_size) && (hidden_size % load_size == 0);
+
+        if (config.valid) {
+            float block_ratio = static_cast<float>(config.block_n) / config.block_k;
+            config.score = 1.0f / (1.0f + fabsf(block_ratio - target_ratio));
+        }
+
+        // printf("n: %d, k: %d, block_n %d, block_k %d score %f vaild %d\n", reduce_size, hidden_size, config.block_n, config.block_k, config.score, config.valid);
+    }
+
+    BlockConfig* best_config = new BlockConfig{32, 1, -1.0f, false};
+    for (auto& config : configs) {
+        if (config.valid && config.score > best_config->score) {
+            best_config = &config;
+        }
+    }
+
+    switch (best_config->block_n) {
+        case 4:
+            switch (best_config->block_k) {
+                case 32: GEN_LAUNCH_KERN_GEMV(4, 32); break;
+                default: TORCH_CHECK(false, "Unsupported block_k for block_n=4");
+            }
+            break;
+        case 8:
+            switch (best_config->block_k) {
+                case 16: GEN_LAUNCH_KERN_GEMV(8, 16); break;
+                default: TORCH_CHECK(false, "Unsupported block_k for block_n=8");
+            }
+            break;
+        case 16:
+            switch (best_config->block_k) {
+                case 8: GEN_LAUNCH_KERN_GEMV(16, 8); break;
+                default: TORCH_CHECK(false, "Unsupported block_k for block_n=16");
+            }
+            break;
+        case 32:
+            switch (best_config->block_k) {
+                case 4: GEN_LAUNCH_KERN_GEMV(32, 4); break;
+                case 2: GEN_LAUNCH_KERN_GEMV(32, 2); break;
+                case 1: GEN_LAUNCH_KERN_GEMV(32, 1); break;
+                default: TORCH_CHECK(false, "Unsupported block_k for block_n=32");
+            }
+            break;
+        default:
+            TORCH_CHECK(false, "Unsupported block configuration");
+    }
+
+    if (use_int4_w4a16) GEN_LAUNCH_KERN_GEMV(32, 4);
+    // printf("n: %d, k: %d, block_n %d, block_k %d score %f\n", reduce_size, hidden_size, best_config->block_n, best_config->block_k, best_config->score);
+    launch_kernel();
+}
+
+void fused_moe_gemv(
+    torch::Tensor &A,
+    torch::Tensor &B,
+    torch::Tensor &C,
+    const c10::optional<torch::Tensor> &A_scale,
+    const c10::optional<torch::Tensor> &B_scale,
+    torch::Tensor &topk_weights,
+    torch::Tensor &topk_ids,
+    torch::Tensor &sorted_token_ids,
+    torch::Tensor &expert_ids,
+    torch::Tensor &num_tokens_post_padded,
+    bool mul_routed_weight,
+    int topk,
+    bool use_int4_w4a16,
+    bool use_swigelu) {
+
+    TORCH_CHECK(A.dim() == 2, "A must be dim 2.")
+    TORCH_CHECK(B.dim() == 3, "B must be dim 3.")
+
+    int32_t bseqlen = A.size(0);
+    bool is_fp8 = false;
+    if (B.dtype() == torch::kFloat8_e4m3fn) {
+        is_fp8 = true;
+    }
+
+    bool use_rms_norm = false;
+    void *rms_gamma_ptr = nullptr;
+    void *rms_sum_out_ptr = nullptr;
+    void *rms_count_ptr = nullptr;
+    float eps = 1e-6;
+
+    int current_arch = at::musa::getMUSAArch();
+    if (current_arch < 300) {
+        if (is_fp8) {
+            TORCH_CHECK(false, "gemv moe not support Float8_e4m3fn on MUSA arch ", current_arch);
+        }
+    }
+
+    int32_t hidden_size = A.size(1);
+    int32_t num_experts = B.size(0);
+    int32_t reduce_size = B.size(1);
+
+    const at::musa::OptionalMUSAGuard device_guard(device_of(A));
+    musaStream_t stream = at::musa::getCurrentMUSAStream();
+
+    void *topk_ids_ptr = topk_ids.data_ptr();
+    void *topk_weights_ptr = topk_weights.data_ptr();
+    void *a_scale_ptr = nullptr;
+    void *b_scale_ptr = nullptr;
+
+    if (A_scale.has_value()) {
+        a_scale_ptr = A_scale.value().data_ptr();
+    }
+    if (B_scale.has_value()) {
+        b_scale_ptr = B_scale.value().data_ptr();
+    }
+
+    int device;
+    musaGetDevice(&device);
+    musaDeviceProp device_prop;
+    musaGetDeviceProperties(&device_prop, device);
+    int num_mp = device_prop.multiProcessorCount;
+    int expert_offset_stride = reduce_size * hidden_size;
+    int half_n_idx = reduce_size / 2;
+    int scale_k_len = 1;
+    int scale_k_group_tile = 128;
+
+    if (use_int4_w4a16 || is_fp8) {
+        scale_k_len = B_scale->size(2);
+        if (scale_k_len != 1) {
+            scale_k_group_tile = ceil_div(hidden_size, scale_k_len);
+            TORCH_CHECK(scale_k_group_tile == 128 || scale_k_group_tile == 64, "scale_k_group_tile only support 128 or 64");
+        }
+    }
+
+    bool is_pergroup_scale = true; // TODO should from params
+    int nr_n = use_swigelu ? reduce_size / 2 : reduce_size;
+
+
+    int blk_n_list[4] = {8, 16, 32, 64};
+    int blk_scores[4] = {0, 0, 0, 0};
+    for (int i = 0; i < 4; i++)
+    {
+        int blks = ceil_div(nr_n, blk_n_list[i]);
+        if (blks < num_mp)
+        {
+            blk_scores[i] = 1;
+            break;
+        }
+        float ratio = (float)blks / num_mp;
+        blk_scores[i] = ratio - (int)ratio;
+    }
+
+    int chose_block = 0;
+    int chose_score = 0;
+    for (int i = 0; i < 4; i++)
+    {
+        if (blk_scores[i] > chose_score)
+        {
+            chose_score = blk_scores[i];
+            chose_block = i;
+        }
+    }
+
+    constexpr int iobit = 128;
+    std::function<void()> launch_kernel;
+    switch (chose_block)
+    {
+    case 0:
+        GEN_LAUNCH_KERN(8, 16);
+        break;
+    case 1:
+        GEN_LAUNCH_KERN(16, 8);
+        break;
+    case 2:
+        GEN_LAUNCH_KERN(32, 4);
+        break;
+    case 3:
+    default:
+        GEN_LAUNCH_KERN(32, 4);
+    }
+    // musaDeviceSynchronize();
+    // auto error = musaGetLastError();
+    // if (error != musaSuccess)
+    // {
+    //     printf("musaGetDevice error before launch %s\n", musaGetErrorString(error));
+    // }
+    launch_kernel();
+    // musaDeviceSynchronize();
+    // error = musaGetLastError();
+    // if (error != musaSuccess)
+    // {
+    //     // printf("musaGetDevice error before launch %s\n", musaGetErrorString(error));
+    // }
+#if 0
+    int bench_iter = 10;
+    GpuTimer timer;
+    timer.start(stream);
+    for (int i = 0; i < bench_iter; i++)
+    {
+        launch_kernel();
+    }
+    timer.stop();
+    musaDeviceSynchronize();
+    float flops = 2.0 * bseqlen * topk * reduce_size * hidden_size;
+    float int4_dtype_ratio = use_int4_w4a16 ? 8 : 1;
+    float bandwidth = (float)bseqlen * topk * hidden_size * at::elementSize(A.scalar_type()) + (float)bseqlen * topk * reduce_size * hidden_size * at::elementSize(B.scalar_type()) / int4_dtype_ratio + (float)bseqlen * topk * reduce_size * at::elementSize(C.scalar_type());
+    float avg_ms = (float)timer.elapsed_millis() / bench_iter;
+    printf("moe_gemm_swiglu_w4a16: %.4f ms, %.2f GFLOPS, %.2f GB/s\n", avg_ms, flops / avg_ms / 1e6, bandwidth / avg_ms / 1e6);
+#endif
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/moe/mp22_fused_grouped_topk_softmax.mu b/vllm_musa/csrc/moe/mp22_fused_grouped_topk_softmax.mu
new file mode 100644
index 000000000..da7866286
--- /dev/null
+++ b/vllm_musa/csrc/moe/mp22_fused_grouped_topk_softmax.mu
@@ -0,0 +1,1139 @@
+#include <torch/all.h>
+#include <ATen/Dispatch.h>
+#include <torch_musa/csrc/aten/musa/MUSAContext.h>
+#include <torch_musa/csrc/core/MUSAGuard.h>
+
+
+#define MT_WARP_SIZE 128
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+namespace vllm_musa {
+namespace moe {
+
+template <typename T>
+struct AccumulateTypeTraits;
+
+template <>
+struct AccumulateTypeTraits<float> {
+    using type = float;
+};
+
+template <>
+struct AccumulateTypeTraits<__half> {
+    using type = float;
+};
+
+template <>
+struct AccumulateTypeTraits<__mt_bfloat16> {
+    using type = __mt_bfloat16;
+};
+
+
+/// Aligned array type
+template <
+    typename T,
+    /// Number of elements in the array
+    int N,
+    /// Alignment requirement in bytes
+    int Alignment = sizeof(T) * N
+>
+struct alignas(Alignment) AlignedArray {
+    __device__ AlignedArray() {}
+    __device__ AlignedArray(const AlignedArray& rhs) {
+        #pragma unroll
+        for (int i = 0; i < N; i++) {
+            this->data[i] = rhs.data[i];
+        }
+    }
+    __device__ AlignedArray(const T& value) {
+        #pragma unroll
+        for (int i = 0; i < N; i++) {
+            this->data[i] = value;
+        }
+    }
+
+    T data[N];
+};
+
+namespace {
+
+template <typename T>
+struct Add {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a + b;
+  }
+};
+
+template <typename T>
+struct Max {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a < b ? b : a;
+  }
+};
+
+template <typename T, int width>
+__device__ __forceinline__ T musa_shfl_xor_sync(T val, int lane_mask) {
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 220))
+  return __shfl_xor_sync(0xffffffff, val, lane_mask, width);
+#elif (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 210))
+  static_assert((width >= 2) && (width <= 128) && ((width & (width - 1)) == 0));
+
+  auto shfl_func = [&](int& var) {
+    if constexpr (width == 128) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x7f, ((~(128 - 1)) & 0x7f) << 7 | (128 - 1));
+    } else if constexpr (width == 64) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x3f, ((~(64 - 1)) & 0x7f) << 7 | (64 - 1));
+    } else if constexpr (width == 32) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x1f, ((~(32 - 1)) & 0x7f) << 7 | (32 - 1));
+    } else if constexpr (width == 16) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0xf, ((~(16 - 1)) & 0x7f) << 7 | (16 - 1));
+    } else if constexpr (width == 8) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x7, ((~(8 - 1)) & 0x7f) << 7 | (8 - 1));
+    } else if constexpr (width == 4) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x3, ((~(4 - 1)) & 0x7f) << 7 | (4 - 1));
+    } else if constexpr (width == 2) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x1, ((~(2 - 1)) & 0x7f) << 7 | (2 - 1));
+    }
+  };
+
+  if constexpr (sizeof(T) == 4) {
+    int var = *(reinterpret_cast<int32_t*>(&val));
+    int ret = shfl_func(var);
+    return *(reinterpret_cast<T*>(&ret));
+  } else {
+    struct __Bits {
+      int __a, __b;
+    };
+    __Bits __tmp;
+    memcpy(&__tmp, &val, sizeof(val));
+    __tmp.__a = shfl_func(__tmp.__a);
+    __tmp.__b = shfl_func(__tmp.__b);
+    int64_t ret = *(reinterpret_cast<int64_t*>(&__tmp));
+    return *(reinterpret_cast<T*>(&ret));
+  }
+#else
+  return 0;
+#endif
+}
+
+
+template <typename T, template <typename> class ReduceOp, int WIDTH=128>
+__device__ __forceinline__ T warpReduce(T val) {
+    ReduceOp<T> r;
+    #pragma unroll
+    for (int offset = WIDTH / 2; offset > 0; offset >>= 1) {
+        // TODO: parameterization width
+        T b = musa_shfl_xor_sync<T, WIDTH>(val, offset);
+        val = r(val, b);
+    }
+    return val;
+}
+
+
+__device__ __host__ __forceinline__ uint32_t next_power_of_two(uint32_t n) {
+    if (n == 0)
+        return 1;
+    n -= 1;
+    n |= n >> 1;
+    n |= n >> 2;
+    n |= n >> 4;
+    n |= n >> 8;
+    n |= n >> 16;
+    return -~n;
+}
+}  // anonymous namespace
+
+
+// ====================== Grouped TopK softmax things ===============================
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    typename std::enable_if<std::is_same<float, T>::value>::type* = nullptr
+>
+__global__ void mp22_grouped_topk_softmax(const T* input, T* softmax_output, int* group_indices, const int experts_per_group, const int topk_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+
+    // threadIdx.y: indexing EXPERT_GROUP
+
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+
+    // 1. Run softmax
+    // NUM_EXPERTS = NUM_EXPERT_GROUP * experts_per_group
+    const int global_offset = row_id * NUM_EXPERT_GROUP * experts_per_group +
+                              experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = input + global_offset;
+    T* thread_write_base_ptr = softmax_output + global_offset;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+
+    T row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccessType* row_chunk_vec_ptr = reinterpret_cast<AccessType*>(&row_chunk);
+    for (int ii = 0; ii < VPT; ii++) {
+        int offset = local_col_id + ii * ELTS_PER_SUB_WARP;
+        row_chunk_vec_ptr[ii] = offset < experts_per_group ?
+                                *reinterpret_cast<const AccessType*>(thread_read_base_ptr + offset) :
+                                (T)(-INFINITY);
+    }
+
+    // 1.1 calculate maxval
+    T maxval = row_chunk[0];  // global max value
+    for (int ii = 1; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval = fmaxf(maxval, row_chunk[ii]);
+    }
+
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask *= 2) {
+        T b = musa_shfl_xor_sync<T, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < b ? b : maxval;
+    }
+    // we got maxval in each group now
+    T maxval_per_group = maxval;
+    #pragma unroll
+    for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+        T b = musa_shfl_xor_sync<T, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < b ? b : maxval;
+    }
+
+    T sum = static_cast<T>(0.0f);
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = expf(row_chunk[ii] - maxval);
+        sum += row_chunk[ii];
+    }
+
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        T b = musa_shfl_xor_sync<T, MT_WARP_SIZE>(sum, mask);
+        sum += b;
+    }
+
+    const T reciprocal_sum = static_cast<T>(1.f) / sum;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_sum;
+    }
+
+    for (int ii = local_col_id; ii < experts_per_group; ii += ELTS_PER_SUB_WARP) {
+        *reinterpret_cast<AccessType*>(thread_write_base_ptr + ii) = row_chunk_vec_ptr[ii / ELTS_PER_SUB_WARP];
+    }
+
+    // Now, we get max value within each expert group (softmaxed value)
+    maxval_per_group = expf(maxval_per_group - maxval) * reciprocal_sum;
+
+
+    // 2. Find the topk_group
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        T max_val = maxval_per_group;
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+            T other_maxval = musa_shfl_xor_sync<T, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.x == 0 && threadIdx.y == 0) {
+            group_indices[row_id * topk_group + k_idx] = group_idx;
+        }
+
+        // Finally, we clear the value in the thread with the current max if there is another iteration to run.
+        if (k_idx + 1 < topk_group) {
+            if (threadIdx.y == group_idx) {
+                maxval_per_group = static_cast<T>(-10000.f);
+            }
+        }
+
+    }
+}
+
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    typename std::enable_if<std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void mp22_grouped_topk_softmax(const T* input, T* softmax_output, int* group_indices, const int experts_per_group, const int topk_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+    using AccumulateType = float;
+    using AccumulateAccessType = AlignedArray<AccumulateType, ELTS_PER_LDG>;
+
+    // threadIdx.y: indexing EXPERT_GROUP
+
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+
+    // 1. Run softmax
+    // NUM_EXPERTS = NUM_EXPERT_GROUP * experts_per_group
+    const int global_offset = row_id * NUM_EXPERT_GROUP * experts_per_group +
+                              experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = input + global_offset;
+    T* thread_write_base_ptr = softmax_output + global_offset;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+    AccumulateType row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccumulateAccessType* row_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&row_chunk);
+    for (int ii = local_col_id; ii < experts_per_group; ii += ELTS_PER_SUB_WARP) {
+        AccessType regs = *reinterpret_cast<const AccessType*>(thread_read_base_ptr + ii);
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii / ELTS_PER_SUB_WARP;
+        #pragma unroll
+        for (int k = 0; k < ELTS_PER_LDG; k++) {
+            row_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(regs.data[k]);
+        }
+    }
+    // 1.1 calculate maxval
+    AccumulateType maxval = row_chunk[0];  // global max value
+    for (int ii = 1; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval = fmaxf(maxval, row_chunk[ii]);
+    }
+
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask *= 2) {
+        AccumulateType b = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < b ? b : maxval;
+    }
+    AccumulateType maxval_per_group = maxval;
+    #pragma unroll
+    for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+        AccumulateType b = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < b ? b : maxval;
+    }
+
+    AccumulateType sum = static_cast<AccumulateType>(0.0f);
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        // row_chunk[ii] = expf(row_chunk[ii] - maxval);
+        row_chunk[ii] = threadIdx.x * ELTS_PER_LDG < experts_per_group ?
+                        expf(row_chunk[ii] - maxval) :
+                        (AccumulateType)(0.0f);
+        sum += row_chunk[ii];
+    }
+
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType b = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(sum, mask);
+        sum += b;
+    }
+
+    const AccumulateType reciprocal_sum = static_cast<AccumulateType>(1.f) / sum;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_sum;
+    }
+
+    for (int ii = local_col_id; ii < experts_per_group; ii += ELTS_PER_SUB_WARP) {
+        AccessType regs;
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii / ELTS_PER_SUB_WARP;
+        for (int k = 0; k < ELTS_PER_LDG; k++) {
+            regs.data[k] = static_cast<T>(row_chunk_cur_vec_ptr->data[k]);
+        }
+        *reinterpret_cast<AccessType*>(thread_write_base_ptr + ii) = regs;
+    }
+
+    // Now, we get max value within each expert group
+    maxval_per_group = expf(maxval_per_group - maxval) * reciprocal_sum;
+
+    // 2. Find the topk_group
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        // if (threadIdx.x == 0 && blockIdx.x == 0 && k_idx == 0) {
+        //     printf("threadIdx.y: %d, maxval: %f\n", threadIdx.y, maxval);
+        // }
+        // low precison dtype will be fine
+        T max_val = static_cast<T>(maxval_per_group);
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+            T other_maxval = musa_shfl_xor_sync<T, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.x == 0 && threadIdx.y == 0) {
+            group_indices[row_id * topk_group + k_idx] = group_idx;
+        }
+
+        // Finally, we clear the value in the thread with the current max if there is another iteration to run.
+        if (k_idx + 1 < topk_group) {
+            if (threadIdx.y == group_idx) {
+                maxval_per_group = static_cast<T>(-10000.f);
+            }
+        }
+    }
+}
+
+// ==================== Grouped Topk Things ====================
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    bool RENORMALIZE,
+    typename std::enable_if<std::is_same<float, T>::value || std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void mp22_grouped_gating_topk(const T* gating_output, T* topk_weights, int* topk_indices, const int topk_group, const int topk, const int experts_per_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+    using AccumulateType = typename AccumulateTypeTraits<T>::type;
+    using AccumulateAccessType = AlignedArray<AccumulateType, ELTS_PER_LDG>;
+
+    // we split a warp into multiple sub-warps, each sub-warp is responsible
+    // for calculating the maximum value within an expert-group.
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+    // 1. Run softmax along expert dimension
+    const int base_offset = row_id * NUM_EXPERT_GROUP * experts_per_group + experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = gating_output + base_offset;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+    // all elements will be loaded into thread registers, this is acceptable
+    // because the kernel assumes a small number of experts (maybe <= 256)
+    AccumulateType row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccumulateAccessType* row_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&row_chunk);
+    for (int ii = 0; ii < VPT; ii++) {
+        int offset = local_col_id + ii * ELTS_PER_SUB_WARP;
+        const AccessType regs = offset < experts_per_group ?
+                                *reinterpret_cast<const AccessType*>(thread_read_base_ptr + offset) :
+                                (T)(-INFINITY);
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii;
+        if constexpr (std::is_same<AccumulateType, T>::value) {
+            *row_chunk_cur_vec_ptr = regs;
+        } else {
+            #pragma unroll
+            for (int k = 0; k < ELTS_PER_LDG; k++) {
+                row_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(regs.data[k]);
+            }
+        }
+    }
+    // 1.1 calculate max value
+    AccumulateType maxval = row_chunk[0];
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval = fmaxf(maxval, row_chunk[ii]);
+    }
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask *= 2) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < other_maxval ? other_maxval : maxval;
+    }
+    AccumulateType maxval_per_group = maxval;
+    #pragma unroll
+    for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < other_maxval ? other_maxval : maxval;
+    }
+
+    // 1.2 calculate sum value
+    AccumulateType sum = static_cast<AccumulateType>(0.f);
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        // TODO: zero sum value on invalid lanes
+        row_chunk[ii] = expf(row_chunk[ii] - maxval);
+        sum += row_chunk[ii];
+    }
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(sum, mask);
+        sum += other_sum;
+    }
+
+    // 1.3 calcluate softmax final result
+    const AccumulateType reciprocal_sum = static_cast<AccumulateType>(1.f) / sum;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_sum;
+    }
+
+    // 2. Find the topk group
+    AccumulateType mask_per_group = static_cast<AccumulateType>(0.f);
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        AccumulateType max_val = maxval_per_group;
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.y == group_idx) {
+            mask_per_group = static_cast<AccumulateType>(1.f);
+            maxval_per_group = static_cast<AccumulateType>(-10000.f);
+        }
+    }
+
+    // 3. Find the topk
+    __shared__ AccumulateType topk_weights_shm[8];
+    for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+        int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+        for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+            row_chunk[vpt_id * ELTS_PER_LDG + ii] = (local_base_offset + ii) < experts_per_group ?
+                                                    row_chunk[vpt_id * ELTS_PER_LDG + ii] * mask_per_group :
+                                                    static_cast<AccumulateType>(-10.0f);
+        }
+    }
+    for (int k_idx = 0; k_idx < topk; k_idx++) {
+        AccumulateType max_val = static_cast<AccumulateType>(-1.0f);
+        int expert = 10000;
+
+        // not all threads' row_chunk value are valid
+        for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+            int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+            for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+                if (row_chunk[vpt_id * ELTS_PER_LDG + ii] > max_val) {
+                    max_val = row_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    expert = threadIdx.y * experts_per_group + local_base_offset + ii;
+                }
+            }
+        }
+
+        #pragma unroll
+        for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            int other_expert = musa_shfl_xor_sync<int, MT_WARP_SIZE>(expert, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_expert < expert)) {
+                max_val = other_maxval;
+                expert = other_expert;
+            }
+        }
+
+        if (threadIdx.x == 0 && threadIdx.y == 0) {
+            // write something to global memory
+            const int idx = topk * row_id + k_idx;
+            topk_indices[idx] = expert;
+            if constexpr (RENORMALIZE) {
+                topk_weights_shm[k_idx] = static_cast<T>(max_val);
+            } else {
+                topk_weights[idx] = static_cast<T>(max_val);
+            }
+        }
+
+        if (k_idx + 1 < topk) {
+            const int local_expert_idx = expert % experts_per_group;
+            const int thread_to_clear = (local_expert_idx / ELTS_PER_LDG) % THREADS_PER_EXPERT_GROUP;
+            if (threadIdx.y == (expert / experts_per_group) && threadIdx.x == thread_to_clear) {
+                int row_chunk_idx_to_clear = (local_expert_idx / ELTS_PER_SUB_WARP) * ELTS_PER_LDG + local_expert_idx % ELTS_PER_LDG;
+                row_chunk[row_chunk_idx_to_clear] = static_cast<AccumulateType>(-1.0f);
+            }
+        }
+    }
+
+    // 4. renormalization
+    if constexpr (RENORMALIZE) {
+        int lane_id = (threadIdx.x + threadIdx.y * blockDim.x) % MT_WARP_SIZE;
+        AccumulateType topk_weight_reg = lane_id < 8 ? topk_weights_shm[lane_id] : static_cast<AccumulateType>(0.0f);
+        __syncthreads();
+        AccumulateType topk_sum = topk_weight_reg;
+        for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+            AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(topk_sum, mask);
+            topk_sum += other_sum;
+        }
+
+        if (lane_id < 8) {
+            topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg / topk_sum);
+        }
+    }
+}
+
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    bool RENORMALIZE,
+    typename std::enable_if<std::is_same<float, T>::value || std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void mp22_grouped_gating_topk_kernel(const T* gating_output, const T* e_score_correction_bias, T* topk_weights, int* topk_indices, const int topk_group, const int topk, const int experts_per_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+    using AccumulateType = typename AccumulateTypeTraits<T>::type;
+    using AccumulateAccessType = AlignedArray<AccumulateType, ELTS_PER_LDG>;
+
+    // USE_BIASED_SCORE: if true, use biased scores for expert selection but
+    // original scores for routing weights.
+
+
+    // we split a warp into multiple sub-warps, each sub-warp is responsible
+    // for calculating the maximum value within an expert-group.
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+    // 1. Run softmax along expert dimension
+    const int base_offset = row_id * NUM_EXPERT_GROUP * experts_per_group + experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = gating_output + base_offset;
+    const T* thread_read_score_bias_base_ptr = e_score_correction_bias + experts_per_group * threadIdx.y;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+    // all elements will be loaded into thread registers, this is acceptable
+    // because the kernel assumes a small number of experts (maybe <= 256)
+    AccumulateType row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccumulateType score_correction_bias_chunk[VPT * ELTS_PER_LDG] = {(T)(0.0f)};
+    AccumulateAccessType* row_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&row_chunk);
+    AccumulateAccessType* score_correction_bias_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&score_correction_bias_chunk);
+    for (int ii = 0; ii < VPT; ii++) {
+        int offset = local_col_id + ii * ELTS_PER_SUB_WARP;
+        bool valid_pos = offset < experts_per_group;
+        const AccessType regs = valid_pos ?
+                                *reinterpret_cast<const AccessType*>(thread_read_base_ptr + offset) :
+                                (T)(-INFINITY);
+        const AccessType score_bias_regs = valid_pos ?
+                                           *reinterpret_cast<const AccessType*>(thread_read_score_bias_base_ptr + offset) :
+                                           (T)(0.0f);
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii;
+        AccumulateAccessType* score_correction_bias_chunk_cur_vec_ptr = score_correction_bias_chunk_vec_ptr + ii;
+        if constexpr (std::is_same<AccumulateType, T>::value) {
+            *row_chunk_cur_vec_ptr = regs;
+            *score_correction_bias_chunk_cur_vec_ptr = score_bias_regs;
+        } else {
+            #pragma unroll
+            for (int k = 0; k < ELTS_PER_LDG; k++) {
+                row_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(regs.data[k]);
+                score_correction_bias_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(score_bias_regs.data[k]);
+            }
+        }
+    }
+    // 1.1 calculate max value
+    AccumulateType maxval = row_chunk[0];
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval = fmaxf(maxval, row_chunk[ii]);
+    }
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < other_maxval ? other_maxval : maxval;
+    }
+
+    // 1.2 calculate sum value
+    AccumulateType sum = static_cast<AccumulateType>(0.f);
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        // TODO: zero sum value on invalid lanes
+        row_chunk[ii] = expf(row_chunk[ii] - maxval);
+        sum += row_chunk[ii];
+    }
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(sum, mask);
+        sum += other_sum;
+    }
+
+    // 1.3 calcluate softmax final result and apply correction bias
+    const AccumulateType reciprocal_sum = static_cast<AccumulateType>(1.f) / sum;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_sum;
+        score_correction_bias_chunk[ii] += row_chunk[ii];
+    }
+
+    // 1.4 find max value in each expert group
+    AccumulateType maxval_per_group = score_correction_bias_chunk[0];
+    for (int ii = 1; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval_per_group = fmaxf(maxval_per_group, score_correction_bias_chunk[ii]);
+    }
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask <<= 1) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval_per_group, mask);
+        maxval_per_group = maxval_per_group < other_maxval ? other_maxval : maxval_per_group;
+    }
+
+    // 2. Find the topk group
+    AccumulateType mask_per_group = static_cast<AccumulateType>(0.f);
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        AccumulateType max_val = maxval_per_group;
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask <<= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.y == group_idx) {
+            mask_per_group = static_cast<AccumulateType>(1.f);
+            maxval_per_group = static_cast<AccumulateType>(-10000.f);
+        }
+    }
+
+    // 3. Find the topk
+    AccumulateType topk_weight_reg = (AccumulateType)(0.0f);
+    int topk_id_reg;
+    int lane_id = (threadIdx.x + threadIdx.y * blockDim.x) % MT_WARP_SIZE;
+    for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+        int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+        for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+            score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] = (local_base_offset + ii) < experts_per_group ?
+                                                    score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] * mask_per_group :
+                                                    static_cast<AccumulateType>(-INFINITY);
+        }
+    }
+    for (int k_idx = 0; k_idx < topk; k_idx++) {
+        AccumulateType max_val = static_cast<AccumulateType>(-INFINITY);
+        AccumulateType original_max_val = static_cast<AccumulateType>(-INFINITY);
+        int expert = 10000;
+
+        // not all threads' row_chunk value are valid
+        for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+            int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+            for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+                if (score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] > max_val) {
+                    max_val = score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    original_max_val = row_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    expert = threadIdx.y * experts_per_group + local_base_offset + ii;
+                }
+            }
+        }
+
+        #pragma unroll
+        for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            AccumulateType other_original_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(original_max_val, mask);
+            int other_expert = musa_shfl_xor_sync<int, MT_WARP_SIZE>(expert, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_expert < expert)) {
+                max_val = other_maxval;
+                original_max_val = other_original_maxval;
+                expert = other_expert;
+            }
+        }
+
+        // now we find maximum value and expert index
+        const int local_expert_idx = expert % experts_per_group;
+        const int thread_to_clear = (local_expert_idx / ELTS_PER_LDG) % THREADS_PER_EXPERT_GROUP;
+        const int chunk_idx_to_clear = (local_expert_idx / ELTS_PER_SUB_WARP) * ELTS_PER_LDG + local_expert_idx % ELTS_PER_LDG;
+        if (k_idx + 1 < topk) {
+            if (threadIdx.y == (expert / experts_per_group) && threadIdx.x == thread_to_clear) {
+                score_correction_bias_chunk[chunk_idx_to_clear] = static_cast<AccumulateType>(-INFINITY);
+            }
+        }
+
+        if (lane_id == k_idx) {
+            topk_weight_reg = static_cast<AccumulateType>(original_max_val);
+            topk_id_reg = expert;
+        }
+    }
+
+    // 4. renormalization
+    if constexpr (RENORMALIZE) {
+        AccumulateType topk_sum = topk_weight_reg;
+        for (int mask = 1; mask < next_power_of_two(topk); mask <<= 1) {
+            AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(topk_sum, mask);
+            topk_sum += other_sum;
+        }
+
+        if (lane_id < topk) {
+            topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg / topk_sum);
+            topk_indices[topk * row_id + lane_id] = topk_id_reg;
+        }
+    } else {
+        if (lane_id < topk) {
+            topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg);
+            topk_indices[topk * row_id + lane_id] = topk_id_reg;
+        }
+    }
+}
+
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    bool RENORMALIZE,
+    typename std::enable_if<std::is_same<float, T>::value || std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void mp22_grouped_gating_topk_sigmoid_kernel(const T* gating_output, const T* e_score_correction_bias, T* topk_weights, int* topk_indices, const int topk_group, const int topk, const int experts_per_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+    using AccumulateType = typename AccumulateTypeTraits<T>::type;
+    using AccumulateAccessType = AlignedArray<AccumulateType, ELTS_PER_LDG>;
+
+    // we split a warp into multiple sub-warps, each sub-warp is responsible
+    // for calculating the maximum value within an expert-group.
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+    // 1. Run softmax along expert dimension
+    const int base_offset = row_id * NUM_EXPERT_GROUP * experts_per_group + experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = gating_output + base_offset;
+    const T* thread_read_score_bias_base_ptr = e_score_correction_bias + experts_per_group * threadIdx.y;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+    // all elements will be loaded into thread registers, this is acceptable
+    // because the kernel assumes a small number of experts (maybe <= 256)
+    AccumulateType row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccumulateType score_correction_bias_chunk[VPT * ELTS_PER_LDG] = {(T)(0.0f)};
+    AccumulateAccessType* row_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&row_chunk);
+    AccumulateAccessType* score_correction_bias_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&score_correction_bias_chunk);
+    for (int ii = 0; ii < VPT; ii++) {
+        int offset = local_col_id + ii * ELTS_PER_SUB_WARP;
+        bool valid_pos = offset < experts_per_group;
+        const AccessType regs = valid_pos ?
+                                *reinterpret_cast<const AccessType*>(thread_read_base_ptr + offset) :
+                                (T)(-INFINITY);
+        const AccessType score_bias_regs = valid_pos ?
+                                           *reinterpret_cast<const AccessType*>(thread_read_score_bias_base_ptr + offset) :
+                                           (T)(0.0f);
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii;
+        AccumulateAccessType* score_correction_bias_chunk_cur_vec_ptr = score_correction_bias_chunk_vec_ptr + ii;
+        if constexpr (std::is_same<AccumulateType, T>::value) {
+            *row_chunk_cur_vec_ptr = regs;
+            *score_correction_bias_chunk_cur_vec_ptr = score_bias_regs;
+        } else {
+            #pragma unroll
+            for (int k = 0; k < ELTS_PER_LDG; k++) {
+                row_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(regs.data[k]);
+                score_correction_bias_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(score_bias_regs.data[k]);
+            }
+        }
+    }
+
+    // 1.1 run sigmoid and apply bias correction
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = 1.f / (1.f + __expf(-row_chunk[ii]));
+        score_correction_bias_chunk[ii] += row_chunk[ii];
+    }
+
+    // 1.2 find max value in each expert group
+    AccumulateType maxval_per_group = score_correction_bias_chunk[0];
+    for (int ii = 1; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval_per_group = fmaxf(maxval_per_group, score_correction_bias_chunk[ii]);
+    }
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask <<= 1) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval_per_group, mask);
+        maxval_per_group = maxval_per_group < other_maxval ? other_maxval : maxval_per_group;
+    }
+
+    // 2. Find the topk group
+    AccumulateType mask_per_group = static_cast<AccumulateType>(0.f);
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        AccumulateType max_val = maxval_per_group;
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask <<= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.y == group_idx) {
+            mask_per_group = static_cast<AccumulateType>(1.f);
+            maxval_per_group = static_cast<AccumulateType>(-10000.f);
+        }
+    }
+
+    // 3. Find the topk
+    AccumulateType topk_weight_reg = (AccumulateType)(0.0f);
+    int topk_id_reg;
+    int lane_id = (threadIdx.x + threadIdx.y * blockDim.x) % MT_WARP_SIZE;
+    for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+        int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+        for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+            score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] = (local_base_offset + ii) < experts_per_group ?
+                                                    score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] * mask_per_group :
+                                                    static_cast<AccumulateType>(-INFINITY);
+        }
+    }
+    for (int k_idx = 0; k_idx < topk; k_idx++) {
+        AccumulateType max_val = static_cast<AccumulateType>(-INFINITY);
+        AccumulateType original_max_val = static_cast<AccumulateType>(-INFINITY);
+        int expert = 10000;
+
+        // not all threads' row_chunk value are valid
+        for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+            int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+            for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+                if (score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] > max_val) {
+                    max_val = score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    original_max_val = row_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    expert = threadIdx.y * experts_per_group + local_base_offset + ii;
+                }
+            }
+        }
+
+        #pragma unroll
+        for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            AccumulateType other_original_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(original_max_val, mask);
+            int other_expert = musa_shfl_xor_sync<int, MT_WARP_SIZE>(expert, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_expert < expert)) {
+                max_val = other_maxval;
+                original_max_val = other_original_maxval;
+                expert = other_expert;
+            }
+        }
+
+        // now we find maximum value and expert index
+        const int local_expert_idx = expert % experts_per_group;
+        const int thread_to_clear = (local_expert_idx / ELTS_PER_LDG) % THREADS_PER_EXPERT_GROUP;
+        const int chunk_idx_to_clear = (local_expert_idx / ELTS_PER_SUB_WARP) * ELTS_PER_LDG + local_expert_idx % ELTS_PER_LDG;
+        if (k_idx + 1 < topk) {
+            if (threadIdx.y == (expert / experts_per_group) && threadIdx.x == thread_to_clear) {
+                score_correction_bias_chunk[chunk_idx_to_clear] = static_cast<AccumulateType>(-INFINITY);
+            }
+        }
+
+        if (lane_id == k_idx) {
+            topk_weight_reg = static_cast<AccumulateType>(original_max_val);
+            topk_id_reg = expert;
+        }
+    }
+
+    // 4. renormalization
+    if constexpr (RENORMALIZE) {
+        AccumulateType topk_sum = topk_weight_reg;
+        for (int mask = 1; mask < next_power_of_two(topk); mask <<= 1) {
+            AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(topk_sum, mask);
+            topk_sum += other_sum;
+        }
+
+        if (lane_id < topk) {
+            topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg / topk_sum);
+            topk_indices[topk * row_id + lane_id] = topk_id_reg;
+        }
+    } else {
+        if (lane_id < topk) {
+            topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg);
+            topk_indices[topk * row_id + lane_id] = topk_id_reg;
+        }
+    }
+}
+
+
+// topk_weights won't be large in decode stage, element wise kernel will be fine
+template <
+    typename T,
+    typename std::enable_if<std::is_same<float, T>::value || std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void renormalize_topk_weights_kernel(T* topk_weights, int num_tokens, int num_cols) {
+    using AccumulateType = typename AccumulateTypeTraits<T>::type;
+
+    const uint32_t bid = blockIdx.x * blockDim.y + threadIdx.y;
+    if (bid >= num_tokens) {
+        return;
+    }
+    const T* thread_read_base_ptr = topk_weights + bid * num_cols;
+    T* thread_write_base_ptr = topk_weights + bid * num_cols;
+
+    AccumulateType sum = static_cast<AccumulateType>(0.f);
+    for (int tid = threadIdx.x; tid < num_cols; tid += blockDim.x) {
+        sum += static_cast<AccumulateType>(thread_read_base_ptr[tid]);
+    }
+
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(sum, mask);
+        sum += other_sum;
+    }
+
+    for (int tid = threadIdx.x; tid < num_cols; tid += blockDim.x) {
+        thread_write_base_ptr[tid] = static_cast<T>(static_cast<AccumulateType>(thread_write_base_ptr[tid]) / sum);
+    }
+}
+
+
+#define LAUNCH_TOPK_GATING_SOFTMAX(TYPE, NUM_EXPERTS, NUM_EXPERT_GROUP)           \
+{\
+    constexpr std::size_t MAX_BYTES_PER_LDG = sizeof(TYPE);                    \
+    constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(TYPE) * (NUM_EXPERTS / NUM_EXPERT_GROUP));\
+    constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(TYPE);\
+    constexpr int VPT = MAX(1, NUM_EXPERTS / ELTS_PER_LDG / MT_WARP_SIZE);\
+    dim3 block(MT_WARP_SIZE / NUM_EXPERT_GROUP, NUM_EXPERT_GROUP);                                    \
+    dim3 grid(num_tokens);                 \
+    moe::mp22_grouped_topk_softmax<TYPE, VPT, NUM_EXPERT_GROUP, MT_WARP_SIZE / NUM_EXPERT_GROUP, ELTS_PER_LDG><<<grid, block, 0, stream>>>(    \
+        gating_output, softmax_res, group_indices,                     \
+        NUM_EXPERTS / NUM_EXPERT_GROUP, topk_group);\
+}
+
+
+#define LAUNCH_GROUPED_GATING_TOPK(TYPE, NUM_EXPERTS, NUM_EXPERT_GROUP)           \
+{\
+    constexpr std::size_t MAX_BYTES_PER_LDG = sizeof(TYPE);                    \
+    constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(TYPE) * (NUM_EXPERTS / NUM_EXPERT_GROUP));\
+    constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(TYPE);\
+    constexpr int VPT = MAX(1, NUM_EXPERTS / ELTS_PER_LDG / MT_WARP_SIZE);\
+    dim3 block(MT_WARP_SIZE / NUM_EXPERT_GROUP, NUM_EXPERT_GROUP);                                    \
+    dim3 grid(num_tokens);                 \
+    if (renormalize) {\
+        int shm_size = topk * sizeof(typename AccumulateTypeTraits<TYPE>::type);\
+        moe::mp22_grouped_gating_topk_sigmoid_kernel<TYPE, VPT, NUM_EXPERT_GROUP, MT_WARP_SIZE / NUM_EXPERT_GROUP, ELTS_PER_LDG, true><<<grid, block, 0, stream>>>(    \
+        gating_output, e_score_correction_bias, topk_weights, topk_indices,                   \
+        topk_group, topk, NUM_EXPERTS / NUM_EXPERT_GROUP);\
+    } else {\
+        moe::mp22_grouped_gating_topk_sigmoid_kernel<TYPE, VPT, NUM_EXPERT_GROUP, MT_WARP_SIZE / NUM_EXPERT_GROUP, ELTS_PER_LDG, false><<<grid, block, 0, stream>>>(    \
+        gating_output, e_score_correction_bias, topk_weights, topk_indices,                   \
+        topk_group, topk, NUM_EXPERTS / NUM_EXPERT_GROUP);\
+    }\
+}
+
+
+template <typename T>
+void MP22TopkGatingSoftmaxKernelLauncher(
+    const T* gating_output,
+    T* softmax_res,
+    int* group_indices,
+    const int num_tokens,
+    const int num_experts,
+    const int num_expert_group,
+    const int topk_group,
+    musaStream_t stream) {
+    static constexpr int WARPS_PER_TB = 1;
+
+    if (num_experts == 256) {
+        switch (num_expert_group) {
+            case 2:
+                LAUNCH_TOPK_GATING_SOFTMAX(T, 256, 2);
+                break;
+            case 4:
+                LAUNCH_TOPK_GATING_SOFTMAX(T, 256, 4);
+                break;
+            case 8:
+                LAUNCH_TOPK_GATING_SOFTMAX(T, 256, 8);
+                break;
+            case 16:
+                LAUNCH_TOPK_GATING_SOFTMAX(T, 256, 16);
+                break;
+            default: {
+                TORCH_CHECK(false, "Not implemented error");
+            }
+        }
+    } else {
+        TORCH_CHECK(false, "Not implemented error");
+    }
+}
+
+
+template <typename T>
+void MP22GroupedGatingTopkKernelLauncher(
+    T* topk_weights,
+    int* topk_indices,
+    const T* gating_output,
+    const T* e_score_correction_bias,
+    const int num_tokens,
+    const int num_experts,
+    const int num_expert_group,
+    const int topk_group,
+    const int topk,
+    const bool renormalize,
+    musaStream_t stream) {
+    static constexpr int WARPS_PER_TB = 1;
+
+    if (num_experts == 256) {
+        switch (num_expert_group) {
+            case 2:
+                LAUNCH_GROUPED_GATING_TOPK(T, 256, 2);
+                break;
+            case 4:
+                LAUNCH_GROUPED_GATING_TOPK(T, 256, 4);
+                break;
+            case 8:
+                LAUNCH_GROUPED_GATING_TOPK(T, 256, 8);
+                break;
+            case 16:
+                LAUNCH_GROUPED_GATING_TOPK(T, 256, 16);
+                break;
+            default: {
+                TORCH_CHECK(false, "Not implemented error");
+            }
+        }
+    } else {
+        TORCH_CHECK(false, "Not implemented error");
+    }
+}
+
+} // namespace moe
+} // namespace vllm_musa
+
+
+void mp22_fused_grouped_topk(
+    const torch::Tensor& topk_weights,
+    const torch::Tensor& topk_indices,
+    const torch::Tensor& gating_output,
+    const c10::optional<torch::Tensor>& e_score_correction_bias,
+    int num_expert_group,
+    int topk_group,
+    int topk,
+    bool renormalize) {
+
+    TORCH_CHECK(e_score_correction_bias.has_value());
+
+    const int num_experts = gating_output.size(-1);
+    const int num_tokens = gating_output.numel() / num_experts;
+
+    const at::musa::OptionalMUSAGuard device_guard(device_of(gating_output));
+    const musaStream_t stream = at::musa::getCurrentMUSAStream();
+
+#define DISPATCH(CPP_TYPE) \
+    vllm_musa::moe::MP22GroupedGatingTopkKernelLauncher(\
+    static_cast<CPP_TYPE*>(topk_weights.data_ptr()),\
+    topk_indices.data_ptr<int>(),\
+    static_cast<CPP_TYPE*>(gating_output.data_ptr()),\
+    static_cast<CPP_TYPE*>(e_score_correction_bias.value().data_ptr()),\
+    num_tokens,\
+    num_experts,\
+    num_expert_group,\
+    topk_group,\
+    topk,\
+    renormalize,\
+    stream);
+
+    if (gating_output.scalar_type() == at::ScalarType::Float) {
+        DISPATCH(float);
+    }
+    else if (gating_output.scalar_type() == at::ScalarType::Half) {
+        DISPATCH(__half);
+    } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
+        DISPATCH(__mt_bfloat16);
+    }
+    else {
+        TORCH_CHECK(false, "only support Float, Half and BFloat16 dtype");
+    }
+
+#undef DISPATCH
+}
diff --git a/vllm_musa/csrc/moe/mp31_fused_grouped_topk_softmax.mu b/vllm_musa/csrc/moe/mp31_fused_grouped_topk_softmax.mu
new file mode 100644
index 000000000..265d8482c
--- /dev/null
+++ b/vllm_musa/csrc/moe/mp31_fused_grouped_topk_softmax.mu
@@ -0,0 +1,1180 @@
+#include <torch/all.h>
+#include <ATen/Dispatch.h>
+#include <torch_musa/csrc/aten/musa/MUSAContext.h>
+#include <torch_musa/csrc/core/MUSAGuard.h>
+
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ <= 220
+#define MT_WARP_SIZE 128
+#else
+#define MT_WARP_SIZE 32
+#endif
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+namespace vllm_musa {
+namespace moe {
+
+template <typename T>
+struct AccumulateTypeTraits;
+
+template <>
+struct AccumulateTypeTraits<float> {
+    using type = float;
+};
+
+template <>
+struct AccumulateTypeTraits<__half> {
+    using type = float;
+};
+
+template <>
+struct AccumulateTypeTraits<__mt_bfloat16> {
+    // using type = __mt_bfloat16;
+    // high precison even produce unexcepted decoding results ?
+    using type = float;
+};
+
+
+/// Aligned array type
+template <
+    typename T,
+    /// Number of elements in the array
+    int N,
+    /// Alignment requirement in bytes
+    int Alignment = sizeof(T) * N
+>
+struct alignas(Alignment) AlignedArray {
+    __device__ AlignedArray() {}
+    __device__ AlignedArray(const AlignedArray& rhs) {
+        #pragma unroll
+        for (int i = 0; i < N; i++) {
+            this->data[i] = rhs.data[i];
+        }
+    }
+    __device__ AlignedArray(const T& value) {
+        #pragma unroll
+        for (int i = 0; i < N; i++) {
+            this->data[i] = value;
+        }
+    }
+
+    T data[N];
+};
+
+namespace {
+
+template <typename T>
+struct Add {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a + b;
+  }
+};
+
+template <typename T>
+struct Max {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a < b ? b : a;
+  }
+};
+
+template <typename T, int width>
+__device__ __forceinline__ T musa_shfl_xor_sync(T val, int lane_mask) {
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 220))
+  return __shfl_xor_sync(0xffffffff, val, lane_mask, width);
+#elif (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 210))
+  static_assert((width >= 2) && (width <= 128) && ((width & (width - 1)) == 0));
+
+  auto shfl_func = [&](int& var) {
+    if constexpr (width == 128) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x7f, ((~(128 - 1)) & 0x7f) << 7 | (128 - 1));
+    } else if constexpr (width == 64) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x3f, ((~(64 - 1)) & 0x7f) << 7 | (64 - 1));
+    } else if constexpr (width == 32) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x1f, ((~(32 - 1)) & 0x7f) << 7 | (32 - 1));
+    } else if constexpr (width == 16) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0xf, ((~(16 - 1)) & 0x7f) << 7 | (16 - 1));
+    } else if constexpr (width == 8) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x7, ((~(8 - 1)) & 0x7f) << 7 | (8 - 1));
+    } else if constexpr (width == 4) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x3, ((~(4 - 1)) & 0x7f) << 7 | (4 - 1));
+    } else if constexpr (width == 2) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x1, ((~(2 - 1)) & 0x7f) << 7 | (2 - 1));
+    }
+  };
+
+  if constexpr (sizeof(T) == 4) {
+    int var = *(reinterpret_cast<int32_t*>(&val));
+    int ret = shfl_func(var);
+    return *(reinterpret_cast<T*>(&ret));
+  } else {
+    struct __Bits {
+      int __a, __b;
+    };
+    __Bits __tmp;
+    memcpy(&__tmp, &val, sizeof(val));
+    __tmp.__a = shfl_func(__tmp.__a);
+    __tmp.__b = shfl_func(__tmp.__b);
+    int64_t ret = *(reinterpret_cast<int64_t*>(&__tmp));
+    return *(reinterpret_cast<T*>(&ret));
+  }
+#else
+  return 0;
+#endif
+}
+
+
+template <typename T, template <typename> class ReduceOp, int WIDTH=128>
+__device__ __forceinline__ T warpReduce(T val) {
+    ReduceOp<T> r;
+    #pragma unroll
+    for (int offset = WIDTH / 2; offset > 0; offset >>= 1) {
+        // TODO: parameterization width
+        T b = musa_shfl_xor_sync<T, WIDTH>(val, offset);
+        val = r(val, b);
+    }
+    return val;
+}
+
+
+__device__ __host__ __forceinline__ uint32_t next_power_of_two(uint32_t n) {
+    if (n == 0)
+        return 1;
+    n -= 1;
+    n |= n >> 1;
+    n |= n >> 2;
+    n |= n >> 4;
+    n |= n >> 8;
+    n |= n >> 16;
+    return -~n;
+}
+}  // anonymous namespace
+
+
+// ====================== Grouped TopK softmax things ===============================
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    typename std::enable_if<std::is_same<float, T>::value>::type* = nullptr
+>
+__global__ void grouped_topk_softmax(const T* input, T* softmax_output, int* group_indices, const int experts_per_group, const int topk_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+
+    // threadIdx.y: indexing EXPERT_GROUP
+
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+
+    // 1. Run softmax
+    // NUM_EXPERTS = NUM_EXPERT_GROUP * experts_per_group
+    const int global_offset = row_id * NUM_EXPERT_GROUP * experts_per_group +
+                              experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = input + global_offset;
+    T* thread_write_base_ptr = softmax_output + global_offset;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+
+    T row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccessType* row_chunk_vec_ptr = reinterpret_cast<AccessType*>(&row_chunk);
+    for (int ii = 0; ii < VPT; ii++) {
+        int offset = local_col_id + ii * ELTS_PER_SUB_WARP;
+        row_chunk_vec_ptr[ii] = offset < experts_per_group ?
+                                *reinterpret_cast<const AccessType*>(thread_read_base_ptr + offset) :
+                                (T)(-INFINITY);
+    }
+
+    // 1.1 calculate maxval
+    T maxval = row_chunk[0];  // global max value
+    for (int ii = 1; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval = fmaxf(maxval, row_chunk[ii]);
+    }
+
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask *= 2) {
+        T b = musa_shfl_xor_sync<T, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < b ? b : maxval;
+    }
+    // we got maxval in each group now
+    T maxval_per_group = maxval;
+    #pragma unroll
+    for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+        T b = musa_shfl_xor_sync<T, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < b ? b : maxval;
+    }
+
+    T sum = static_cast<T>(0.0f);
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = expf(row_chunk[ii] - maxval);
+        sum += row_chunk[ii];
+    }
+
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        T b = musa_shfl_xor_sync<T, MT_WARP_SIZE>(sum, mask);
+        sum += b;
+    }
+
+    const T reciprocal_sum = static_cast<T>(1.f) / sum;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_sum;
+    }
+
+    for (int ii = local_col_id; ii < experts_per_group; ii += ELTS_PER_SUB_WARP) {
+        *reinterpret_cast<AccessType*>(thread_write_base_ptr + ii) = row_chunk_vec_ptr[ii / ELTS_PER_SUB_WARP];
+    }
+
+    // Now, we get max value within each expert group (softmaxed value)
+    maxval_per_group = expf(maxval_per_group - maxval) * reciprocal_sum;
+
+
+    // 2. Find the topk_group
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        T max_val = maxval_per_group;
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+            T other_maxval = musa_shfl_xor_sync<T, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.x == 0 && threadIdx.y == 0) {
+            group_indices[row_id * topk_group + k_idx] = group_idx;
+        }
+
+        // Finally, we clear the value in the thread with the current max if there is another iteration to run.
+        if (k_idx + 1 < topk_group) {
+            if (threadIdx.y == group_idx) {
+                maxval_per_group = static_cast<T>(-10000.f);
+            }
+        }
+
+    }
+}
+
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    typename std::enable_if<std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void grouped_topk_softmax(const T* input, T* softmax_output, int* group_indices, const int experts_per_group, const int topk_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+    using AccumulateType = float;
+    using AccumulateAccessType = AlignedArray<AccumulateType, ELTS_PER_LDG>;
+
+    // threadIdx.y: indexing EXPERT_GROUP
+
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+
+    // 1. Run softmax
+    // NUM_EXPERTS = NUM_EXPERT_GROUP * experts_per_group
+    const int global_offset = row_id * NUM_EXPERT_GROUP * experts_per_group +
+                              experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = input + global_offset;
+    T* thread_write_base_ptr = softmax_output + global_offset;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+    AccumulateType row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccumulateAccessType* row_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&row_chunk);
+    for (int ii = local_col_id; ii < experts_per_group; ii += ELTS_PER_SUB_WARP) {
+        AccessType regs = *reinterpret_cast<const AccessType*>(thread_read_base_ptr + ii);
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii / ELTS_PER_SUB_WARP;
+        #pragma unroll
+        for (int k = 0; k < ELTS_PER_LDG; k++) {
+            row_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(regs.data[k]);
+        }
+    }
+    // 1.1 calculate maxval
+    AccumulateType maxval = row_chunk[0];  // global max value
+    for (int ii = 1; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval = fmaxf(maxval, row_chunk[ii]);
+    }
+
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask *= 2) {
+        AccumulateType b = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < b ? b : maxval;
+    }
+    AccumulateType maxval_per_group = maxval;
+    #pragma unroll
+    for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+        AccumulateType b = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < b ? b : maxval;
+    }
+
+    AccumulateType sum = static_cast<AccumulateType>(0.0f);
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        // row_chunk[ii] = expf(row_chunk[ii] - maxval);
+        row_chunk[ii] = threadIdx.x * ELTS_PER_LDG < experts_per_group ?
+                        expf(row_chunk[ii] - maxval) :
+                        (AccumulateType)(0.0f);
+        sum += row_chunk[ii];
+    }
+
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType b = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(sum, mask);
+        sum += b;
+    }
+
+    const AccumulateType reciprocal_sum = static_cast<AccumulateType>(1.f) / sum;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_sum;
+    }
+
+    for (int ii = local_col_id; ii < experts_per_group; ii += ELTS_PER_SUB_WARP) {
+        AccessType regs;
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii / ELTS_PER_SUB_WARP;
+        for (int k = 0; k < ELTS_PER_LDG; k++) {
+            regs.data[k] = static_cast<T>(row_chunk_cur_vec_ptr->data[k]);
+        }
+        *reinterpret_cast<AccessType*>(thread_write_base_ptr + ii) = regs;
+    }
+
+    // Now, we get max value within each expert group
+    maxval_per_group = expf(maxval_per_group - maxval) * reciprocal_sum;
+
+    // 2. Find the topk_group
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        // if (threadIdx.x == 0 && blockIdx.x == 0 && k_idx == 0) {
+        //     printf("threadIdx.y: %d, maxval: %f\n", threadIdx.y, maxval);
+        // }
+        // low precison dtype will be fine
+        T max_val = static_cast<T>(maxval_per_group);
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+            T other_maxval = musa_shfl_xor_sync<T, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.x == 0 && threadIdx.y == 0) {
+            group_indices[row_id * topk_group + k_idx] = group_idx;
+        }
+
+        // Finally, we clear the value in the thread with the current max if there is another iteration to run.
+        if (k_idx + 1 < topk_group) {
+            if (threadIdx.y == group_idx) {
+                maxval_per_group = static_cast<T>(-10000.f);
+            }
+        }
+    }
+}
+
+// ==================== Grouped Topk Things ====================
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    bool RENORMALIZE,
+    typename std::enable_if<std::is_same<float, T>::value || std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void grouped_gating_topk(const T* gating_output, T* topk_weights, int* topk_indices, const int topk_group, const int topk, const int experts_per_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+    using AccumulateType = typename AccumulateTypeTraits<T>::type;
+    using AccumulateAccessType = AlignedArray<AccumulateType, ELTS_PER_LDG>;
+
+    // we split a warp into multiple sub-warps, each sub-warp is responsible
+    // for calculating the maximum value within an expert-group.
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+    // 1. Run softmax along expert dimension
+    const int base_offset = row_id * NUM_EXPERT_GROUP * experts_per_group + experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = gating_output + base_offset;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+    // all elements will be loaded into thread registers, this is acceptable
+    // because the kernel assumes a small number of experts (maybe <= 256)
+    AccumulateType row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccumulateAccessType* row_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&row_chunk);
+    for (int ii = 0; ii < VPT; ii++) {
+        int offset = local_col_id + ii * ELTS_PER_SUB_WARP;
+        const AccessType regs = offset < experts_per_group ?
+                                *reinterpret_cast<const AccessType*>(thread_read_base_ptr + offset) :
+                                (T)(-INFINITY);
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii;
+        if constexpr (std::is_same<AccumulateType, T>::value) {
+            *row_chunk_cur_vec_ptr = regs;
+        } else {
+            #pragma unroll
+            for (int k = 0; k < ELTS_PER_LDG; k++) {
+                row_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(regs.data[k]);
+            }
+        }
+    }
+    // 1.1 calculate max value
+    AccumulateType maxval = row_chunk[0];
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval = fmaxf(maxval, row_chunk[ii]);
+    }
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask *= 2) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < other_maxval ? other_maxval : maxval;
+    }
+    AccumulateType maxval_per_group = maxval;
+    #pragma unroll
+    for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < other_maxval ? other_maxval : maxval;
+    }
+
+    // 1.2 calculate sum value
+    AccumulateType sum = static_cast<AccumulateType>(0.f);
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        // TODO: zero sum value on invalid lanes
+        row_chunk[ii] = expf(row_chunk[ii] - maxval);
+        sum += row_chunk[ii];
+    }
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(sum, mask);
+        sum += other_sum;
+    }
+
+    // 1.3 calcluate softmax final result
+    const AccumulateType reciprocal_sum = static_cast<AccumulateType>(1.f) / sum;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_sum;
+    }
+
+    // 2. Find the topk group
+    AccumulateType mask_per_group = static_cast<AccumulateType>(0.f);
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        AccumulateType max_val = maxval_per_group;
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask *= 2) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.y == group_idx) {
+            mask_per_group = static_cast<AccumulateType>(1.f);
+            maxval_per_group = static_cast<AccumulateType>(-10000.f);
+        }
+    }
+
+    // 3. Find the topk
+    __shared__ AccumulateType topk_weights_shm[8];
+    for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+        int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+        for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+            row_chunk[vpt_id * ELTS_PER_LDG + ii] = (local_base_offset + ii) < experts_per_group ?
+                                                    row_chunk[vpt_id * ELTS_PER_LDG + ii] * mask_per_group :
+                                                    static_cast<AccumulateType>(-10.0f);
+        }
+    }
+    for (int k_idx = 0; k_idx < topk; k_idx++) {
+        AccumulateType max_val = static_cast<AccumulateType>(-1.0f);
+        int expert = 10000;
+
+        // not all threads' row_chunk value are valid
+        for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+            int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+            for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+                if (row_chunk[vpt_id * ELTS_PER_LDG + ii] > max_val) {
+                    max_val = row_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    expert = threadIdx.y * experts_per_group + local_base_offset + ii;
+                }
+            }
+        }
+
+        #pragma unroll
+        for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            int other_expert = musa_shfl_xor_sync<int, MT_WARP_SIZE>(expert, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_expert < expert)) {
+                max_val = other_maxval;
+                expert = other_expert;
+            }
+        }
+
+        if (threadIdx.x == 0 && threadIdx.y == 0) {
+            // write something to global memory
+            const int idx = topk * row_id + k_idx;
+            topk_indices[idx] = expert;
+            if constexpr (RENORMALIZE) {
+                topk_weights_shm[k_idx] = static_cast<T>(max_val);
+            } else {
+                topk_weights[idx] = static_cast<T>(max_val);
+            }
+        }
+
+        if (k_idx + 1 < topk) {
+            const int local_expert_idx = expert % experts_per_group;
+            const int thread_to_clear = (local_expert_idx / ELTS_PER_LDG) % THREADS_PER_EXPERT_GROUP;
+            if (threadIdx.y == (expert / experts_per_group) && threadIdx.x == thread_to_clear) {
+                int row_chunk_idx_to_clear = (local_expert_idx / ELTS_PER_SUB_WARP) * ELTS_PER_LDG + local_expert_idx % ELTS_PER_LDG;
+                row_chunk[row_chunk_idx_to_clear] = static_cast<AccumulateType>(-1.0f);
+            }
+        }
+    }
+
+    // 4. renormalization
+    if constexpr (RENORMALIZE) {
+        int lane_id = (threadIdx.x + threadIdx.y * blockDim.x) % MT_WARP_SIZE;
+        AccumulateType topk_weight_reg = lane_id < 8 ? topk_weights_shm[lane_id] : static_cast<AccumulateType>(0.0f);
+        __syncthreads();
+        AccumulateType topk_sum = topk_weight_reg;
+        for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+            AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(topk_sum, mask);
+            topk_sum += other_sum;
+        }
+
+        if (lane_id < 8) {
+            topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg / topk_sum);
+        }
+    }
+}
+
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    bool RENORMALIZE,
+    typename std::enable_if<std::is_same<float, T>::value || std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void grouped_gating_topk_kernel(const T* gating_output, const T* e_score_correction_bias, T* topk_weights, int* topk_indices, const int topk_group, const int topk, const int experts_per_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+    using AccumulateType = typename AccumulateTypeTraits<T>::type;
+    using AccumulateAccessType = AlignedArray<AccumulateType, ELTS_PER_LDG>;
+
+    // USE_BIASED_SCORE: if true, use biased scores for expert selection but
+    // original scores for routing weights.
+
+
+    // we split a warp into multiple sub-warps, each sub-warp is responsible
+    // for calculating the maximum value within an expert-group.
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+    // 1. Run softmax along expert dimension
+    const int base_offset = row_id * NUM_EXPERT_GROUP * experts_per_group + experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = gating_output + base_offset;
+    const T* thread_read_score_bias_base_ptr = e_score_correction_bias + experts_per_group * threadIdx.y;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+    // all elements will be loaded into thread registers, this is acceptable
+    // because the kernel assumes a small number of experts (maybe <= 256)
+    AccumulateType row_chunk[VPT * ELTS_PER_LDG] = {(T)(-INFINITY)};
+    AccumulateType score_correction_bias_chunk[VPT * ELTS_PER_LDG] = {(T)(0.0f)};
+    AccumulateAccessType* row_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&row_chunk);
+    AccumulateAccessType* score_correction_bias_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&score_correction_bias_chunk);
+    for (int ii = 0; ii < VPT; ii++) {
+        int offset = local_col_id + ii * ELTS_PER_SUB_WARP;
+        bool valid_pos = offset < experts_per_group;
+        const AccessType regs = valid_pos ?
+                                *reinterpret_cast<const AccessType*>(thread_read_base_ptr + offset) :
+                                (T)(-INFINITY);
+        const AccessType score_bias_regs = valid_pos ?
+                                           *reinterpret_cast<const AccessType*>(thread_read_score_bias_base_ptr + offset) :
+                                           (T)(0.0f);
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii;
+        AccumulateAccessType* score_correction_bias_chunk_cur_vec_ptr = score_correction_bias_chunk_vec_ptr + ii;
+        if constexpr (std::is_same<AccumulateType, T>::value) {
+            *row_chunk_cur_vec_ptr = regs;
+            *score_correction_bias_chunk_cur_vec_ptr = score_bias_regs;
+        } else {
+            #pragma unroll
+            for (int k = 0; k < ELTS_PER_LDG; k++) {
+                row_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(regs.data[k]);
+                score_correction_bias_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(score_bias_regs.data[k]);
+            }
+        }
+    }
+    // 1.1 calculate max value
+    AccumulateType maxval = row_chunk[0];
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval = fmaxf(maxval, row_chunk[ii]);
+    }
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval, mask);
+        maxval = maxval < other_maxval ? other_maxval : maxval;
+    }
+
+    // 1.2 calculate sum value
+    AccumulateType sum = static_cast<AccumulateType>(0.f);
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        // TODO: zero sum value on invalid lanes
+        row_chunk[ii] = expf(row_chunk[ii] - maxval);
+        sum += row_chunk[ii];
+    }
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(sum, mask);
+        sum += other_sum;
+    }
+
+    // 1.3 calcluate softmax final result and apply correction bias
+    const AccumulateType reciprocal_sum = static_cast<AccumulateType>(1.f) / sum;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_sum;
+        score_correction_bias_chunk[ii] += row_chunk[ii];
+    }
+
+    // 1.4 find max value in each expert group
+    AccumulateType maxval_per_group = score_correction_bias_chunk[0];
+    for (int ii = 1; ii < VPT * ELTS_PER_LDG; ii++) {
+        maxval_per_group = fmaxf(maxval_per_group, score_correction_bias_chunk[ii]);
+    }
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask <<= 1) {
+        AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval_per_group, mask);
+        maxval_per_group = maxval_per_group < other_maxval ? other_maxval : maxval_per_group;
+    }
+
+    // 2. Find the topk group
+    AccumulateType mask_per_group = static_cast<AccumulateType>(0.f);
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        AccumulateType max_val = maxval_per_group;
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask <<= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.y == group_idx) {
+            mask_per_group = static_cast<AccumulateType>(1.f);
+            maxval_per_group = static_cast<AccumulateType>(-10000.f);
+        }
+    }
+
+    // 3. Find the topk
+    AccumulateType topk_weight_reg = (AccumulateType)(0.0f);
+    int topk_id_reg;
+    int lane_id = (threadIdx.x + threadIdx.y * blockDim.x) % MT_WARP_SIZE;
+    for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+        int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+        for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+            score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] = (local_base_offset + ii) < experts_per_group ?
+                                                    score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] * mask_per_group :
+                                                    static_cast<AccumulateType>(-INFINITY);
+        }
+    }
+    for (int k_idx = 0; k_idx < topk; k_idx++) {
+        AccumulateType max_val = static_cast<AccumulateType>(-INFINITY);
+        AccumulateType original_max_val = static_cast<AccumulateType>(-INFINITY);
+        int expert = 10000;
+
+        // not all threads' row_chunk value are valid
+        for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+            int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+            for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+                if (score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] > max_val) {
+                    max_val = score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    original_max_val = row_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    expert = threadIdx.y * experts_per_group + local_base_offset + ii;
+                }
+            }
+        }
+
+        #pragma unroll
+        for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            AccumulateType other_original_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(original_max_val, mask);
+            int other_expert = musa_shfl_xor_sync<int, MT_WARP_SIZE>(expert, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_expert < expert)) {
+                max_val = other_maxval;
+                original_max_val = other_original_maxval;
+                expert = other_expert;
+            }
+        }
+
+        // now we find maximum value and expert index
+        const int local_expert_idx = expert % experts_per_group;
+        const int thread_to_clear = (local_expert_idx / ELTS_PER_LDG) % THREADS_PER_EXPERT_GROUP;
+        const int chunk_idx_to_clear = (local_expert_idx / ELTS_PER_SUB_WARP) * ELTS_PER_LDG + local_expert_idx % ELTS_PER_LDG;
+        if (k_idx + 1 < topk) {
+            if (threadIdx.y == (expert / experts_per_group) && threadIdx.x == thread_to_clear) {
+                score_correction_bias_chunk[chunk_idx_to_clear] = static_cast<AccumulateType>(-INFINITY);
+            }
+        }
+
+        if (lane_id == k_idx) {
+            topk_weight_reg = static_cast<AccumulateType>(original_max_val);
+            topk_id_reg = expert;
+        }
+    }
+
+    // 4. renormalization
+    if constexpr (RENORMALIZE) {
+        AccumulateType topk_sum = topk_weight_reg;
+        for (int mask = 1; mask < next_power_of_two(topk); mask <<= 1) {
+            AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(topk_sum, mask);
+            topk_sum += other_sum;
+        }
+
+        if (lane_id < topk) {
+            topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg / topk_sum);
+            topk_indices[topk * row_id + lane_id] = topk_id_reg;
+        }
+    } else {
+        if (lane_id < topk) {
+            topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg);
+            topk_indices[topk * row_id + lane_id] = topk_id_reg;
+        }
+    }
+}
+
+
+template <
+    typename T,
+    int VPT,
+    int NUM_EXPERT_GROUP,
+    int THREADS_PER_EXPERT_GROUP,
+    int ELTS_PER_LDG,
+    bool RENORMALIZE,
+    typename std::enable_if<std::is_same<float, T>::value || std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void mp31_grouped_gating_topk_sigmoid_kernel(const T* gating_output, const T* e_score_correction_bias, float* topk_weights, int* topk_indices, const int topk_group, const int topk, const int experts_per_group) {
+    using AccessType = AlignedArray<T, ELTS_PER_LDG>;
+    using AccumulateType = typename AccumulateTypeTraits<T>::type;
+    using AccumulateAccessType = AlignedArray<AccumulateType, ELTS_PER_LDG>;
+
+    // we split a warp into multiple sub-warps, each sub-warp is responsible
+    // for calculating the maximum value within an expert-group.
+    constexpr uint32_t ELTS_PER_SUB_WARP = (MT_WARP_SIZE / NUM_EXPERT_GROUP) * ELTS_PER_LDG;
+
+    const uint32_t row_id = blockIdx.x;
+
+    // 1. Run sigmoid along expert dimension
+    const int base_offset = row_id * NUM_EXPERT_GROUP * experts_per_group + experts_per_group * threadIdx.y;
+    const T* thread_read_base_ptr = gating_output + base_offset;
+    const T* thread_read_score_bias_base_ptr = e_score_correction_bias + experts_per_group * threadIdx.y;
+
+    // !!! 2-dim threadblock
+    uint32_t local_col_id = threadIdx.x * ELTS_PER_LDG;
+
+    // all elements will be loaded into thread registers, this is acceptable
+    // because the kernel assumes a small number of experts (maybe <= 256)
+    AccumulateType row_chunk[VPT * ELTS_PER_LDG] = {(AccumulateType)(-INFINITY)};
+    AccumulateType score_correction_bias_chunk[VPT * ELTS_PER_LDG] = {(AccumulateType)(0.0f)};
+    AccumulateAccessType* row_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&row_chunk);
+    AccumulateAccessType* score_correction_bias_chunk_vec_ptr = reinterpret_cast<AccumulateAccessType*>(&score_correction_bias_chunk);
+    for (int ii = 0; ii < VPT; ii++) {
+        int offset = local_col_id + ii * ELTS_PER_SUB_WARP;
+        bool valid_pos = offset < experts_per_group;
+        const AccessType regs = valid_pos ?
+                                *reinterpret_cast<const AccessType*>(thread_read_base_ptr + offset) :
+                                (T)(-INFINITY);
+        const AccessType score_bias_regs = valid_pos ?
+                                           *reinterpret_cast<const AccessType*>(thread_read_score_bias_base_ptr + offset) :
+                                           (T)(0.0f);
+        AccumulateAccessType* row_chunk_cur_vec_ptr = row_chunk_vec_ptr + ii;
+        AccumulateAccessType* score_correction_bias_chunk_cur_vec_ptr = score_correction_bias_chunk_vec_ptr + ii;
+        if constexpr (std::is_same<AccumulateType, T>::value) {
+            *row_chunk_cur_vec_ptr = regs;
+            *score_correction_bias_chunk_cur_vec_ptr = score_bias_regs;
+        } else {
+            #pragma unroll
+            for (int k = 0; k < ELTS_PER_LDG; k++) {
+                row_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(regs.data[k]);
+                score_correction_bias_chunk_cur_vec_ptr->data[k] = static_cast<AccumulateType>(score_bias_regs.data[k]);
+            }
+        }
+    }
+
+    // 1.1 run sigmoid and apply bias correction
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        row_chunk[ii] = 1.f / (1.f + __expf(-row_chunk[ii]));
+        score_correction_bias_chunk[ii] += row_chunk[ii];
+    }
+
+    // 1.2 find max value in each expert group
+    // naive sort should be fine on small numel
+    AccumulateType maxval_per_group = static_cast<AccumulateType>(-INFINITY);
+    AccumulateType submaxval_per_group = static_cast<AccumulateType>(-INFINITY);
+    int max_idx = -1;
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        if (score_correction_bias_chunk[ii] > maxval_per_group) {
+            maxval_per_group = score_correction_bias_chunk[ii];
+            max_idx = ii;
+        }
+    }
+    for (int ii = 0; ii < VPT * ELTS_PER_LDG; ii++) {
+        if (ii != max_idx && (score_correction_bias_chunk[ii] > submaxval_per_group)) {
+            submaxval_per_group = score_correction_bias_chunk[ii];
+        }
+    }
+
+    AccumulateType other_maxvals[2];
+    #pragma unroll
+    for (int mask = 1; mask < THREADS_PER_EXPERT_GROUP; mask <<= 1) {
+        other_maxvals[0] = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(maxval_per_group, mask);
+        other_maxvals[1] = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(submaxval_per_group, mask);
+        for (int i = 0; i < 2; i++) {
+            if (other_maxvals[i] > maxval_per_group) {
+                submaxval_per_group = maxval_per_group;
+                maxval_per_group = other_maxvals[i];
+            } else if (other_maxvals[i] > submaxval_per_group) {
+                submaxval_per_group = other_maxvals[i];
+            }
+        }
+    }
+
+    maxval_per_group += submaxval_per_group;
+
+
+    // 2. Find the topk group
+    // use mask_per_group to mark the un-selected values as 0.0f, which keeps
+    // consistence with this implementation:
+    // tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), 0.0)
+    AccumulateType mask_per_group = static_cast<AccumulateType>(0.f);
+    for (int k_idx = 0; k_idx < topk_group; k_idx++) {
+        AccumulateType max_val = maxval_per_group;
+        int group_idx = threadIdx.y;
+
+        #pragma unroll
+        for (int mask = THREADS_PER_EXPERT_GROUP; mask < MT_WARP_SIZE; mask <<= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            int other_group_idx = musa_shfl_xor_sync<int, MT_WARP_SIZE>(group_idx, mask);
+
+            if (other_maxval > max_val || (other_maxval == max_val && other_group_idx < group_idx)) {
+                max_val = other_maxval;
+                group_idx = other_group_idx;
+            }
+        }
+
+        if (threadIdx.y == group_idx) {
+            mask_per_group = static_cast<AccumulateType>(1.f);
+            maxval_per_group = static_cast<AccumulateType>(-10000.f);
+        }
+    }
+
+    // 3. Find the topk
+    AccumulateType topk_weight_reg = (AccumulateType)(0.0f);
+    int topk_id_reg;
+    int lane_id = (threadIdx.x + threadIdx.y * blockDim.x) % MT_WARP_SIZE;
+    for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+        int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+        for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+            score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] = (local_base_offset + ii) < experts_per_group ?
+                                                    score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] * mask_per_group :
+                                                    static_cast<AccumulateType>(-INFINITY);
+        }
+    }
+    for (int k_idx = 0; k_idx < topk; k_idx++) {
+        AccumulateType max_val = static_cast<AccumulateType>(-INFINITY);
+        AccumulateType original_max_val = static_cast<AccumulateType>(-INFINITY);
+        int expert = threadIdx.y * experts_per_group + local_col_id;  // expert id will be updated later
+
+        // not all threads' row_chunk value are valid
+        for (int vpt_id = 0; vpt_id < VPT; vpt_id++) {
+            int local_base_offset = local_col_id + vpt_id * ELTS_PER_SUB_WARP;
+            for (int ii = 0; ii < ELTS_PER_LDG; ii++) {
+                if (score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii] > max_val) {
+                    max_val = score_correction_bias_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    original_max_val = row_chunk[vpt_id * ELTS_PER_LDG + ii];
+                    expert = threadIdx.y * experts_per_group + local_base_offset + ii;
+                }
+            }
+        }
+
+        #pragma unroll
+        for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+            AccumulateType other_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(max_val, mask);
+            AccumulateType other_original_maxval = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(original_max_val, mask);
+            int other_expert = musa_shfl_xor_sync<int, MT_WARP_SIZE>(expert, mask);
+
+            if (other_maxval > max_val || ((!isinf(other_maxval) && !isinf(max_val)) && other_maxval == max_val && other_expert < expert)) {
+                max_val = other_maxval;
+                original_max_val = other_original_maxval;
+                expert = other_expert;
+            }
+        }
+
+        // now we find maximum value and expert index
+        const int local_expert_idx = expert % experts_per_group;
+        const int thread_to_clear = (local_expert_idx / ELTS_PER_LDG) % THREADS_PER_EXPERT_GROUP;
+        const int chunk_idx_to_clear = (local_expert_idx / ELTS_PER_SUB_WARP) * ELTS_PER_LDG + local_expert_idx % ELTS_PER_LDG;
+        if (k_idx + 1 < topk) {
+            if (threadIdx.y == (expert / experts_per_group) && threadIdx.x == thread_to_clear) {
+                score_correction_bias_chunk[chunk_idx_to_clear] = static_cast<AccumulateType>(-INFINITY);
+            }
+        }
+
+        if (lane_id == k_idx) {
+            topk_weight_reg = static_cast<AccumulateType>(original_max_val);
+            topk_id_reg = expert;
+        }
+    }
+
+    // 4. renormalization
+    if constexpr (RENORMALIZE) {
+        AccumulateType topk_sum = topk_weight_reg;
+        for (int mask = 1; mask < next_power_of_two(topk); mask <<= 1) {
+            AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(topk_sum, mask);
+            topk_sum += other_sum;
+        }
+
+        if (lane_id < topk) {
+            // topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg / topk_sum);
+            topk_weights[topk * row_id + lane_id] = static_cast<float>(topk_weight_reg / topk_sum);
+            topk_indices[topk * row_id + lane_id] = topk_id_reg;
+        }
+    } else {
+        if (lane_id < topk) {
+            // topk_weights[topk * row_id + lane_id] = static_cast<T>(topk_weight_reg);
+            topk_weights[topk * row_id + lane_id] = static_cast<float>(topk_weight_reg);
+            topk_indices[topk * row_id + lane_id] = topk_id_reg;
+        }
+    }
+}
+
+
+// topk_weights won't be large in decode stage, element wise kernel will be fine
+template <
+    typename T,
+    typename std::enable_if<std::is_same<float, T>::value || std::is_same<__half, T>::value || std::is_same<__mt_bfloat16, T>::value>::type* = nullptr
+>
+__global__ void renormalize_topk_weights_kernel(T* topk_weights, int num_tokens, int num_cols) {
+    using AccumulateType = typename AccumulateTypeTraits<T>::type;
+
+    const uint32_t bid = blockIdx.x * blockDim.y + threadIdx.y;
+    if (bid >= num_tokens) {
+        return;
+    }
+    const T* thread_read_base_ptr = topk_weights + bid * num_cols;
+    T* thread_write_base_ptr = topk_weights + bid * num_cols;
+
+    AccumulateType sum = static_cast<AccumulateType>(0.f);
+    for (int tid = threadIdx.x; tid < num_cols; tid += blockDim.x) {
+        sum += static_cast<AccumulateType>(thread_read_base_ptr[tid]);
+    }
+
+    #pragma unroll
+    for (int mask = MT_WARP_SIZE / 2; mask > 0; mask >>= 1) {
+        AccumulateType other_sum = musa_shfl_xor_sync<AccumulateType, MT_WARP_SIZE>(sum, mask);
+        sum += other_sum;
+    }
+
+    for (int tid = threadIdx.x; tid < num_cols; tid += blockDim.x) {
+        thread_write_base_ptr[tid] = static_cast<T>(static_cast<AccumulateType>(thread_write_base_ptr[tid]) / sum);
+    }
+}
+
+
+#define LAUNCH_TOPK_GATING_SOFTMAX(TYPE, NUM_EXPERTS, NUM_EXPERT_GROUP)           \
+{\
+    constexpr std::size_t MAX_BYTES_PER_LDG = sizeof(TYPE);                    \
+    constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(TYPE) * (NUM_EXPERTS / NUM_EXPERT_GROUP));\
+    constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(TYPE);\
+    constexpr int VPT = MAX(1, NUM_EXPERTS / ELTS_PER_LDG / MT_WARP_SIZE);\
+    dim3 block(MT_WARP_SIZE / NUM_EXPERT_GROUP, NUM_EXPERT_GROUP);                                    \
+    dim3 grid(num_tokens);                 \
+    moe::grouped_topk_softmax<TYPE, VPT, NUM_EXPERT_GROUP, MT_WARP_SIZE / NUM_EXPERT_GROUP, ELTS_PER_LDG><<<grid, block, 0, stream>>>(    \
+        gating_output, softmax_res, group_indices,                     \
+        NUM_EXPERTS / NUM_EXPERT_GROUP, topk_group);\
+}
+
+
+#define LAUNCH_GROUPED_GATING_TOPK(TYPE, NUM_EXPERTS, NUM_EXPERT_GROUP)           \
+{\
+    constexpr std::size_t MAX_BYTES_PER_LDG = sizeof(TYPE);                    \
+    constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(TYPE) * (NUM_EXPERTS / NUM_EXPERT_GROUP));\
+    constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(TYPE);\
+    constexpr int VPT = MAX(1, NUM_EXPERTS / ELTS_PER_LDG / MT_WARP_SIZE);\
+    dim3 block(MT_WARP_SIZE / NUM_EXPERT_GROUP, NUM_EXPERT_GROUP);                                    \
+    dim3 grid(num_tokens);                 \
+    if (renormalize) {\
+        int shm_size = topk * sizeof(typename AccumulateTypeTraits<TYPE>::type);\
+        moe::mp31_grouped_gating_topk_sigmoid_kernel<TYPE, VPT, NUM_EXPERT_GROUP, MT_WARP_SIZE / NUM_EXPERT_GROUP, ELTS_PER_LDG, true><<<grid, block, 0, stream>>>(    \
+        gating_output, e_score_correction_bias, topk_weights, topk_indices,                   \
+        topk_group, topk, NUM_EXPERTS / NUM_EXPERT_GROUP);\
+    } else {\
+        moe::mp31_grouped_gating_topk_sigmoid_kernel<TYPE, VPT, NUM_EXPERT_GROUP, MT_WARP_SIZE / NUM_EXPERT_GROUP, ELTS_PER_LDG, false><<<grid, block, 0, stream>>>(    \
+        gating_output, e_score_correction_bias, topk_weights, topk_indices,                   \
+        topk_group, topk, NUM_EXPERTS / NUM_EXPERT_GROUP);\
+    }\
+}
+
+
+template <typename T>
+void topkGatingSoftmaxKernelLauncher(
+    const T* gating_output,
+    T* softmax_res,
+    int* group_indices,
+    const int num_tokens,
+    const int num_experts,
+    const int num_expert_group,
+    const int topk_group,
+    musaStream_t stream) {
+    static constexpr int WARPS_PER_TB = 1;
+
+    if (num_experts == 256) {
+        switch (num_expert_group) {
+            case 2:
+                LAUNCH_TOPK_GATING_SOFTMAX(T, 256, 2);
+                break;
+            case 4:
+                LAUNCH_TOPK_GATING_SOFTMAX(T, 256, 4);
+                break;
+            case 8:
+                LAUNCH_TOPK_GATING_SOFTMAX(T, 256, 8);
+                break;
+            case 16:
+                LAUNCH_TOPK_GATING_SOFTMAX(T, 256, 16);
+                break;
+            default: {
+                TORCH_CHECK(false, "Not implemented error");
+            }
+        }
+    } else {
+        TORCH_CHECK(false, "Not implemented error");
+    }
+}
+
+
+template <typename T>
+void MP31GroupedGatingTopkKernelLauncher(
+    float* topk_weights,
+    int* topk_indices,
+    const T* gating_output,
+    const T* e_score_correction_bias,
+    const int num_tokens,
+    const int num_experts,
+    const int num_expert_group,
+    const int topk_group,
+    const int topk,
+    const bool renormalize,
+    musaStream_t stream) {
+    static constexpr int WARPS_PER_TB = 1;
+
+    if (num_experts == 256) {
+        switch (num_expert_group) {
+            case 2:
+                LAUNCH_GROUPED_GATING_TOPK(T, 256, 2);
+                break;
+            case 4:
+                LAUNCH_GROUPED_GATING_TOPK(T, 256, 4);
+                break;
+            case 8:
+                LAUNCH_GROUPED_GATING_TOPK(T, 256, 8);
+                break;
+            case 16:
+                LAUNCH_GROUPED_GATING_TOPK(T, 256, 16);
+                break;
+            default: {
+                TORCH_CHECK(false, "Not implemented error");
+            }
+        }
+    } else {
+        TORCH_CHECK(false, "Not implemented error");
+    }
+}
+
+} // namespace moe
+} // namespace vllm_musa
+
+
+std::tuple<torch::Tensor, torch::Tensor> mp31_fused_grouped_topk(
+    const torch::Tensor& gating_output,
+    const c10::optional<torch::Tensor>& e_score_correction_bias,
+    int num_expert_group,
+    int topk_group,
+    int topk,
+    bool renormalize) {
+
+    TORCH_CHECK(e_score_correction_bias.has_value());
+
+    const int num_experts = gating_output.size(-1);
+    const int num_tokens = gating_output.numel() / num_experts;
+
+    auto topk_weights_option = gating_output.options().memory_format(at::MemoryFormat::Contiguous).dtype(at::kFloat);
+    auto topk_indices_option = gating_output.options().memory_format(at::MemoryFormat::Contiguous).dtype(at::kInt);
+
+    auto topk_weights = at::empty({num_tokens, topk}, topk_weights_option);
+    auto topk_indices = at::empty({num_tokens, topk}, topk_indices_option);
+
+    const at::musa::OptionalMUSAGuard device_guard(device_of(gating_output));
+    const musaStream_t stream = at::musa::getCurrentMUSAStream();
+
+#define DISPATCH(CPP_TYPE) \
+    vllm_musa::moe::MP31GroupedGatingTopkKernelLauncher(\
+    topk_weights.data_ptr<float>(),\
+    topk_indices.data_ptr<int>(),\
+    static_cast<CPP_TYPE*>(gating_output.data_ptr()),\
+    static_cast<CPP_TYPE*>(e_score_correction_bias.value().data_ptr()),\
+    num_tokens,\
+    num_experts,\
+    num_expert_group,\
+    topk_group,\
+    topk,\
+    renormalize,\
+    stream);
+
+    if (gating_output.scalar_type() == at::ScalarType::Float) {
+        DISPATCH(float);
+    }
+    else if (gating_output.scalar_type() == at::ScalarType::Half) {
+        DISPATCH(__half);
+    } else if (gating_output.scalar_type() == at::ScalarType::BFloat16) {
+        DISPATCH(__mt_bfloat16);
+    }
+    else {
+        TORCH_CHECK(false, "only support Float, Half and BFloat16 dtype");
+    }
+
+    return std::tuple<at::Tensor, at::Tensor>(topk_weights, topk_indices);
+
+#undef DISPATCH
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/moe/mutlass/fused_moe_impl/mp22_group_kernels.hpp b/vllm_musa/csrc/moe/mutlass/fused_moe_impl/mp22_group_kernels.hpp
new file mode 100644
index 000000000..2f842b3da
--- /dev/null
+++ b/vllm_musa/csrc/moe/mutlass/fused_moe_impl/mp22_group_kernels.hpp
@@ -0,0 +1,491 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (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 <mutlass/mutlass.h>
+#include <mutlass/fast_math.h>
+#include <mute/tensor.hpp>
+#include <mutlass/util/packed_stride.hpp>
+#include <mute/tensor.hpp>
+#include <mute/atom/copy_atom.hpp>
+#include <mutlass/gemm/collective/collective_builder.hpp>
+#include <vector>
+
+template <class ElementDst, bool IsPerGroupScale, class TensorScale, class Tensor0, class Tensor1>
+MUTE_DEVICE
+void
+dequant_w4(TensorScale scale_tensor, Tensor0& w4_tensor, Tensor1& f16_tensor) {
+  using ElementB = mute::conditional_t<IsPerGroupScale, uint8_t, int8_t>;
+  using ElementLH = mute::conditional_t<IsPerGroupScale, float, int8_t>;
+  MUTE_STATIC_ASSERT_V(size(w4_tensor) == size(f16_tensor));
+  // using v2f16_t = _Float16 __attribute__((vector_size(4)));
+  // using v2f16_t = mutlass::Array<mutlass::bfloat16_t, 2>;
+  constexpr bool is_half = is_same_v<ElementDst, mutlass::half_t>;
+  using v2f16_t = mute::conditional_t<is_half,
+                                      _Float16 __attribute__((vector_size(4))),
+                                      mutlass::Array<mutlass::bfloat16_t, 2>>;
+  using v2f32_t = float __attribute__((vector_size(8)));
+
+  Tensor s8_view = recast<ElementB>(w4_tensor);
+  Tensor v2f16_view = recast<v2f16_t>(f16_tensor);
+  // if (thread(0)) {
+  //   print("f16_tensor = ");print(f16_tensor);print("\n");
+  //   print("v2f16_view = ");print(v2f16_view);print("\n");
+  //   print("scale_tensor = ");print(scale_tensor);print("\n");
+  // }
+  asm("CTRL.NOP");
+  MUTE_UNROLL
+  for (int mi = 0; mi < size<1>(s8_view); ++mi) {
+    MUTE_UNROLL
+    for (int ni = 0; ni < size<2>(s8_view); ++ni) {
+      // alignment data
+      MUTE_UNROLL
+      for (int vi = 0; vi < size<0>(s8_view); ++vi) {
+        ElementLH lo;
+        ElementLH hi;
+        if constexpr (IsPerGroupScale) {
+          lo = (float)((s8_view(vi, mi, ni)) & 0xF) - 8.f;
+          hi = (float)((s8_view(vi, mi, ni) >> 4)) - 8.f;
+        } else {
+          lo = (s8_view(vi, mi, ni) << 4) & 0xF0;
+          hi = s8_view(vi, mi, ni) & 0xF0;
+        }
+        if constexpr (is_half) {
+          if constexpr (IsPerGroupScale) {
+            v2f32_t value{float(lo) * scale_tensor(vi, mi, ni), float(hi) * scale_tensor(vi, mi, ni)};
+            v2f16_view(vi, mi, ni) = __builtin_convertvector(value, v2f16_t);
+          } else {
+            v2f32_t value{float(lo), float(hi)};
+            v2f16_view(vi, mi, ni) = __builtin_convertvector(value, v2f16_t);
+          }
+        } else {
+          if constexpr (IsPerGroupScale) {
+            v2f32_t value{float(lo) * scale_tensor(vi, mi, ni), float(hi)* scale_tensor(vi, mi, ni)};
+            float2 f32x2 = make_float2(value[0], value[1]);
+            __mt_bfloat162 res = __float22bfloat162_rn(f32x2);
+
+            v2f16_t f16x2;
+            f16x2[0] = mutlass::bfloat16_t(res.x);
+            f16x2[1] = mutlass::bfloat16_t(res.y);
+
+            v2f16_view(vi, mi, ni) = f16x2;
+          } else {
+            v2f32_t value{float(lo), float(hi)};
+            float2 f32x2 = make_float2(value[0], value[1]);
+            __mt_bfloat162 res = __float22bfloat162_rn(f32x2);
+
+            v2f16_t f16x2;
+            f16x2[0] = mutlass::bfloat16_t(res.x);
+            f16x2[1] = mutlass::bfloat16_t(res.y);
+
+            v2f16_view(vi, mi, ni) = f16x2;
+          }
+        }
+      }
+    }
+  }
+  // if (thread(0)) {
+  //   print("f16_tensor = ");print(f16_tensor(0));print("\n");
+  // }
+  asm("CTRL.NOP");
+}
+
+
+template <class LayoutAccum>
+MUTE_DEVICE
+auto
+convert_layout_acc_rowcol(LayoutAccum acc_layout) {
+  if constexpr (decltype(rank<0>(acc_layout))::value == 3) { // MP31
+    // from ((2, VN, VM), MMA_M, MMA_N) to ((VM, MMA_M),(2, VN, MMA_N));
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(rank(acc_layout))::value == 3);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0,2>(l), get<1>(l)), make_layout(get<0,0>(l), get<0,1>(l), get<2>(l)));
+  } else {  // MP22
+    // from((2,4), MMA_M, MMA_N) to ((4, MMA_M), (2, MMA_N))
+    static_assert(decltype(size<0, 0>(acc_layout))::value == 2);
+    static_assert(decltype(size<0, 1>(acc_layout))::value == 4);
+    auto l = acc_layout;
+    return make_layout(make_layout(get<0,1>(l), get<1>(l)), make_layout(get<0,0>(l), get<2>(l)));
+  }
+}
+
+
+template<
+class ElementA,
+class ElementB,
+class ElementD,
+class ElementScale,
+class ElementIdx,
+class ElementTopK
+>
+struct MoeParams {
+    using IndexT = int64_t;
+    using LayoutA = Layout<Shape<IndexT, IndexT>, Stride<IndexT, _1>>;
+    using LayoutB = Layout<Shape<IndexT, IndexT, IndexT>, Stride<IndexT, _1, IndexT>>;
+    using LayoutD = Layout<Shape<IndexT, IndexT>, Stride<IndexT, _1>>;
+    using LayoutScale = Layout<Shape<IndexT, IndexT, IndexT>, Stride<IndexT, _1, IndexT>>;
+    using LayoutOffset = Layout<Shape<IndexT>, Stride<_1>>;
+
+    ElementA* ptr_A;
+    ElementB* ptr_B;
+    ElementD* ptr_D;
+    ElementScale* ptr_scale;
+    ElementIdx* ptr_num_tokens_post_padded;
+    ElementIdx* ptr_offsetA;
+    ElementIdx* ptr_offsetB;
+    ElementTopK* ptr_topk_weights;
+    LayoutA    gmem_layout_A;
+    LayoutB    gmem_layout_B;
+    LayoutD    gmem_layout_D;
+    LayoutScale gmem_layout_scale;
+    LayoutOffset gmem_layout_A_offset;
+    LayoutOffset gmem_layout_B_offset;
+    int32_t EM;
+    const int32_t num_valid_tokens;
+    const int topk;
+    int stride_bse;
+    int stride_bsn;
+    int stride_cm;
+    int stride_cn;
+};
+
+
+template <
+class Params,
+class KernelTraits
+>
+__global__ __launch_bounds__(KernelTraits::NumThreads, 1) void mp22_grouped_gemm_kernel(
+  Params params
+) {
+  using namespace mute;
+  using namespace mutlass;
+
+  using X = Underscore;
+  auto thread_idx = threadIdx.x;
+  using ElementA = typename KernelTraits::ElementA;
+  using ElementB = typename KernelTraits::ElementB;
+  using ElementD = typename KernelTraits::ElementD;
+  using ElementScale = typename KernelTraits::ElementScale;
+  using ElementIdx = typename KernelTraits::ElementIdx;
+  using ElementAccum = typename KernelTraits::ElementAccum;
+  using TiledMma = typename KernelTraits::TiledMma;
+  using TileShapeMNK = typename KernelTraits::TileShapeMNK;
+  using UnderlyingElementB = mute::conditional_t<KernelTraits::WithQuant, mute::uint4_t, ElementB>;
+  constexpr int TileM = size<0>(TileShapeMNK{});
+  constexpr int TileN = size<1>(TileShapeMNK{});
+  constexpr int TileK = size<2>(TileShapeMNK{});
+  constexpr int QuantSizeK = KernelTraits::QuantSizeK;
+  constexpr int num_block_per_quant = QuantSizeK / TileK;
+  constexpr bool is_per_group_scale = KernelTraits::GroupScale;
+  int problem_m = size<0>(params.gmem_layout_D);
+  int problem_n = size<1>(params.gmem_layout_D);
+  int problem_k = size<1>(params.gmem_layout_A);
+  int tile_m_num = mute::ceil_div(params.EM, TileM); //TODO:
+  int tile_n_num = mute::ceil_div(problem_n, TileN);
+  int TileNumInGroup = KernelTraits::GroupSizeM * tile_n_num;
+  int num_tokens_post_padded = params.ptr_num_tokens_post_padded[0];
+
+  auto [n_coord, m_coord, l_coord] = static_cast<uint3>(blockIdx);
+
+  int pid = m_coord * gridDim.x + n_coord;
+  int num_pid_in_group = KernelTraits::GroupSizeM * tile_n_num;
+  int group_id = pid / num_pid_in_group;
+  int first_pid_m = group_id * KernelTraits::GroupSizeM;
+  int group_size_m = min(tile_m_num - first_pid_m, KernelTraits::GroupSizeM);
+  int pid_m = first_pid_m + (pid % num_pid_in_group % group_size_m);
+  int pid_n = (pid % num_pid_in_group) / group_size_m;
+  if (pid_m * TileM >= num_tokens_post_padded) return; //TODO
+  auto mBOffset = make_tensor(make_gmem_ptr(params.ptr_offsetB), params.gmem_layout_B_offset);
+  int offset_expert = mBOffset(pid_m);
+
+  __shared__ ElementA __attribute__((aligned(128))) smem_buf[sizeof(typename KernelTraits::SharedStorage{})];
+  auto shared_storage = reinterpret_cast<typename KernelTraits::SharedStorage*>(smem_buf);
+  Tensor sA = make_tensor(make_smem_ptr(shared_storage->smem_a.data()), typename KernelTraits::SmemLayoutA{}); // (BLK_M,BLK_K)
+  Tensor sB = make_tensor(make_smem_ptr(shared_storage->smem_b.data()), typename KernelTraits::SmemLayoutB{}); // (BLK_N,BLK_K)
+
+  auto mAOffset = make_tensor(make_gmem_ptr(params.ptr_offsetA), params.gmem_layout_A_offset); //(m, batch)
+  auto gAOffset = local_tile(mAOffset, make_shape(TileM), make_coord(pid_m)); //(BLK_M)
+  auto mA = make_tensor(make_gmem_ptr(params.ptr_A), params.gmem_layout_A);
+  auto gA_k = local_tile(mA, make_shape(1, TileK), make_coord(_, _)); // (1, TileK, problem_m, problem_k/TileK)
+  auto mB = make_tensor(make_gmem_ptr<UnderlyingElementB>(params.ptr_B), params.gmem_layout_B);
+  auto gB = local_tile(mB, make_shape(TileN, TileK), make_coord(pid_n, _, offset_expert)); // (BLK_N, Blk_K, problem_k/Blk_K)
+
+  using compose_layout_offsetAA = Layout<
+      Shape<Int<TileM>, Int<KernelTraits::TileKOffsetA>>,
+      Stride<_1, _0>>;
+  using compose_layout_A = Layout<
+      Shape<Int<TileM>, Int<TileK>>,
+      Stride<_0, _1>>;
+
+  using GmemTiledCopyA = typename KernelTraits::GmemTiledCopyA;
+  using GmemTiledCopyAoffset = typename KernelTraits::GmemTiledCopyAoffset;
+  using GmemTiledCopyB = typename KernelTraits::GmemTiledCopyB;
+  using GmemTiledCopyScaleB = typename KernelTraits::GmemTiledCopyScaleB;
+  GmemTiledCopyA gmem_tiled_copy_a;
+  GmemTiledCopyAoffset gmem_tiled_copy_a_offset;
+  GmemTiledCopyB gmem_tiled_copy_b;
+  GmemTiledCopyScaleB gmem_tiled_copy_scale_b;
+  auto gmem_thr_copy_a = gmem_tiled_copy_a.get_slice(threadIdx.x);
+  auto gmem_thr_copy_a_offset = gmem_tiled_copy_a_offset.get_slice(threadIdx.x);
+  auto gmem_thr_copy_b = gmem_tiled_copy_b.get_slice(threadIdx.x);
+  auto gmem_thr_copy_scale_b = gmem_tiled_copy_scale_b.get_slice(threadIdx.x);
+
+  auto tAgAOffset = gmem_thr_copy_a_offset.partition_S(gAOffset.compose(compose_layout_offsetAA{})); //(COPY_ATOM, COPY_M, COPY_K)
+  auto tArAOffset = make_tensor_like<ElementIdx>(tAgAOffset); //(COPY_ATOM, COPY_M, COPY_K)
+
+  // auto tAgA = gmem_thr_copy_a.partition_S(gA);
+  auto tAsA = gmem_thr_copy_a.partition_D(sA); // (ACPY,ACPY_M,ACPY_K)
+  auto tArA = make_fragment_like(tAsA); // (ACPY,ACPY_M,ACPY_K)
+  auto tBgB = gmem_thr_copy_b.partition_S(gB); /// (BCPY,BCPY_N,BCPY_K, problem_k / TileK)
+  auto tBsB = gmem_thr_copy_b.partition_D(sB); // (BCPY,BCPY_N,BCPY_K)
+  auto tBrB = make_fragment_like(tBsB); // (BCPY,BCPY_N,BCPY_K)
+  auto tBrB_quant = make_fragment_like<UnderlyingElementB>(tBsB); // (BCPY,BCPY_N,BCPY_K)
+  auto tBrB_dq = make_fragment_like(tBsB); // (BCPY,BCPY_N,BCPY_K)
+
+
+  //
+  // PREDICATES
+  //
+
+
+  // Construct identity layout for sA and sB
+  Tensor cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA)));    // (BLK_M,BLK_K) -> (blk_m,blk_k)
+  Tensor cB = make_identity_tensor(make_shape(size<0>(sB), size<1>(sB)));    // (BLK_N,BLK_K) -> (blk_n,blk_k)
+
+  // Repeat the partitioning with identity layouts
+  Tensor tAcA = gmem_thr_copy_a.partition_S(cA);                             // (ACPY,ACPY_M,ACPY_K) -> (blk_m,blk_k)
+  Tensor tBcB = gmem_thr_copy_b.partition_S(cB);                             // (BCPY,BCPY_N,BCPY_K) -> (blk_n,blk_k)
+
+  // Set predicates for m bounds
+  auto residue_mnk = make_coord(params.EM - TileM * pid_m, problem_n - TileN * pid_n, problem_k - TileK * size<2>(gB));
+
+  // Tile MMA compute thread partitions and allocate accum
+
+  TiledMma tiled_mma;
+  auto thr_mma = tiled_mma.get_thread_slice(thread_idx);
+  Tensor tCrA  = thr_mma.make_fragment_A(thr_mma.partition_A(sA));           // (MMA,MMA_M,MMA_K)
+  Tensor tCrB  = thr_mma.make_fragment_B(thr_mma.partition_B(sB));           // (MMA,MMA_M,MMA_K)
+
+  auto thr_copy_A       = make_tiled_copy_A(typename KernelTraits::SmemCopyAtomA{}, tiled_mma).get_thread_slice(thread_idx);
+  Tensor tCsA           = thr_copy_A.partition_S(sA);
+  Tensor tCrA_copy_view = thr_copy_A.retile_D(tCrA);
+  MUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));            // M
+
+  auto thr_copy_B       = make_tiled_copy_B(typename KernelTraits::SmemCopyAtomB{}, tiled_mma).get_thread_slice(thread_idx);
+  Tensor tCsB           = thr_copy_B.partition_S(sB);
+  Tensor tCrB_copy_view = thr_copy_B.retile_D(tCrB);
+  MUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));            // N
+
+  auto accum = partition_fragment_C(thr_mma, take<0, 2>(TileShapeMNK{})); // (MMA,MMA_M,MMA_N)
+
+
+  auto m_scaleB = make_tensor(
+    make_gmem_ptr(params.ptr_scale),
+    params.gmem_layout_scale
+  )(_, _, offset_expert);
+
+  auto g_scaleB = local_tile(m_scaleB, make_tile(TileN), make_coord(pid_n, _));
+  auto tBgScaleB = gmem_thr_copy_scale_b.partition_S(g_scaleB(_, 0).compose(Layout<Shape<Int<TileN>, Int<TileK>>, Stride<_1, _0>>{}));
+  auto tBrScaleB = make_tensor_like<ElementScale>(tBgScaleB);
+
+  int k_tile_iter = 0;
+  int k_tile_count = mute::ceil_div(problem_k, TileK);
+  clear(tArA);
+  if constexpr (KernelTraits::WithScale) {
+    clear(tBrB_quant);
+  } else {
+    clear(tBrB);
+  }
+  // Start async loads for 0th k-tile, where we take care of the k residue
+  {
+    // fill(tArAOffset, params.num_valid_tokens);
+    copy(gmem_tiled_copy_a_offset, tAgAOffset, tArAOffset);
+    MUTLASS_PRAGMA_UNROLL
+    for (int m = 0; m < size<1>(tArA); ++m) {
+      if (tArAOffset(0, m, 0) < params.num_valid_tokens) {
+        auto mA_with_offset = gA_k(0, _, tArAOffset(0, m, 0) / params.topk, k_tile_iter).compose(compose_layout_A{}); // (TileK) -> (TileM, TileK)
+        auto tAgA = gmem_thr_copy_a.partition_S(mA_with_offset);
+        copy(gmem_tiled_copy_a, tAgA(_,m,_), tArA(_,m,_));
+
+      }
+    }
+
+    if constexpr (KernelTraits::WithQuant) {
+      copy(gmem_tiled_copy_b, tBgB(_,_,_,k_tile_iter), tBrB_quant(_,_,_));
+    } else {
+      copy(gmem_tiled_copy_b, tBgB(_,_,_,k_tile_iter), tBrB(_,_,_));
+    }
+
+    if constexpr (is_per_group_scale) {
+      auto tBgScaleB = gmem_thr_copy_scale_b.partition_S(g_scaleB(_, k_tile_iter / num_block_per_quant).compose(Layout<Shape<Int<TileN>, Int<TileK>>, Stride<_1, _0>>{}));
+      copy(tBgScaleB, tBrScaleB);
+    }
+
+    ++k_tile_iter;
+    --k_tile_count;
+  }
+
+
+
+  //
+  // Prologue
+  //
+  // Copy rmem to smem
+  copy(tArA, tAsA);
+  if constexpr (KernelTraits::WithQuant) {
+    dequant_w4<ElementA, is_per_group_scale>(tBrScaleB, tBrB_quant, tBrB);
+    // if constexpr (is_per_group_scale) {
+    //   MUTLASS_PRAGMA_UNROLL
+    //   for (int i = 0; i < size(tBrB); ++i) {
+    //     // tBrB(i) = ElementA(tBrB(i) * tBrScaleB(i));
+    //   }
+    // }
+  }
+
+  copy(tBrB, tBsB);
+
+  // Clear accum
+  __syncthreads();
+
+  // Load A, B smem->rmem for k=0
+  copy(tCsA(_,_,0), tCrA_copy_view(_,_,0));
+  copy(tCsB(_,_,0), tCrB_copy_view(_,_,0));
+  //
+  // Mainloop
+  //
+
+  // Size of the k-tiles's outer product mode (k)
+  auto K_BLOCK_MAX = size<2>(tCrA);
+
+  MUTLASS_PRAGMA_NO_UNROLL
+  while (k_tile_count > -1)
+  {
+    // Pipeline the outer products with a static for loop
+    for_each(make_int_sequence<K_BLOCK_MAX>{}, [&] (auto k_block)
+    {
+      // if ((k_block == K_BLOCK_MAX - 1) && k_tile_count != 0)
+      if ((k_block == K_BLOCK_MAX - 1))
+      {
+        __syncthreads();
+
+        // Copy rmem to smem
+        copy(tArA, tAsA);
+        if constexpr (KernelTraits::WithQuant) {
+          dequant_w4<ElementA, is_per_group_scale>(tBrScaleB, tBrB_quant, tBrB);
+          // if constexpr (is_per_group_scale) {
+          //   MUTLASS_PRAGMA_UNROLL
+          //   for (int i = 0; i < size(tBrB); ++i) {
+          //     // tBrB(i) = ElementA(tBrB(i) * tBrScaleB(i));
+          //   }
+          // }
+        }
+        copy(tBrB, tBsB);
+        __syncthreads();
+      }
+
+      // Load A, B smem->rmem for k+1
+      int k_block_next = (k_block + Int<1>{}) % K_BLOCK_MAX;    // static
+      copy(tCsA(_,_,k_block_next), tCrA_copy_view(_,_,k_block_next));
+      copy(tCsB(_,_,k_block_next), tCrB_copy_view(_,_,k_block_next));
+      // if (k_block == 0 && k_tile_count > 0) //TODO: will hang
+      if (k_block == 0)
+      {
+        MUTLASS_PRAGMA_UNROLL
+        for (int m = 0; m < size<1>(tArAOffset); ++m) {
+          if (tArAOffset(0, m, 0) < params.num_valid_tokens) {
+            auto mA_with_offset = gA_k(0, _, tArAOffset(0, m, 0) / params.topk, k_tile_iter).compose(compose_layout_A{});
+            auto tAgA = gmem_thr_copy_a.partition_S(mA_with_offset);
+            copy(gmem_tiled_copy_a, tAgA(_,m,_), tArA(_,m,_));
+          }
+        }
+        if constexpr (KernelTraits::WithQuant) {
+          copy(gmem_tiled_copy_b, tBgB(_,_,_,k_tile_iter), tBrB_quant(_,_,_));
+        } else {
+          copy(gmem_tiled_copy_b, tBgB(_,_,_,k_tile_iter), tBrB(_,_,_));
+        }
+
+        if constexpr (is_per_group_scale) {
+          if (k_tile_iter % num_block_per_quant == 0) {
+            auto tBgScaleB = gmem_thr_copy_scale_b.partition_S(g_scaleB(_, k_tile_iter / num_block_per_quant).compose(Layout<Shape<Int<TileN>, Int<TileK>>, Stride<_1, _0>>{}));
+            copy(tBgScaleB, tBrScaleB);
+          }
+        }
+        if (--k_tile_count > 0) ++k_tile_iter;
+      }
+      mute::gemm(tiled_mma, accum, tCrA(_,_,k_block), tCrB(_,_,k_block), accum);
+    });
+  }
+
+  if constexpr (true) {
+    Tensor cO = make_identity_tensor(select<0,1>(TileShapeMNK{}));
+    Tensor tOcO = thr_mma.partition_C(cO);
+
+    Tensor tOcO_mn = make_tensor(tOcO.data(), convert_layout_acc_rowcol(tOcO.layout()));
+    Tensor accum_mn = make_tensor(accum.data(), convert_layout_acc_rowcol(accum.layout()));
+
+    Tensor scale = make_tensor<ElementScale>(size<1>(accum_mn));
+    if constexpr (KernelTraits::WithQuant) {
+      MUTE_UNROLL
+      for (int ni = 0; ni < size<1>(accum_mn); ++ni) {
+        int64_t offs_scale_n = offset_expert * params.stride_bse + (pid_n * TileN + get<1>(tOcO_mn(0,ni))) * params.stride_bsn;
+        scale(ni) = params.ptr_scale[offs_scale_n];
+      }
+    }
+
+    MUTE_UNROLL
+    for (int mi = 0; mi < size<0>(accum_mn); ++mi) {
+      int64_t offs_token_id = pid_m * TileM + get<0>(tOcO_mn(mi, 0));
+      int64_t offs_token = params.ptr_offsetA[offs_token_id];
+      float moe_weight = offs_token < params.num_valid_tokens ? params.ptr_topk_weights[offs_token] : 0.0f;
+
+      MUTE_UNROLL
+      for (int ni = 0; ni < size<1>(accum_mn); ++ni) {
+
+        int64_t offs_n = pid_n * TileN + get<1>(tOcO_mn(0,ni));
+        if constexpr (KernelTraits::WithQuant) {
+          constexpr float rescale = 1.0f / 16.0f;
+          if constexpr (!is_per_group_scale) {
+            accum_mn(mi, ni) = accum_mn(mi, ni) * float(scale[ni]) * rescale;
+          }
+        }
+        if (offs_token < params.num_valid_tokens) {
+
+          if constexpr (KernelTraits::MulRoutedWeight) {
+            params.ptr_D[offs_token * params.stride_cm + offs_n * params.stride_cn] = ElementD(accum_mn(mi, ni) * moe_weight);
+          } else {
+            params.ptr_D[offs_token * params.stride_cm + offs_n * params.stride_cn] = ElementD(accum_mn(mi, ni));
+          }
+        }
+      }
+    }
+  }
+
+}
diff --git a/vllm_musa/csrc/moe/mutlass/fused_moe_impl/mp22_kernel_traits.hpp b/vllm_musa/csrc/moe/mutlass/fused_moe_impl/mp22_kernel_traits.hpp
new file mode 100644
index 000000000..46ba8576f
--- /dev/null
+++ b/vllm_musa/csrc/moe/mutlass/fused_moe_impl/mp22_kernel_traits.hpp
@@ -0,0 +1,162 @@
+/***************************************************************************************************
+ * Copyright (c) 2024 - 2025 Moore Threads Technology Co., Ltd("Moore Threads"). All rights reserved.
+ * SPDX-License-Identifier: BSD-3-Clause
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright notice, this
+ * list of conditions and the following disclaimer.
+ *
+ * 2. 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.
+ *
+ * 3. Neither the name of the copyright holder 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 THE COPYRIGHT HOLDER OR CONTRIBUTORS 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 TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ **************************************************************************************************/
+
+ #pragma once
+
+ #include <mute/tensor.hpp>
+ #include <mute/atom/mma_atom.hpp>
+ #include <mute/atom/copy_atom.hpp>
+
+ #include <mutlass/mutlass.h>
+ #include <mutlass/gemm/gemm.h>
+ #include <mutlass/gemm/collective/collective_builder.hpp>
+ #include "mutlass/epilogue/thread/linear_combination.h"
+ #include <mute/atom/copy_atom.hpp>
+
+ using namespace mute;
+
+ template<
+ class ElementA,
+ class ElementB,
+ class SmemLayoutA,
+ class SmemLayoutB
+ >
+ struct SharedStorage {
+   mute::array_aligned<ElementA, mute::cosize_v<SmemLayoutA>, 128> smem_a;
+   mute::array_aligned<ElementB, mute::cosize_v<SmemLayoutB>, 128> smem_b;
+ };
+
+ template <
+   class TileShapeMNK_,
+   class ElementA_,
+   class ElementB_,
+   class ElementD_,
+   class ElementIdx_,
+   class ElementScale_,
+   class ElementTopk_,
+   int AlignmentA,
+   int AlignmentB,
+   int GroupSizeM_,
+   bool WithScale_,
+   bool WithQuant_,
+   bool MulRoutedWeight_,
+   bool GroupScale_,
+   int QuantSizeK_
+ >
+ struct GroupGemmKernelTraits {
+
+   using ElementA = mute::conditional_t<mute::is_same_v<ElementA_, float>, tfloat32_t, ElementA_>;
+   using ElementB = mute::conditional_t<mute::is_same_v<ElementB_, float>, tfloat32_t, ElementB_>;
+   using ElementD = ElementD_;
+   using ElementScale = ElementScale_;
+   using ElementIdx = ElementIdx_;
+   using ElementAccum = float;
+   using ElementTopk = ElementTopk_;
+
+   using StrideA = mutlass::detail::TagToStrideA_t<mutlass::layout::RowMajor>;
+   using StrideB = mutlass::detail::TagToStrideB_t<mutlass::layout::ColumnMajor>;
+   using StrideD = mutlass::detail::TagToStrideC_t<mutlass::layout::RowMajor>;
+
+   using TileShapeMNK = TileShapeMNK_;
+   static constexpr int TileSizeM = size<0>(TileShapeMNK{});
+   static constexpr int TileSizeN = size<1>(TileShapeMNK{});
+   static constexpr int TileSizeK = size<2>(TileShapeMNK{});
+   static constexpr int GroupSizeM = GroupSizeM_;
+   static constexpr bool WithScale = WithScale_;
+   static constexpr bool WithQuant = WithQuant_;
+   static constexpr bool MulRoutedWeight = MulRoutedWeight_;
+   static constexpr bool GroupScale = GroupScale_;
+   static const int QuantSizeK = QuantSizeK_;
+
+   using MmaOp = decltype(mutlass::gemm::collective::detail::mp22_mma_operation_select<ElementA, StrideA, ElementA, StrideB>());
+
+   using TiledMma = decltype(mutlass::gemm::collective::detail::mp22_make_tensorop_tiled_mma<TileShapeMNK, MmaOp>());
+   static constexpr int NumThreads = size(TiledMma{});
+
+   static constexpr int ThreadCount = size(TiledMma{});
+   // A
+   using SmemLayoutAtomA = decltype(mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementA, StrideA>());
+   using SmemCopyAtomA = Copy_Atom<DefaultCopy, ElementA>;
+   using SmemLayoutA = decltype(tile_to_shape(
+       SmemLayoutAtomA{},
+       make_shape(shape<0>(TileShapeMNK{}), shape<2>(TileShapeMNK{}))));
+
+   // B
+   using SmemLayoutAtomB = decltype(mutlass::gemm::collective::detail::make_mp22_smem_atom_layout<ElementA, StrideB>());
+   using SmemCopyAtomB = Copy_Atom<DefaultCopy, ElementA>;
+   using SmemLayoutB = decltype(tile_to_shape(
+     SmemLayoutAtomB{},
+     make_shape(shape<1>(TileShapeMNK{}), shape<2>(TileShapeMNK{}))));
+
+   using SharedStorage = SharedStorage<ElementA, ElementA, SmemLayoutA, SmemLayoutB>;
+
+   static constexpr int threads_major_a = TileSizeK   / AlignmentA;
+   static constexpr int threads_minor_a = ThreadCount / threads_major_a;
+   using GmemTiledCopyAOperation = UniversalCopy<uint_bit_t<AlignmentA * sizeof_bits_v<ElementA>>>;
+   using GmemTiledCopyA = decltype(make_tiled_copy(
+     Copy_Atom<GmemTiledCopyAOperation, ElementA>{},
+     Layout<Shape <Int<threads_minor_a>, Int<threads_major_a>>,
+            Stride<Int<threads_major_a>,                _1>>{},
+     Layout<Shape<_1, Int<AlignmentA>>>{}));
+
+   static constexpr int threads_major_b = TileSizeK   / AlignmentB;
+   static constexpr int threads_minor_b = ThreadCount / threads_major_b;
+   using GmemTiledCopyBOperation = UniversalCopy<uint_bit_t<AlignmentB * sizeof_bits_v<ElementB>>>;
+   using GmemTiledCopyB = decltype(make_tiled_copy(
+     Copy_Atom<GmemTiledCopyBOperation, ElementB>{},
+     Layout<Shape <Int<threads_minor_b>, Int<threads_major_b>>,
+            Stride<Int<threads_major_b>,                _1>>{},
+     Layout<Shape<_1, Int<AlignmentB>>>{}));
+
+   using GmemTiledCopyScaleBOperation = UniversalCopy<uint_bit_t<AlignmentB * sizeof_bits_v<ElementScale>>>;
+   using GmemTiledCopyScaleB = decltype(make_tiled_copy(
+     Copy_Atom<GmemTiledCopyScaleBOperation, ElementScale>{},
+     Layout<Shape <Int<threads_minor_b>, Int<threads_major_b>>,
+            Stride<Int<threads_major_b>,                _1>>{},
+     Layout<Shape<_1, Int<AlignmentB>>>{}));
+
+   // using ScalePermTile = Layout<Shape<_32, _2, _4>, Stride<_1, _128, _32>>;
+   using GmemTiledCopyAoffsetOperation = UniversalCopy<uint_bit_t<sizeof_bits_v<ElementIdx>>>;
+   using GmemTiledCopyAoffset = decltype(make_tiled_copy(
+     Copy_Atom<GmemTiledCopyAoffsetOperation, ElementIdx>{},
+     Layout<Shape <Int<threads_minor_a>, Int<threads_major_a>>,
+            Stride<Int<threads_major_a>,                _1>>{},
+     Layout<Shape<_1, _1>>{}));
+   static constexpr int  TileKOffsetA = threads_major_a;
+
+  //  using GmemTiledCopyBoffsetOperation = UniversalCopy<uint_bit_t<AlignmentB * sizeof_bits_v<ElementB>>>;
+  //  using GmemTiledCopyB = decltype(make_tiled_copy(
+  //    Copy_Atom<GmemTiledCopyBoffsetOperation, ElementB>{},
+  //    Layout<Shape <Int<threads_minor_b>, Int<threads_major_b>>,
+  //           Stride<Int<threads_major_b>,                _1>>{},
+  //    Layout<Shape<_1, Int<AlignmentB>>>{}));
+
+ };
diff --git a/vllm_musa/csrc/moe/mutlass/moe_mutlass.mu b/vllm_musa/csrc/moe/mutlass/moe_mutlass.mu
new file mode 100644
index 000000000..d0b404ebb
--- /dev/null
+++ b/vllm_musa/csrc/moe/mutlass/moe_mutlass.mu
@@ -0,0 +1,326 @@
+#include <torch/torch.h>
+#include <iostream>
+#include "fused_moe_impl/mp22_kernel_traits.hpp"
+#include "fused_moe_impl/mp22_group_kernels.hpp"
+
+
+#include "mutlass/mutlass.h"
+
+#include "mute/tensor.hpp"
+#include "mutlass/tensor_ref.h"
+#include "mutlass/epilogue/collective/default_epilogue.hpp"
+#include "mutlass/epilogue/thread/linear_combination.h"
+#include "mutlass/gemm/dispatch_policy.hpp"
+#include "mutlass/gemm/collective/collective_builder.hpp"
+#include "mutlass/epilogue/collective/collective_builder.hpp"
+#include "mutlass/gemm/device/gemm_universal_adapter.h"
+#include "mutlass/gemm/kernel/gemm_universal.hpp"
+
+#include "mutlass/util/command_line.h"
+#include "mutlass/util/distribution.h"
+#include "mutlass/util/host_tensor.h"
+#include "mutlass/util/packed_stride.hpp"
+#include "mutlass/util/tensor_view_io.h"
+#include "mutlass/util/reference/device/gemm.h"
+#include "mutlass/util/reference/device/tensor_compare.h"
+#include "mutlass/util/reference/device/tensor_fill.h"
+
+#include "../../type_convert.muh"
+#include "../../dispatch_utils.h"
+
+using namespace mute;
+
+
+/////////////////////////////////////////////////////////////////////////////////////////////////
+/// GEMM kernel configurations
+/////////////////////////////////////////////////////////////////////////////////////////////////
+
+using ElementScale       = float;                // Element type for scaling factor
+using ElementIdx         = int32_t;
+constexpr int AlignmentA  = 8;    // Alignment of A matrix in units of elements (up to 16 bytes)
+constexpr int AlignmentB  = 8;    // Alignment of B matrix in units of elements (up to 16 bytes)
+constexpr int AlignmentC  = 1;    // Alignment of C matrix in units of elements (up to 16 bytes)
+constexpr int AlignmentD  = 1;    // Alignment of C matrix in units of elements (up to 16 bytes)
+constexpr int Stages   = 2;
+
+using TileShape           = Shape<_32, _32, _64>;   // Threadblock-level tile size
+constexpr int GroupSizeM   = 1;
+
+constexpr bool check_error = false;
+
+template <typename AType, typename BType, typename CType, typename ScoreType, bool WITH_W4A16_QUANT>
+void fused_moe_mutlass_kernel_launcher(
+    AType *A, BType *B, CType *C, float *B_scale,
+    int32_t *num_tokens_post_padded, int32_t *sorted_token_ids, int32_t *expert_ids,
+    ScoreType *topk_weights, int32_t *topk_ids,
+    const int32_t M, const int32_t N, const int32_t K, const int32_t E, const int EM, const int topk_ids_elements,
+    const int stride_am, const int stride_ak, const int stride_be, const int stride_bk, const int stride_bn,
+    const int stride_cm, const int stride_cn, const int stride_asm, const int stride_ask,
+    const int stride_bse, const int stride_bsk, const int stride_bsn,
+    const int m_offset, const int expert_offset, const int topk, bool mul_routed_weight, int group_size) {
+
+    if (K % size<2>(TileShape{}) != 0 || N % size<1>(TileShape{}) != 0) {
+        std::cerr << "K must be divisible by AlignmentA and AlignmentB and N must divisible by TileShapeN" << std::endl;
+        std::abort();
+    }
+
+    using MoeParams = MoeParams<AType, BType, CType, ElementScale, ElementIdx, ScoreType>;
+
+    typename MoeParams::LayoutA g_layout_a = make_layout(make_shape(int64_t(M), int64_t(K)),
+                            make_stride(int64_t(stride_am), _1{}));
+    typename MoeParams::LayoutB g_layout_b = make_layout(make_shape(int64_t(N), int64_t(K), int64_t(E)),
+                            make_stride(int64_t(K), _1{}, int64_t(N*K)));
+    typename MoeParams::LayoutD g_layout_d = make_layout(make_shape(int64_t(M*topk), int64_t(N)),
+                            make_stride(int64_t(stride_cm), _1{}));
+    typename MoeParams::LayoutScale g_layout_scale;
+    if (group_size != -1) {
+        g_layout_scale = make_layout(make_shape(int64_t(N), int64_t(K / group_size), int64_t(E)),
+                                     make_stride(int64_t(K / group_size), _1{}, int64_t(K / group_size * N)));
+    } else {
+        g_layout_scale = make_layout(make_shape(int64_t(N), int64_t(1), int64_t(E)),
+                                    make_stride(int64_t(1), _1{}, int64_t(stride_bse)));
+    }
+    typename MoeParams::LayoutOffset g_layout_m_offset = make_layout(make_shape(int64_t(m_offset)),
+                                        make_stride(_1{}));
+    typename MoeParams::LayoutOffset g_layout_expert_offset = make_layout(make_shape(int64_t(expert_offset)),
+                                                        make_stride(_1{}));
+    if(check_error) {
+        print("EM = ");print(EM);print("\n");
+        print("topk_ids_elements = ");print(topk_ids_elements);print("\n");
+        print("g_layout_a = ");print(g_layout_a);print("\n");
+        print("g_layout_b = ");print(g_layout_b);print("\n");
+        print("g_layout_d = ");print(g_layout_d);print("\n");
+        print("g_layout_scale = ");print(g_layout_scale);print("\n");
+        print("g_layout_m_offset = ");print(g_layout_m_offset);print("\n");
+        print("g_layout_expert_offset = ");print(g_layout_expert_offset);print("\n");
+        print("WITH_W4A16_QUANT = ");print(WITH_W4A16_QUANT);print("\n");
+        print("mul_routed_weight = ");print(mul_routed_weight);print("\n");
+        print("GroupScale = ");print(group_size);print("\n");
+
+    }
+
+    auto params = MoeParams{
+        A,
+        B,
+        C,
+        B_scale, // scale
+        num_tokens_post_padded, // num_tokens_post_padded
+        sorted_token_ids, // m_offset_ptr
+        expert_ids, // expert_offset_ptr
+        topk_weights,
+        g_layout_a,
+        g_layout_b,
+        g_layout_d,
+        g_layout_scale,
+        g_layout_m_offset,
+        g_layout_expert_offset,
+        EM,
+        topk_ids_elements,
+        topk,
+        stride_bse,
+        stride_bsn,
+        stride_cm,
+        stride_cn
+    };
+
+    if constexpr (check_error) {
+        cudaDeviceSynchronize();
+        auto error = cudaGetLastError();
+        if (error != cudaSuccess) {
+            printf("before launch mutlass moe [ERROR] %s\n", cudaGetErrorString(error));
+            exit(-1);
+        }
+
+    }
+
+    if (mul_routed_weight) {
+        if (group_size == 128){
+            using KernelTraits = GroupGemmKernelTraits<TileShape, AType, BType, CType, ElementIdx, ElementScale, ScoreType, AlignmentA, AlignmentB, GroupSizeM, WITH_W4A16_QUANT, WITH_W4A16_QUANT, true, true, 128>;
+            dim3 block_dim(KernelTraits::NumThreads);
+            dim3 grid_dims(ceil_div(N, size<1>(TileShape{})), ceil_div(EM, size<0>(TileShape{})), 1);
+
+            const musaStream_t stream = at::cuda::getCurrentCUDAStream();
+            mp22_grouped_gemm_kernel<MoeParams,  KernelTraits><<<grid_dims, block_dim, 0, stream>>>(
+                params
+            );
+        } else if (group_size == 64) {
+            using KernelTraits = GroupGemmKernelTraits<TileShape, AType, BType, CType, ElementIdx, ElementScale, ScoreType, AlignmentA, AlignmentB, GroupSizeM, WITH_W4A16_QUANT, WITH_W4A16_QUANT, true, true, 64>;
+            dim3 block_dim(KernelTraits::NumThreads);
+            dim3 grid_dims(ceil_div(N, size<1>(TileShape{})), ceil_div(EM, size<0>(TileShape{})), 1);
+
+            const musaStream_t stream = at::cuda::getCurrentCUDAStream();
+            mp22_grouped_gemm_kernel<MoeParams,  KernelTraits><<<grid_dims, block_dim, 0, stream>>>(
+                params
+            );
+        } else {
+            using KernelTraits = GroupGemmKernelTraits<TileShape, AType, BType, CType, ElementIdx, ElementScale, ScoreType, AlignmentA, AlignmentB, GroupSizeM, WITH_W4A16_QUANT, WITH_W4A16_QUANT, true, false, 64>;
+            dim3 block_dim(KernelTraits::NumThreads);
+            dim3 grid_dims(ceil_div(N, size<1>(TileShape{})), ceil_div(EM, size<0>(TileShape{})), 1);
+
+            const musaStream_t stream = at::cuda::getCurrentCUDAStream();
+            mp22_grouped_gemm_kernel<MoeParams,  KernelTraits><<<grid_dims, block_dim, 0, stream>>>(
+                params
+            );
+        }
+    }
+    else
+    {
+        if (group_size == 128){
+            using KernelTraits = GroupGemmKernelTraits<TileShape, AType, BType, CType, ElementIdx, ElementScale, ScoreType, AlignmentA, AlignmentB, GroupSizeM, WITH_W4A16_QUANT, WITH_W4A16_QUANT, false, true, 128>;
+            dim3 block_dim(KernelTraits::NumThreads);
+            dim3 grid_dims(ceil_div(N, size<1>(TileShape{})), ceil_div(EM, size<0>(TileShape{})), 1);
+
+            const musaStream_t stream = at::cuda::getCurrentCUDAStream();
+            mp22_grouped_gemm_kernel<MoeParams, KernelTraits><<<grid_dims, block_dim, 0, stream>>>(
+                params
+            );
+        } else if (group_size == 64){
+            using KernelTraits = GroupGemmKernelTraits<TileShape, AType, BType, CType, ElementIdx, ElementScale, ScoreType, AlignmentA, AlignmentB, GroupSizeM, WITH_W4A16_QUANT, WITH_W4A16_QUANT, false, true, 64>;
+            dim3 block_dim(KernelTraits::NumThreads);
+            dim3 grid_dims(ceil_div(N, size<1>(TileShape{})), ceil_div(EM, size<0>(TileShape{})), 1);
+
+            const musaStream_t stream = at::cuda::getCurrentCUDAStream();
+            mp22_grouped_gemm_kernel<MoeParams, KernelTraits><<<grid_dims, block_dim, 0, stream>>>(
+                params
+            );
+        } else {
+            using KernelTraits = GroupGemmKernelTraits<TileShape, AType, BType, CType, ElementIdx, ElementScale, ScoreType, AlignmentA, AlignmentB, GroupSizeM, WITH_W4A16_QUANT, WITH_W4A16_QUANT, false, false, 64>;
+            dim3 block_dim(KernelTraits::NumThreads);
+            dim3 grid_dims(ceil_div(N, size<1>(TileShape{})), ceil_div(EM, size<0>(TileShape{})), 1);
+
+            const musaStream_t stream = at::cuda::getCurrentCUDAStream();
+            mp22_grouped_gemm_kernel<MoeParams, KernelTraits><<<grid_dims, block_dim, 0, stream>>>(
+                params
+            );
+        }
+    }
+
+    if constexpr (check_error) {
+        cudaDeviceSynchronize();
+        auto error = cudaGetLastError();
+        if (error != cudaSuccess) {
+            printf("after launch mutlass moe [ERROR] %s\n", cudaGetErrorString(error));
+            exit(-1);
+        }
+
+    }
+}
+
+
+
+#define CALL_MUTLASS_MOE_LAUNCH(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, WITH_W4A16_QUANT)                 \
+fused_moe_mutlass_kernel_launcher<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, WITH_W4A16_QUANT>(     \
+    static_cast<_ADTYPE*>(A.data_ptr()),                                                                \
+    static_cast<_BDTYPE*>(B.data_ptr()),                                                                \
+    static_cast<_ADTYPE*>(C.data_ptr()),                                                                \
+    B_scale ? static_cast<float*>(B_scale->data_ptr()) : nullptr,                                       \
+    static_cast<int32_t *>(num_tokens_post_padded.data_ptr()),                                          \
+    static_cast<int32_t *>(sorted_token_ids.data_ptr()),                                                \
+    static_cast<int32_t *>(expert_ids.data_ptr()),                                                      \
+    static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights.data_ptr()),                                          \
+    static_cast<int32_t *>(topk_ids.data_ptr()),                                                        \
+    bseqlen, reduce_size, hidden_size, num_experts, EM, topk_ids.numel(),                               \
+    A.stride(0), A.stride(1), B.stride(0), B.stride(2), B.stride(1),                                    \
+    C.stride(1), C.stride(2), 0, 0, B_scale ? B_scale->stride(0) : 0,                                   \
+    B_scale ? B_scale->stride(2) : 0, B_scale ? B_scale->stride(1) : 0,                                 \
+    m_offset, expert_offset, topk, mul_routed_weight, group_size);
+
+
+void fused_moe_mutlass(
+    torch::Tensor &A,   // [32768, 2048]    [bseqlen, hidden_size]
+    torch::Tensor &B,   // [64, 2816, 256]  []
+    torch::Tensor &C,   // [32768, 6, 2846]
+    const c10::optional<torch::Tensor> &A_scale,    // use for fp8
+    const c10::optional<torch::Tensor> &B_scale,    // use for fp8/int
+    torch::Tensor &topk_weights,
+    torch::Tensor &topk_ids,
+    torch::Tensor &sorted_token_ids,
+    torch::Tensor &expert_ids,
+    torch::Tensor &num_tokens_post_padded,
+    bool mul_routed_weight,
+    int topk,
+    bool use_int4_w4a16,
+    int BLOCK_SIZE_M,
+    int BLOCK_SIZE_N,
+    int BLOCK_SIZE_K,
+    int GROUP_SIZE_M,
+    int group_size)
+{
+    TORCH_CHECK(A.dim() == 2, "A must be dim 2.")
+    TORCH_CHECK(B.dim() == 3, "B must be dim 3.")
+    TORCH_CHECK(C.dim() == 3, "C must be dim 3.")
+
+    TORCH_CHECK(group_size == 128 || group_size == 64 || group_size == -1,
+    "group_size only support 128 or 64 for per-group w4a16, and -1 for w16a16 and per-channel w4a16");
+
+    int32_t bseqlen = A.size(0);        // M
+    int32_t hidden_size = A.size(1);    // K
+    int32_t num_experts = B.size(0);    // E
+    int32_t reduce_size = B.size(1);    // N
+    int32_t EM = sorted_token_ids.size(0);
+    int32_t m_offset = sorted_token_ids.size(0);
+    int32_t expert_offset = expert_ids.size(0);
+
+    bool tile_shape_check = (size<0>(TileShape{}) == BLOCK_SIZE_M) &&
+                   (size<1>(TileShape{}) == BLOCK_SIZE_N) &&
+                   (size<2>(TileShape{}) == BLOCK_SIZE_K) &&
+                   (GROUP_SIZE_M == GroupSizeM);
+    assert(tile_shape_check && "Error: TileShape dimensions do not match the expected sizes.");
+
+    const at::musa::OptionalMUSAGuard device_guard(device_of(A));
+
+    // CALL_MUTLASS_MOE_LAUNCH(mutlass::bfloat16_t, mutlass::bfloat16_t, float, false)
+    // CALL_MUTLASS_MOE_LAUNCH(mutlass::half_t, mutlass::int4b_t, float, true)
+
+    if (A.scalar_type() == at::ScalarType::BFloat16)
+    {
+        if (topk_weights.scalar_type() == at::ScalarType::Float)
+        {
+            if (use_int4_w4a16) {
+                CALL_MUTLASS_MOE_LAUNCH(mutlass::bfloat16_t, mutlass::int4b_t, float, true)
+            }
+            else
+            {
+                CALL_MUTLASS_MOE_LAUNCH(mutlass::bfloat16_t, mutlass::bfloat16_t, float, false)
+            }
+        }
+        else if (topk_weights.scalar_type() == at::ScalarType::BFloat16)
+        {
+            if (use_int4_w4a16) {
+                CALL_MUTLASS_MOE_LAUNCH(mutlass::bfloat16_t, mutlass::int4b_t, mutlass::bfloat16_t, true)
+            }
+            else
+            {
+                CALL_MUTLASS_MOE_LAUNCH(mutlass::bfloat16_t, mutlass::bfloat16_t, mutlass::bfloat16_t, false)
+            }
+        }
+    }
+    else if (A.scalar_type() == at::ScalarType::Half)
+    {
+        if (topk_weights.scalar_type() == at::ScalarType::Float)
+        {
+            if (use_int4_w4a16) {
+                CALL_MUTLASS_MOE_LAUNCH(mutlass::half_t, mutlass::int4b_t, float, true)
+            }
+            else
+            {
+                CALL_MUTLASS_MOE_LAUNCH(mutlass::half_t, mutlass::half_t, float, false)
+            }
+        }
+        else if (topk_weights.scalar_type() == at::ScalarType::Half)
+        {
+            if (use_int4_w4a16) {
+                CALL_MUTLASS_MOE_LAUNCH(mutlass::half_t, mutlass::int4b_t, mutlass::half_t, true)
+            }
+            else
+            {
+                CALL_MUTLASS_MOE_LAUNCH(mutlass::half_t, mutlass::half_t, mutlass::half_t, false)
+            }
+        }
+    }
+    else
+    {
+        printf("no support on mutlass moe\n");
+        assert(false);
+    }
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/moe/topk_softmax_kernels.mu b/vllm_musa/csrc/moe/topk_softmax_kernels.mu
new file mode 100644
index 000000000..aebc1334b
--- /dev/null
+++ b/vllm_musa/csrc/moe/topk_softmax_kernels.mu
@@ -0,0 +1,500 @@
+/*
+ * Adapted from https://github.com/NVIDIA/TensorRT-LLM/blob/v0.7.1/cpp/tensorrt_llm/kernels/mixtureOfExperts/moe_kernels.cu
+ * Copyright (c) 2024, The vLLM team.
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <torch/all.h>
+
+#include "../musa.h"
+#include <cub/cub.cuh>
+
+#include "../cuda_compat.h"
+
+#define MAX(a, b) ((a) > (b) ? (a) : (b))
+#define MIN(a, b) ((a) < (b) ? (a) : (b))
+
+namespace vllm_musa {
+namespace moe {
+
+/// Aligned array type
+template <
+    typename T,
+    /// Number of elements in the array
+    int N,
+    /// Alignment requirement in bytes
+    int Alignment = sizeof(T) * N
+>
+class alignas(Alignment) AlignedArray {
+    float data[N];
+};
+
+// ====================== Softmax things ===============================
+// We have our own implementation of softmax here so we can support transposing the output
+// in the softmax kernel when we extend this module to support expert-choice routing.
+template <int TPB>
+__launch_bounds__(TPB) __global__
+    void moeSoftmax(const float* input, const bool* finished, float* output, const int num_cols)
+{
+    using BlockReduce = cub::BlockReduce<float, TPB>;
+    __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+    __shared__ float normalizing_factor;
+    __shared__ float float_max;
+
+    const int thread_row_offset = blockIdx.x * num_cols;
+
+    cub::Sum sum;
+    float threadData(-FLT_MAX);
+
+    // Don't touch finished rows.
+    if ((finished != nullptr) && finished[blockIdx.x])
+    {
+        return;
+    }
+
+    for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
+    {
+        const int idx = thread_row_offset + ii;
+        threadData = max(static_cast<float>(input[idx]), threadData);
+    }
+
+    const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, cub::Max());
+    if (threadIdx.x == 0)
+    {
+        float_max = maxElem;
+    }
+    __syncthreads();
+
+    threadData = 0;
+
+    for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
+    {
+        const int idx = thread_row_offset + ii;
+        threadData += exp((static_cast<float>(input[idx]) - float_max));
+    }
+
+    const auto Z = BlockReduce(tmpStorage).Reduce(threadData, sum);
+
+    if (threadIdx.x == 0)
+    {
+        normalizing_factor = 1.f / Z;
+    }
+    __syncthreads();
+
+    for (int ii = threadIdx.x; ii < num_cols; ii += TPB)
+    {
+        const int idx = thread_row_offset + ii;
+        const float val = exp((static_cast<float>(input[idx]) - float_max)) * normalizing_factor;
+        output[idx] = val;
+    }
+}
+
+template <int TPB>
+__launch_bounds__(TPB) __global__ void moeTopK(const float* inputs_after_softmax, const bool* finished, float* output,
+    int* indices, int* source_rows, const int num_experts, const int k, const int start_expert, const int end_expert)
+{
+
+    using cub_kvp = cub::KeyValuePair<int, float>;
+    using BlockReduce = cub::BlockReduce<cub_kvp, TPB>;
+    __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+    cub_kvp thread_kvp;
+    cub::ArgMax arg_max;
+
+    const int num_rows = gridDim.x;
+    const int block_row = blockIdx.x;
+
+    const bool row_is_active = finished ? !finished[block_row] : true;
+    const int thread_read_offset = blockIdx.x * num_experts;
+    for (int k_idx = 0; k_idx < k; ++k_idx)
+    {
+        thread_kvp.key = 0;
+        thread_kvp.value = -1.f; // This is OK because inputs are probabilities
+
+        cub_kvp inp_kvp;
+        for (int expert = threadIdx.x; expert < num_experts; expert += TPB)
+        {
+            const int idx = thread_read_offset + expert;
+            inp_kvp.key = expert;
+            inp_kvp.value = inputs_after_softmax[idx];
+
+            for (int prior_k = 0; prior_k < k_idx; ++prior_k)
+            {
+                const int prior_winning_expert = indices[k * block_row + prior_k];
+
+                if (prior_winning_expert == expert)
+                {
+                    inp_kvp = thread_kvp;
+                }
+            }
+
+            thread_kvp = arg_max(inp_kvp, thread_kvp);
+        }
+
+        const cub_kvp result_kvp = BlockReduce(tmpStorage).Reduce(thread_kvp, arg_max);
+        if (threadIdx.x == 0)
+        {
+            // Ignore experts the node isn't responsible for with expert parallelism
+            const int expert = result_kvp.key;
+            const bool node_uses_expert = expert >= start_expert && expert < end_expert;
+            const bool should_process_row = row_is_active && node_uses_expert;
+
+            const int idx = k * block_row + k_idx;
+            output[idx] = result_kvp.value;
+            indices[idx] = should_process_row ? (expert - start_expert) : num_experts;
+            assert(indices[idx] >= 0);
+            source_rows[idx] = k_idx * num_rows + block_row;
+        }
+        __syncthreads();
+    }
+}
+
+// ====================== TopK softmax things ===============================
+
+/*
+  A Top-K gating softmax written to exploit when the number of experts in the MoE layers
+  are a small power of 2. This allows us to cleanly share the rows among the threads in
+  a single warp and eliminate communication between warps (so no need to use shared mem).
+
+  It fuses the softmax, max and argmax into a single kernel.
+
+  Limitations:
+  1) This implementation is intended for when the number of experts is a small power of 2.
+  2) This implementation assumes k is small, but will work for any k.
+*/
+
+template <int VPT, int NUM_EXPERTS, int WARPS_PER_CTA, int BYTES_PER_LDG>
+__launch_bounds__(WARPS_PER_CTA* WARP_SIZE) __global__
+    void topkGatingSoftmax(const float* input, const bool* finished, float* output, const int num_rows, int* indices,
+        int* source_rows, const int k, const int start_expert, const int end_expert)
+{
+    // We begin by enforcing compile time assertions and setting up compile time constants.
+    static_assert(VPT == (VPT & -VPT), "VPT must be power of 2");
+    static_assert(NUM_EXPERTS == (NUM_EXPERTS & -NUM_EXPERTS), "NUM_EXPERTS must be power of 2");
+    static_assert(BYTES_PER_LDG == (BYTES_PER_LDG & -BYTES_PER_LDG), "BYTES_PER_LDG must be power of 2");
+    static_assert(BYTES_PER_LDG <= 16, "BYTES_PER_LDG must be leq 16");
+
+    // Number of bytes each thread pulls in per load
+    static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(float);
+    static constexpr int ELTS_PER_ROW = NUM_EXPERTS;
+    static constexpr int THREADS_PER_ROW = ELTS_PER_ROW / VPT;
+    static constexpr int LDG_PER_THREAD = VPT / ELTS_PER_LDG;
+
+    // Restrictions based on previous section.
+    static_assert(VPT % ELTS_PER_LDG == 0, "The elements per thread must be a multiple of the elements per ldg");
+    static_assert(WARP_SIZE % THREADS_PER_ROW == 0, "The threads per row must cleanly divide the threads per warp");
+    static_assert(THREADS_PER_ROW == (THREADS_PER_ROW & -THREADS_PER_ROW), "THREADS_PER_ROW must be power of 2");
+    static_assert(THREADS_PER_ROW <= WARP_SIZE, "THREADS_PER_ROW can be at most warp size");
+
+    // We have NUM_EXPERTS elements per row. We specialize for small #experts
+    static constexpr int ELTS_PER_WARP = WARP_SIZE * VPT;
+    static constexpr int ROWS_PER_WARP = ELTS_PER_WARP / ELTS_PER_ROW;
+    static constexpr int ROWS_PER_CTA = WARPS_PER_CTA * ROWS_PER_WARP;
+
+    // Restrictions for previous section.
+    static_assert(ELTS_PER_WARP % ELTS_PER_ROW == 0, "The elts per row must cleanly divide the total elt per warp");
+
+    // ===================== From this point, we finally start computing run-time variables. ========================
+
+    // Compute CTA and warp rows. We pack multiple rows into a single warp, and a block contains WARPS_PER_CTA warps.
+    // This, each block processes a chunk of rows. We start by computing the start row for each block.
+    const int cta_base_row = blockIdx.x * ROWS_PER_CTA;
+
+    // Now, using the base row per thread block, we compute the base row per warp.
+    const int warp_base_row = cta_base_row + threadIdx.y * ROWS_PER_WARP;
+
+    // The threads in a warp are split into sub-groups that will work on a row.
+    // We compute row offset for each thread sub-group
+    const int thread_row_in_warp = threadIdx.x / THREADS_PER_ROW;
+    const int thread_row = warp_base_row + thread_row_in_warp;
+
+    // Threads with indices out of bounds should early exit here.
+    if (thread_row >= num_rows)
+    {
+        return;
+    }
+    const bool row_is_active = finished ? !finished[thread_row] : true;
+
+    // We finally start setting up the read pointers for each thread. First, each thread jumps to the start of the
+    // row it will read.
+    const float* thread_row_ptr = input + thread_row * ELTS_PER_ROW;
+
+    // Now, we compute the group each thread belong to in order to determine the first column to start loads.
+    const int thread_group_idx = threadIdx.x % THREADS_PER_ROW;
+    const int first_elt_read_by_thread = thread_group_idx * ELTS_PER_LDG;
+    const float* thread_read_ptr = thread_row_ptr + first_elt_read_by_thread;
+
+    // Determine the pointer type to use to read in the data depending on the BYTES_PER_LDG template param. In theory,
+    // this can support all powers of 2 up to 16.
+    // NOTE(woosuk): The original implementation uses CUTLASS aligned array here.
+    // We defined our own aligned array and use it here to avoid the dependency on CUTLASS.
+    using AccessType = AlignedArray<float, ELTS_PER_LDG>;
+
+    // Finally, we pull in the data from global mem
+    float row_chunk[VPT];
+    AccessType* row_chunk_vec_ptr = reinterpret_cast<AccessType*>(&row_chunk);
+    const AccessType* vec_thread_read_ptr = reinterpret_cast<const AccessType*>(thread_read_ptr);
+#pragma unroll
+    for (int ii = 0; ii < LDG_PER_THREAD; ++ii)
+    {
+        row_chunk_vec_ptr[ii] = vec_thread_read_ptr[ii * THREADS_PER_ROW];
+    }
+
+    // First, we perform a max reduce within the thread. We can do the max in fp16 safely (I think) and just
+    // convert to float afterwards for the exp + sum reduction.
+    float thread_max = row_chunk[0];
+#pragma unroll
+    for (int ii = 1; ii < VPT; ++ii)
+    {
+        thread_max = max(thread_max, row_chunk[ii]);
+    }
+
+// Now, we find the max within the thread group and distribute among the threads. We use a butterfly reduce.
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
+    {
+        thread_max = max(thread_max, VLLM_SHFL_XOR_SYNC_WIDTH(thread_max, mask, THREADS_PER_ROW));
+    }
+
+    // From this point, thread max in all the threads have the max within the row.
+    // Now, we subtract the max from each element in the thread and take the exp. We also compute the thread local sum.
+    float row_sum = 0;
+#pragma unroll
+    for (int ii = 0; ii < VPT; ++ii)
+    {
+        row_chunk[ii] = expf(row_chunk[ii] - thread_max);
+        row_sum += row_chunk[ii];
+    }
+
+// Now, we perform the sum reduce within each thread group. Similar to the max reduce, we use a bufferfly pattern.
+#pragma unroll
+    for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
+    {
+        row_sum += VLLM_SHFL_XOR_SYNC_WIDTH(row_sum, mask, THREADS_PER_ROW);
+    }
+
+    // From this point, all threads have the max and the sum for their rows in the thread_max and thread_sum variables
+    // respectively. Finally, we can scale the rows for the softmax. Technically, for top-k gating we don't need to
+    // compute the entire softmax row. We can likely look at the maxes and only compute for the top-k values in the row.
+    // However, this kernel will likely not be a bottle neck and it seems better to closer match torch and find the
+    // argmax after computing the softmax.
+    const float reciprocal_row_sum = 1.f / row_sum;
+
+#pragma unroll
+    for (int ii = 0; ii < VPT; ++ii)
+    {
+        row_chunk[ii] = row_chunk[ii] * reciprocal_row_sum;
+    }
+
+    // Now, softmax_res contains the softmax of the row chunk. Now, I want to find the topk elements in each row, along
+    // with the max index.
+    int start_col = first_elt_read_by_thread;
+    static constexpr int COLS_PER_GROUP_LDG = ELTS_PER_LDG * THREADS_PER_ROW;
+
+    for (int k_idx = 0; k_idx < k; ++k_idx)
+    {
+        // First, each thread does the local argmax
+        float max_val = row_chunk[0];
+        int expert = start_col;
+#pragma unroll
+        for (int ldg = 0, col = start_col; ldg < LDG_PER_THREAD; ++ldg, col += COLS_PER_GROUP_LDG)
+        {
+#pragma unroll
+            for (int ii = 0; ii < ELTS_PER_LDG; ++ii)
+            {
+                float val = row_chunk[ldg * ELTS_PER_LDG + ii];
+
+                // No check on the experts here since columns with the smallest index are processed first and only
+                // updated if > (not >=)
+                if (val > max_val)
+                {
+                    max_val = val;
+                    expert = col + ii;
+                }
+            }
+        }
+
+// Now, we perform the argmax reduce. We use the butterfly pattern so threads reach consensus about the max.
+// This will be useful for K > 1 so that the threads can agree on "who" had the max value. That thread can
+// then blank out their max with -inf and the warp can run more iterations...
+#pragma unroll
+        for (int mask = THREADS_PER_ROW / 2; mask > 0; mask /= 2)
+        {
+            float other_max = VLLM_SHFL_XOR_SYNC_WIDTH(max_val, mask, THREADS_PER_ROW);
+            int other_expert = VLLM_SHFL_XOR_SYNC_WIDTH(expert, mask, THREADS_PER_ROW);
+
+            // We want lower indices to "win" in every thread so we break ties this way
+            if (other_max > max_val || (other_max == max_val && other_expert < expert))
+            {
+                max_val = other_max;
+                expert = other_expert;
+            }
+        }
+
+        // Write the max for this k iteration to global memory.
+        if (thread_group_idx == 0)
+        {
+            // Add a guard to ignore experts not included by this node
+            const bool node_uses_expert = expert >= start_expert && expert < end_expert;
+            const bool should_process_row = row_is_active && node_uses_expert;
+
+            // The lead thread from each sub-group will write out the final results to global memory. (This will be a
+            // single) thread per row of the input/output matrices.
+            const int idx = k * thread_row + k_idx;
+            output[idx] = max_val;
+            indices[idx] = should_process_row ? (expert - start_expert) : NUM_EXPERTS;
+            source_rows[idx] = k_idx * num_rows + thread_row;
+        }
+
+        // Finally, we clear the value in the thread with the current max if there is another iteration to run.
+        if (k_idx + 1 < k)
+        {
+            const int ldg_group_for_expert = expert / COLS_PER_GROUP_LDG;
+            const int thread_to_clear_in_group = (expert / ELTS_PER_LDG) % THREADS_PER_ROW;
+
+            // Only the thread in the group which produced the max will reset the "winning" value to -inf.
+            if (thread_group_idx == thread_to_clear_in_group)
+            {
+                const int offset_for_expert = expert % ELTS_PER_LDG;
+                // Safe to set to any negative value since row_chunk values must be between 0 and 1.
+                row_chunk[ldg_group_for_expert * ELTS_PER_LDG + offset_for_expert] = -10000.f;
+            }
+        }
+    }
+}
+
+namespace detail
+{
+// Constructs some constants needed to partition the work across threads at compile time.
+template <int EXPERTS, int BYTES_PER_LDG>
+struct TopkConstants
+{
+    static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(float);
+    static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 || EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0, "");
+    static constexpr int VECs_PER_THREAD = MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
+    static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
+    static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
+    static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
+};
+} // namespace detail
+
+template <int EXPERTS, int WARPS_PER_TB>
+void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, float* output, int* indices,
+    int* source_row, const int num_rows, const int k, const int start_expert, const int end_expert, cudaStream_t stream)
+{
+    static constexpr std::size_t MAX_BYTES_PER_LDG = 16;
+
+    static constexpr int BYTES_PER_LDG = MIN(MAX_BYTES_PER_LDG, sizeof(float) * EXPERTS);
+    using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG>;
+    static constexpr int VPT = Constants::VPT;
+    static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;
+    const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;
+    const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;
+
+    dim3 block_dim(WARP_SIZE, WARPS_PER_TB);
+    topkGatingSoftmax<VPT, EXPERTS, WARPS_PER_TB, BYTES_PER_LDG><<<num_blocks, block_dim, 0, stream>>>(
+        input, finished, output, num_rows, indices, source_row, k, start_expert, end_expert);
+}
+
+#define LAUNCH_SOFTMAX(NUM_EXPERTS, WARPS_PER_TB)                       \
+    topkGatingSoftmaxLauncherHelper<NUM_EXPERTS, WARPS_PER_TB>(         \
+        gating_output, nullptr, topk_weights, topk_indicies,            \
+        token_expert_indices, num_tokens, topk, 0, num_experts,         \
+        stream);
+
+void topkGatingSoftmaxKernelLauncher(
+    const float* gating_output,
+    float* topk_weights,
+    int* topk_indicies,
+    int* token_expert_indices,
+    float* softmax_workspace,
+    const int num_tokens,
+    const int num_experts,
+    const int topk,
+    cudaStream_t stream) {
+    static constexpr int WARPS_PER_TB = 4;
+    switch (num_experts) {
+        case 1:
+            LAUNCH_SOFTMAX(1, WARPS_PER_TB);
+            break;
+        case 2:
+            LAUNCH_SOFTMAX(2, WARPS_PER_TB);
+            break;
+        case 4:
+            LAUNCH_SOFTMAX(4, WARPS_PER_TB);
+            break;
+        case 8:
+            LAUNCH_SOFTMAX(8, WARPS_PER_TB);
+            break;
+        case 16:
+            LAUNCH_SOFTMAX(16, WARPS_PER_TB);
+            break;
+        case 32:
+            LAUNCH_SOFTMAX(32, WARPS_PER_TB);
+            break;
+        case 64:
+            LAUNCH_SOFTMAX(64, WARPS_PER_TB);
+            break;
+        case 128:
+            LAUNCH_SOFTMAX(128, WARPS_PER_TB);
+            break;
+        case 256:
+            LAUNCH_SOFTMAX(256, WARPS_PER_TB);
+            break;
+        default: {
+            TORCH_CHECK(softmax_workspace != nullptr,
+                "softmax_workspace must be provided for num_experts that are not a power of 2.");
+            static constexpr int TPB = 256;
+            moeSoftmax<TPB><<<num_tokens, TPB, 0, stream>>>(
+                gating_output, nullptr, softmax_workspace, num_experts);
+            moeTopK<TPB><<<num_tokens, TPB, 0, stream>>>(
+                softmax_workspace, nullptr, topk_weights, topk_indicies, token_expert_indices,
+                num_experts, topk, 0, num_experts);
+        }
+    }
+}
+
+} // namespace moe
+} // namespace vllm_musa
+
+void topk_softmax(
+    torch::Tensor& topk_weights,                // [num_tokens, topk]
+    torch::Tensor& topk_indices,                // [num_tokens, topk]
+    torch::Tensor& token_expert_indices,        // [num_tokens, topk]
+    torch::Tensor& gating_output)               // [num_tokens, num_experts]
+{
+    const int num_experts = gating_output.size(-1);
+    const int num_tokens = gating_output.numel() / num_experts;
+    const int topk = topk_weights.size(-1);
+
+    const bool is_pow_2 = (num_experts != 0) && ((num_experts & (num_experts - 1)) == 0);
+    const bool needs_workspace = !is_pow_2 || num_experts > 256;
+    const int64_t workspace_size = needs_workspace ? num_tokens * num_experts : 0;
+
+    const at::cuda::OptionalCUDAGuard device_guard(device_of(gating_output));
+    const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    torch::Tensor softmax_workspace = torch::empty({workspace_size}, gating_output.options());
+    vllm_musa::moe::topkGatingSoftmaxKernelLauncher(
+        gating_output.data_ptr<float>(),
+        topk_weights.data_ptr<float>(),
+        topk_indices.data_ptr<int>(),
+        token_expert_indices.data_ptr<int>(),
+        softmax_workspace.data_ptr<float>(),
+        num_tokens,
+        num_experts,
+        topk,
+        stream);
+}
diff --git a/vllm_musa/csrc/musa.h b/vllm_musa/csrc/musa.h
new file mode 100644
index 000000000..927a02bb3
--- /dev/null
+++ b/vllm_musa/csrc/musa.h
@@ -0,0 +1,176 @@
+#pragma once
+
+#include <musa_runtime.h>
+#include <musa.h>
+#include <murand_kernel.h>
+#include <mublas.h>
+#include <musa_bf16.h>
+#include <musa_fp16.h>
+#include <musa_fp8.h>
+#include <cub/cub.cuh>
+#include "torch_musa/csrc/aten/musa/MUSAContext.h"
+#include "torch_musa/csrc/core/MUSAGuard.h"
+
+using __nv_bfloat16 = __mt_bfloat16;
+using __nv_bfloat162 = __mt_bfloat162;
+
+using nv_bfloat16 = __mt_bfloat16;
+using nv_bfloat162 = __mt_bfloat162;
+
+using __nv_fp8_e4m3 = __mt_fp8_e4m3;
+using __nv_fp8_e5m2 = __mt_fp8_e5m2;
+using __nv_fp8x2_e4m3 = __mt_fp8x2_e4m3;
+using __nv_fp8x4_e4m3 = __mt_fp8x4_e4m3;
+using __nv_fp8x2_e5m2 = __mt_fp8x2_e5m2;
+using __nv_fp8x4_e5m2 = __mt_fp8x4_e5m2;
+using __nv_fp8x2_storage_t = __mt_fp8x2_storage_t;
+using __nv_fp8x4_storage_t = __mt_fp8x4_storage_t;
+
+#define cudaDevAttrMultiProcessorCount musaDevAttrMultiProcessorCount
+#define cudaDeviceGetAttribute musaDeviceGetAttribute
+#define cudaDevAttrMaxThreadsPerBlock musaDevAttrMaxThreadsPerBlock
+#define AT_CUDA_CHECK AT_MUSA_CHECK
+#define cuda musa
+#define getCurrentCUDAStream getCurrentMUSAStream
+#define OptionalCUDAGuard OptionalMUSAGuard
+#define CUBLAS_COMPUTE_16F CUDA_R_16F
+#define CUBLAS_COMPUTE_32F CUDA_R_32F
+#define CUBLAS_COMPUTE_32F_FAST_16F MUBLAS_COMPUTE_32F_FAST_16F
+#define CUBLAS_GEMM_DEFAULT MUBLAS_GEMM_DEFAULT
+#define CUBLAS_GEMM_DEFAULT_TENSOR_OP MUBLAS_GEMM_DEFAULT
+#define CUBLAS_OP_N MUBLAS_OP_N
+#define CUBLAS_OP_T MUBLAS_OP_T
+#define CUBLAS_STATUS_SUCCESS MUBLAS_STATUS_SUCCESS
+#define CUBLAS_TF32_TENSOR_OP_MATH MUBLAS_MATH_MODE_DEFAULT
+#define CUDA_R_16F  MUSA_R_16F
+#define CUDA_R_32F  MUSA_R_32F
+#define cublasComputeType_t cudaDataType_t
+#define cublasCreate mublasCreate
+#define cublasDestroy mublasDestroy
+#define cublasGemmEx mublasGemmEx
+#define cublasGemmBatchedEx mublasGemmBatchedEx
+#define cublasGemmStridedBatchedEx mublasGemmStridedBatchedEx
+#define cublasHandle_t mublasHandle_t
+#define cublasSetMathMode mublasSetMathMode
+#define cublasSetStream mublasSetStream
+#define cublasSgemm mublasSgemm
+#define cublasStatus_t mublasStatus_t
+#define cublasOperation_t mublasOperation_t
+#define cublasGetStatusString mublasStatus_to_string
+#define cudaDataType_t musaDataType_t
+#define cudaDeviceCanAccessPeer musaDeviceCanAccessPeer
+#define cudaDeviceDisablePeerAccess musaDeviceDisablePeerAccess
+#define cudaDeviceEnablePeerAccess musaDeviceEnablePeerAccess
+#define cudaDeviceProp musaDeviceProp
+#define cudaDeviceSynchronize musaDeviceSynchronize
+#define cudaError_t musaError_t
+#define cudaErrorPeerAccessAlreadyEnabled musaErrorPeerAccessAlreadyEnabled
+#define cudaErrorPeerAccessNotEnabled musaErrorPeerAccessNotEnabled
+#define cudaEventCreateWithFlags musaEventCreateWithFlags
+#define cudaEventDisableTiming musaEventDisableTiming
+#define cudaEventRecord musaEventRecord
+#define cudaEventSynchronize musaEventSynchronize
+#define cudaEvent_t musaEvent_t
+#define cudaEventDestroy musaEventDestroy
+#define cudaFree musaFree
+#define cudaFreeHost musaFreeHost
+#define cudaGetDevice musaGetDevice
+#define cudaGetDeviceCount musaGetDeviceCount
+#define cudaGetDeviceProperties musaGetDeviceProperties
+#define cudaGetErrorString musaGetErrorString
+#define cudaGetLastError musaGetLastError
+#define cudaHostRegister musaHostRegister
+#define cudaHostRegisterPortable musaHostRegisterPortable
+#define cudaHostRegisterReadOnly musaHostRegisterReadOnly
+#define cudaHostUnregister musaHostUnregister
+#define cudaLaunchHostFunc musaLaunchHostFunc
+#define cudaMalloc musaMalloc
+#define cudaMallocHost musaMallocHost
+#define cudaMallocManaged musaMallocManaged
+#define cudaMemcpy musaMemcpy
+#define cudaMemcpyAsync musaMemcpyAsync
+#define cudaMemcpyPeerAsync musaMemcpyPeerAsync
+#define cudaMemcpy2DAsync musaMemcpy2DAsync
+#define cudaMemcpyDeviceToDevice musaMemcpyDeviceToDevice
+#define cudaMemcpyDeviceToHost musaMemcpyDeviceToHost
+#define cudaMemcpyHostToDevice musaMemcpyHostToDevice
+#define cudaMemcpyKind musaMemcpyKind
+#define cudaMemset musaMemset
+#define cudaMemsetAsync musaMemsetAsync
+#define cudaMemGetInfo musaMemGetInfo
+#define cudaOccupancyMaxPotentialBlockSize musaOccupancyMaxPotentialBlockSize
+#define cudaSetDevice musaSetDevice
+#define cudaStreamCreateWithFlags musaStreamCreateWithFlags
+#define cudaStreamDestroy musaStreamDestroy
+#define cudaStreamFireAndForget musaStreamFireAndForget
+#define cudaStreamNonBlocking musaStreamNonBlocking
+#define cudaStreamPerThread musaStreamPerThread
+#define cudaStreamSynchronize musaStreamSynchronize
+#define cudaStreamWaitEvent musaStreamWaitEvent
+#define cudaStream_t musaStream_t
+#define cudaSuccess musaSuccess
+
+// Additional mappings for MUSA virtual memory pool
+#define CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED MU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED
+#define CU_MEM_ACCESS_FLAGS_PROT_READWRITE MU_MEM_ACCESS_FLAGS_PROT_READWRITE
+#define CU_MEM_ALLOC_GRANULARITY_RECOMMENDED MU_MEM_ALLOC_GRANULARITY_RECOMMENDED
+#define CU_MEM_ALLOCATION_TYPE_PINNED MU_MEM_ALLOCATION_TYPE_PINNED
+#define CU_MEM_LOCATION_TYPE_DEVICE MU_MEM_LOCATION_TYPE_DEVICE
+#define CUdevice MUdevice
+#define CUdeviceptr MUdeviceptr
+#define CUmemAccessDesc MUmemAccessDesc
+#define CUmemAllocationProp MUmemAllocationProp
+#define CUmemGenericAllocationHandle MUmemGenericAllocationHandle
+#define cuDeviceGet muDeviceGet
+#define cuDeviceGetAttribute muDeviceGetAttribute
+#define cuMemAddressFree muMemAddressFree
+#define cuMemAddressReserve muMemAddressReserve
+#define cuMemCreate muMemCreate
+#define cuMemGetAllocationGranularity muMemGetAllocationGranularity
+#define cuMemMap muMemMap
+#define cuMemRelease muMemRelease
+#define cuMemSetAccess muMemSetAccess
+#define cuMemUnmap muMemUnmap
+#define cudaFuncAttributeMaxDynamicSharedMemorySize musaFuncAttributeMaxDynamicSharedMemorySize
+#define cudaFuncSetAttribute musaFuncSetAttribute
+#define cudaMemcpy3DPeerParms musaMemcpy3DPeerParms
+#define make_cudaExtent make_musaExtent
+#define make_cudaPitchedPtr make_musaPitchedPtr
+
+// Additional mappings for MUSA graphs
+#define CUDA_SUCCESS MUSA_SUCCESS
+#define CUresult MUresult
+#define cuGetErrorString muGetErrorString
+#define cudaErrorGraphExecUpdateFailure musaErrorGraphExecUpdateFailure
+#define cudaErrorInvalidDeviceFunction musaErrorInvalidDeviceFunction
+#define cudaGraphDestroy musaGraphDestroy
+#define cudaGraphExecDestroy musaGraphExecDestroy
+#define cudaGraphExec_t musaGraphExec_t
+#define cudaGraphExecUpdate musaGraphExecUpdate
+#define cudaGraphExecUpdateResultInfo musaGraphExecUpdateResult
+#define cudaGraphGetNodes musaGraphGetNodes
+#define cudaGraphInstantiate musaGraphInstantiate
+#define cudaGraphKernelNodeGetParams musaGraphKernelNodeGetParams
+#define cudaGraphKernelNodeSetParams musaGraphKernelNodeSetParams
+#define cudaGraphLaunch musaGraphLaunch
+#define cudaGraphNodeGetType musaGraphNodeGetType
+#define cudaGraphNode_t musaGraphNode_t
+#define cudaGraphNodeType musaGraphNodeType
+#define cudaGraphNodeTypeKernel musaGraphNodeTypeKernel
+#define cudaGraph_t musaGraph_t
+#define cudaKernelNodeParams musaKernelNodeParams
+#define cudaStreamCaptureModeRelaxed musaStreamCaptureModeRelaxed
+#define cudaStreamEndCapture musaStreamEndCapture
+
+// Additional mappings for custom allreduce, added by qn: NOT used now.
+#define cudaIpcMemHandle_t musaIpcMemHandle_t
+#define cudaIpcMemLazyEnablePeerAccess musaIpcMemLazyEnablePeerAccess
+#define cudaStreamIsCapturing at::musa::musaStreamIsCapturing
+// NotImplementedError: musaStreamIsCapturing is not supported now!
+#define cudaStreamCaptureStatusActive musaStreamCaptureStatusActive
+#define cudaIpcCloseMemHandle musaIpcCloseMemHandle
+#define cudaIpcOpenMemHandle musaIpcOpenMemHandle
+#define cudaIpcGetMemHandle musaIpcGetMemHandle
+#define cudaStreamCaptureStatus musaStreamCaptureStatus
+#define CU_POINTER_ATTRIBUTE_RANGE_START_ADDR MU_POINTER_ATTRIBUTE_RANGE_START_ADDR
+#define cuPointerGetAttribute muPointerGetAttribute
diff --git a/vllm_musa/csrc/ops.mu b/vllm_musa/csrc/ops.mu
new file mode 100644
index 000000000..ca443cbdf
--- /dev/null
+++ b/vllm_musa/csrc/ops.mu
@@ -0,0 +1,497 @@
+#include <optional>
+#include <torch/all.h>
+#include <torch/extension.h>
+#include <torch/library.h>
+
+#include <map>
+#include <vector>
+
+torch::Tensor weak_ref_tensor(torch::Tensor &tensor)
+{
+    if (!tensor.is_privateuseone())
+    {
+        throw std::runtime_error("Tensor must be on MUSA device");
+    }
+
+    void *data_ptr = tensor.data_ptr();
+    std::vector<int64_t> sizes = tensor.sizes().vec();
+    std::vector<int64_t> strides = tensor.strides().vec();
+    auto options = tensor.options();
+
+    auto new_tensor = torch::from_blob(data_ptr, sizes, strides, options);
+    return new_tensor;
+}
+
+// layer norm ops
+void rms_norm(torch::Tensor &out, torch::Tensor &input, torch::Tensor &weight,
+              double epsilon);
+
+void fused_add_rms_norm(torch::Tensor &input, torch::Tensor &residual,
+                        torch::Tensor &weight, double epsilon);
+
+void musa_fused_add_rms_norm(torch::Tensor &input, torch::Tensor &residual,
+                            torch::Tensor &weight,double epsilon);
+
+void rmsnorm_cast(
+    at::Tensor &output_fp8, at::Tensor &inv_scale,
+    const at::Tensor &input, const at::Tensor &gamma,
+    const float eps, const int quant_size);
+
+void rmsnorm_add_cast(
+    torch::Tensor &output_fp8, torch::Tensor &inv_scale,
+    const torch::Tensor &input, const torch::Tensor &resdual, const torch::Tensor &gamma,
+    const float eps, const int quant_size);
+
+// layer rotary_embedding ops
+void rotary_embedding(torch::Tensor &positions, torch::Tensor &query,
+                      torch::Tensor &key, int64_t head_size,
+                      torch::Tensor &cos_sin_cache, bool is_neox);
+
+void batched_rotary_embedding(torch::Tensor &positions, torch::Tensor &query,
+                              torch::Tensor &key, int64_t head_size,
+                              torch::Tensor &cos_sin_cache, bool is_neox,
+                              int64_t rot_dim,
+                              torch::Tensor &cos_sin_cache_offsets);
+void rotary_embedding_contiguous(
+    torch::Tensor &positions, // [num_tokens]
+    torch::Tensor &query,     // [num_tokens, num_heads, head_size]
+    torch::Tensor &key,       // [num_tokens, num_kv_heads, head_size]
+    int64_t head_size,
+    torch::Tensor &cos_sin_cache, // [max_position, rot_dim]
+    bool is_neox);
+
+void batched_rotary_embedding_contiguous(
+    torch::Tensor &positions, // [num_tokens]
+    torch::Tensor &query,     // [num_tokens, num_heads, head_size]
+    torch::Tensor &key,       // [num_tokens, num_kv_heads, head_size]
+    int64_t head_size,
+    torch::Tensor &cos_sin_cache, // [max_position, rot_dim]
+    bool is_neox, int64_t rot_dim,
+    torch::Tensor &cos_sin_cache_offsets // [num_tokens]
+);
+
+// cache
+void reshape_and_cache(torch::Tensor &key, torch::Tensor &value,
+                       torch::Tensor &key_cache, torch::Tensor &value_cache,
+                       torch::Tensor &slot_mapping,
+                       const std::string &kv_cache_dtype,
+                       torch::Tensor &k_scale, torch::Tensor &v_scale);
+
+void reshape_and_cache_flash(torch::Tensor &key, torch::Tensor &value,
+                             torch::Tensor &key_cache,
+                             torch::Tensor &value_cache,
+                             torch::Tensor &slot_mapping,
+                             const std::string &kv_cache_dtype,
+                             torch::Tensor &k_scale, torch::Tensor &v_scale);
+
+void swap_blocks(torch::Tensor &src, torch::Tensor &dst,
+                 const torch::Tensor &block_mapping);
+
+// Note: the key_caches and value_caches vectors are constant but
+// not the Tensors they contain. The vectors need to be const refs
+// in order to satisfy pytorch's C++ operator registration code.
+void copy_blocks(std::vector<torch::Tensor> const &key_caches,
+                 std::vector<torch::Tensor> const &value_caches,
+                 const torch::Tensor &block_mapping);
+
+void copy_blocks_mla(std::vector<torch::Tensor> const &kv_caches,
+                     const torch::Tensor &block_mapping);
+
+void concat_and_cache_mla(
+    torch::Tensor &kv_c,         // [num_tokens, kv_lora_rank]
+    torch::Tensor &k_pe,         // [num_tokens, pe_dim]
+    torch::Tensor &kv_cache,     // [num_blocks, block_size, (kv_lora_rank +
+                                 // pe_dim)]
+    torch::Tensor &slot_mapping, // [num_tokens] or [num_actual_tokens]
+    const std::string &kv_cache_dtype, torch::Tensor &scale);
+
+void gather_cache(
+    torch::Tensor const &src_cache,   // [NUM_BLOCKS, BLOCK_SIZE, ENTRIES...]
+    torch::Tensor const &dst,         // [TOT_TOKENS, ENTRIES...]
+    torch::Tensor const &block_table, // [BATCH, BLOCK_INDICES]
+    torch::Tensor const &cu_seq_lens, // [BATCH+1]
+    int64_t batch_size, std::optional<torch::Tensor> seq_starts = std::nullopt);
+
+// paged attention v1
+void paged_attention_v1(
+    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 &seq_lens, int64_t block_size,
+    int64_t max_seq_len, const std::optional<torch::Tensor> &alibi_slopes,
+    const std::string &kv_cache_dtype, torch::Tensor &k_scale,
+    torch::Tensor &v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step);
+
+void paged_attention_v2(
+    torch::Tensor &out,        // [num_seqs, num_heads, head_size]
+    torch::Tensor &exp_sums,   // [num_seqs, num_heads, max_num_partitions]
+    torch::Tensor &max_logits, // [num_seqs, num_heads, max_num_partitions]
+    torch::Tensor
+        &tmp_out,         // [num_seqs, num_heads, max_num_partitions, head_size]
+    torch::Tensor &query, // [num_seqs, num_heads, head_size]
+    torch::Tensor
+        &key_cache, // [num_blocks, num_heads, head_size/x, block_size, x]
+    torch::Tensor
+        &value_cache,     // [num_blocks, num_heads, head_size, block_size]
+    int64_t num_kv_heads, // [num_heads]
+    double scale,
+    torch::Tensor &block_tables, // [num_seqs, max_num_blocks_per_seq]
+    torch::Tensor &seq_lens,     // [num_seqs]
+    int64_t block_size, int64_t max_seq_len,
+    const std::optional<torch::Tensor> &alibi_slopes,
+    const std::string &kv_cache_dtype, torch::Tensor &k_scale,
+    torch::Tensor &v_scale, const int64_t tp_rank,
+    const int64_t blocksparse_local_blocks,
+    const int64_t blocksparse_vert_stride, const int64_t blocksparse_block_size,
+    const int64_t blocksparse_head_sliding_step);
+
+void merge_attn_states(torch::Tensor &output,
+                       std::optional<torch::Tensor> output_lse,
+                       const torch::Tensor &prefix_output,
+                       const torch::Tensor &prefix_lse,
+                       const torch::Tensor &suffix_output,
+                       const torch::Tensor &suffix_lse);
+
+// moe
+void topk_softmax(torch::Tensor &topk_weights, torch::Tensor &topk_indices,
+                  torch::Tensor &token_expert_indices,
+                  torch::Tensor &gating_output);
+
+void moe_sum(torch::Tensor &input, torch::Tensor &output);
+
+void moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
+                          int64_t block_size, torch::Tensor sorted_token_ids,
+                          torch::Tensor experts_ids,
+                          torch::Tensor num_tokens_post_pad);
+
+void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
+                              int64_t block_size,
+                              torch::Tensor sorted_token_ids,
+                              torch::Tensor experts_ids,
+                              torch::Tensor num_tokens_post_pad);
+
+std::tuple<torch::Tensor, torch::Tensor> mp31_fused_grouped_topk(
+    const torch::Tensor &gating_output,
+    const c10::optional<torch::Tensor> &e_score_correction_bias,
+    int num_expert_group, int topk_group, int topk, bool renormalize);
+
+
+void mp22_fused_grouped_topk(
+    const torch::Tensor& topk_weights,
+    const torch::Tensor& topk_indices,
+    const torch::Tensor& gating_output,
+    const c10::optional<torch::Tensor>& e_score_correction_bias,
+    int num_expert_group,
+    int topk_group,
+    int topk,
+    bool renormalize);
+
+// activation
+void silu_and_mul(torch::Tensor &out,    // [..., d]
+                  torch::Tensor &input); // [..., 2 * d]
+
+void _fused_mul_add(torch::Tensor &output, torch::Tensor &self,
+                    torch::Tensor &bias, const float scale);
+
+void gemm_swiglu(torch::Tensor &output, // [batch, hidden_size/2]
+                 torch::Tensor &input,  // [batch, hidden_size]
+                 torch::Tensor &weight, // [hidden_size, hidden_size/2]
+                 double alpha);
+
+void flash_attn_fwd(
+    const torch::Tensor &Q,            // verlen ? [total_seqlen, num_head, qk_head_dim] : [batch, num_head, seqlen, qk_head_dim]
+    const torch::Tensor &K,            // verlen ? [total_seqlen, kv_num_head, qk_head_dim] : [batch, kv_num_head, seqlen, qk_head_dim]
+    const torch::Tensor &V,            // verlen ? [total_seqlen, kv_num_head, vo_head_dim] : [batch, kv_num_head, seqlen, vo_head_dim]
+    torch::Tensor &O,                  // [total_seqlen, num_head, d_vo]
+    c10::optional<torch::Tensor> &lse, // [num_head, total_seqlen]
+    const torch::Tensor &cu_q_seqlen,
+    const torch::Tensor &cu_kv_seqlen,
+    int max_seqlen_q,
+    int max_seqlen_kv,
+    bool causal,
+    bool varlen,
+    bool return_softmax_lse);
+
+void mp22_flash_attn_varlen_func(torch::Tensor &Q,                  // [seqs_len, num_heads, head_size]
+                                 torch::Tensor &K,                  // [num_block, page_size, num_heads, head_size]
+                                 torch::Tensor &V,                  // [num_block, page_size, num_heads, head_size]
+                                 torch::Tensor &out,                // [seqs_len, num_heads, head_size]
+                                 c10::optional<torch::Tensor> &lse, // [ num_heads, seqs_len]
+                                 torch::Tensor &cu_q_seqlen,        // such as [0, 19, 44]
+                                 torch::Tensor &cu_kv_seqlen,
+                                 torch::Tensor &block_table,
+                                 float sm_scale,
+                                 int max_seqlen_q,
+                                 int max_seqlen_kv,
+                                 bool causal,
+                                 bool return_softmax_lse);
+
+void gemm_swiglu_w4a16(torch::Tensor &output,
+                       torch::Tensor &input,
+                       torch::Tensor &weight,
+                       torch::Tensor &scale,
+                       double alpha);
+
+void gemm_w4a16(torch::Tensor &output,
+                torch::Tensor &input,
+                torch::Tensor &weight,
+                torch::Tensor &scale,
+                double alpha);
+
+void mudnn_w4a16_matmul(torch::Tensor out, torch::Tensor in,
+                        torch::Tensor weight, torch::Tensor scales,
+                        std::optional<torch::Tensor> bias,
+                        const bool split_k = false);
+
+void mudnn_w8a8_scaled_mm(torch::Tensor &c, torch::Tensor const &a,
+                          torch::Tensor const &b, torch::Tensor const &a_scales,
+                          torch::Tensor const &b_scales,
+                          std::optional<torch::Tensor> const &bias);
+
+void fused_gemv(
+    torch::Tensor &A,                            // [bseqlen, in]
+    torch::Tensor &B,                            // [out, in/8]
+    torch::Tensor &C,                            // [bseqlen, out]
+    const c10::optional<torch::Tensor> &A_scale, // [out, in/group_size]
+    const c10::optional<torch::Tensor> &B_scale, // [out, in/group_size]
+    bool use_int4_w4a16,
+    bool use_swigelu,
+    bool use_rms_norm,
+    const c10::optional<torch::Tensor> &gamma,
+    float eps);
+
+void group_scale_gemm(
+    torch::Tensor &A,
+    torch::Tensor &B,
+    torch::Tensor &C,
+    torch::Tensor &B_scale);
+
+// custom ar(all-reduce)
+using fptr_t = int64_t;
+fptr_t init_custom_ar(const std::vector<int64_t> &fake_ipc_ptrs,
+                      torch::Tensor &rank_data, int64_t rank, bool full_nvlink);
+void all_reduce(fptr_t _fa, torch::Tensor &inp, torch::Tensor &out,
+                fptr_t reg_buffer, int64_t reg_buffer_sz_bytes);
+void dispose(fptr_t _fa);
+int64_t meta_size();
+void register_buffer(fptr_t _fa, const std::vector<int64_t> &fake_ipc_ptrs);
+std::tuple<std::vector<int64_t>, std::vector<int64_t>> get_graph_buffer_ipc_meta(fptr_t _fa);
+void register_graph_buffers(fptr_t _fa,
+                            const std::vector<std::vector<int64_t>> &handles,
+                            const std::vector<std::vector<int64_t>> &offsets);
+
+void decode_mla(
+    torch::Tensor &Q, // [batch*seqlen, num_head, head_dim]
+    torch::Tensor &K_Buffer,
+    torch::Tensor &V_Buffer,
+    float sm_scale,
+    int num_splits_kv,
+    torch::Tensor &Req_to_tokens,
+    torch::Tensor &B_Seqlen,
+    torch::Tensor &Att_out);
+
+void mp31_decode_mha(
+    torch::Tensor &Q,        // [batch, num_head, head_dim]
+    torch::Tensor &K_Buffer, // [num_blocks, page_size, h_k, d_k]
+    torch::Tensor &V_Buffer, // [num_blocks, page_size, h_v, d_v]
+    float sm_scale,
+    int num_splits_kv,
+    torch::Tensor &Req_to_tokens,
+    torch::Tensor &B_Seqlen,
+    torch::Tensor &Att_out);
+
+void fused_moe_mutlass(
+    torch::Tensor &A,                            // [32768, 2048]    [bseqlen, hidden_size]
+    torch::Tensor &B,                            // [64, 2816, 256]  []
+    torch::Tensor &C,                            // [32768, 6, 2846]
+    const c10::optional<torch::Tensor> &A_scale, // use for fp8
+    const c10::optional<torch::Tensor> &B_scale, // use for fp8/int
+    torch::Tensor &topk_weights,
+    torch::Tensor &topk_ids,
+    torch::Tensor &sorted_token_ids,
+    torch::Tensor &expert_ids,
+    torch::Tensor &num_tokens_post_padded,
+    bool mul_routed_weight,
+    int topk,
+    bool use_int4_w4a16,
+    int BLOCK_SIZE_M,
+    int BLOCK_SIZE_N,
+    int BLOCK_SIZE_K,
+    int GROUP_SIZE_M,
+    int group_size);
+
+void fused_moe_gemv(
+    torch::Tensor &A,                            // [32768, 2048]    [bseqlen, hidden_size]
+    torch::Tensor &B,                            // [64, 2816, 256]  []
+    torch::Tensor &C,                            // [32768, 6, 2846]
+    const c10::optional<torch::Tensor> &A_scale, // use for fp8
+    const c10::optional<torch::Tensor> &B_scale, // use for fp8/int
+    torch::Tensor &topk_weights,
+    torch::Tensor &topk_ids,
+    torch::Tensor &sorted_token_ids,
+    torch::Tensor &expert_ids,
+    torch::Tensor &num_tokens_post_padded,
+    bool mul_routed_weight,
+    int topk,
+    bool use_int4_w4a16,
+    bool use_swigelu);
+
+void varlen_fa_seqlen_pad(
+    torch::Tensor &q_in,
+    torch::Tensor &k_in,
+    torch::Tensor &v_in,
+    torch::Tensor &q_out,
+    torch::Tensor &k_out,
+    torch::Tensor &v_out,
+    torch::Tensor &accum_lens_q,
+    torch::Tensor &accum_lens_k,
+    int sum_seq,
+    int max_seq,
+    int batch_size);
+
+void varlen_fa_seqlen_unpad(
+    torch::Tensor &in,  // [bs, num_heads, max_prefill_seq_len, head_size]
+    torch::Tensor &out, // [sum_seq, num_heads, head_size]
+    torch::Tensor &accum_lens,
+    int sum_seq,
+    int max_seq,
+    int head_dim,
+    int num_heads,
+    int batch_size);
+
+void varlen_fa_lse_seqlen_unpad(
+    torch::Tensor &in,         //  [batch_size, num_heads, max_seq]
+    torch::Tensor &out,        //  [num_heads, sum_seq]
+    torch::Tensor &accum_lens, //  [seq_0, seq_1,...]
+    int sum_seq,               // (seq_lens)）
+    int max_seq,
+    int num_heads,
+    int batch_size);
+
+void static_scaled_fp8_quant(torch::Tensor &out,          // [..., d]
+                             torch::Tensor const &input,  // [..., d]
+                             torch::Tensor const &scale); // [1]
+
+void dynamic_scaled_fp8_quant(torch::Tensor &out,         // [..., d]
+                              torch::Tensor const &input, // [..., d]
+                              torch::Tensor &scale);      // [1]
+
+void dynamic_per_token_scaled_fp8_quant(
+    torch::Tensor &out,         // [..., d]
+    torch::Tensor const &input, // [..., d]
+    torch::Tensor &scales, std::optional<at::Tensor> const &scale_ub);
+
+void musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
+    const std::tuple<torch::Tensor, torch::Tensor> &lhs,
+    const std::tuple<torch::Tensor, torch::Tensor> &rhs,
+    torch::Tensor &out, const torch::Tensor &m_indices, const int block_m);
+
+void sgl_per_token_group_quant_fp8(
+    at::Tensor input,
+    at::Tensor output_q,
+    at::Tensor output_s,
+    int64_t group_size,
+    double eps,
+    double fp8_min,
+    double fp8_max,
+    bool scale_ue8m0);
+
+void mla_mudnn_bmm(
+    torch::Tensor &c,
+    torch::Tensor const &a,
+    torch::Tensor const &b);
+
+//========== pybind11 ==========
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
+{
+    // tensor
+    m.def("weak_ref_tensor", &weak_ref_tensor, "weak reference a tensor");
+    m.def("mudnn_w4a16_matmul", &mudnn_w4a16_matmul, "muDNN W4A16 matmul");
+    m.def("mla_mudnn_bmm", &mla_mudnn_bmm, "muDNN non-contiguous matmul");
+
+    // norm
+    m.def("rms_norm", &rms_norm, "Root mean square normalization");
+    m.def("fused_add_rms_norm", &fused_add_rms_norm,
+          "Fused add root mean square normalization");
+
+    m.def("musa_fused_add_rms_norm", &musa_fused_add_rms_norm,
+          "Fused add root mean square normalization");
+
+    m.def("rmsnorm_cast", &rmsnorm_cast, "fused rmsnorm and cast to fp8 (row-wise 128-block)");
+    m.def("rmsnorm_add_cast", &rmsnorm_add_cast, "fused rmsnorm add and cast to fp8 (row-wise 128-block)");
+    // embedding
+    m.def("rotary_embedding", &rotary_embedding, "rotary_embedding");
+    m.def("batched_rotary_embedding", &batched_rotary_embedding,
+          "batched_rotary_embedding");
+    m.def("rotary_embedding_contiguous", &rotary_embedding_contiguous, "rotary_embedding");
+    m.def("batched_rotary_embedding_contiguous", &batched_rotary_embedding_contiguous,
+          "batched_rotary_embedding");
+    // cache
+    m.def("reshape_and_cache", &reshape_and_cache, "reshape_and_cache");
+    m.def("reshape_and_cache_flash", &reshape_and_cache_flash, "reshape_and_cache_flash (MUSA)");
+    m.def("swap_blocks", &swap_blocks, "swap_blocks");
+    m.def("copy_blocks", &copy_blocks, "copy_blocks");
+    m.def("copy_blocks_mla", &copy_blocks_mla, "copy_blocks_mla (MUSA)");
+    m.def("concat_and_cache_mla", &concat_and_cache_mla,
+          "concat_and_cache_mla (MUSA)");
+    m.def("gather_cache", &gather_cache, "gather_cache (MUSA)");
+
+    // paged attention v1
+    m.def("paged_attention_v1", &paged_attention_v1, "paged_attention_v1");
+    m.def("paged_attention_v2", &paged_attention_v2, "paged_attention_v2 (MUSA)");
+    m.def("merge_attn_states", &merge_attn_states, "merge_attn_states (MUSA)");
+
+    // moe
+    m.def("topk_softmax", &topk_softmax, "topk_softmax (MUSA)");
+    m.def("moe_sum", &moe_sum, "moe_sum (MUSA)");
+    m.def("moe_align_block_size", &moe_align_block_size,
+          "moe_align_block_size (MUSA)");
+    m.def("sgl_moe_align_block_size", &sgl_moe_align_block_size,
+          "sgl_moe_align_block_size (MUSA)");
+    m.def("mp31_fused_grouped_topk", &mp31_fused_grouped_topk, "mp31_fused_grouped_topk (MUSA)");
+    m.def("mp22_fused_grouped_topk", &mp22_fused_grouped_topk, "mp22_fused_grouped_topk (MUSA)");
+    m.def("fused_moe_mutlass", &fused_moe_mutlass, "fused_moe_mutlass (MUSA)");
+    m.def("musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous", &musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous,
+          "musa_deep_gemm");
+
+    // activation
+    m.def("silu_and_mul", &silu_and_mul, "silu_and_mul (MUSA)");
+
+    // element-wise
+    m.def("_fused_mul_add", &_fused_mul_add, "fused multiply with add (MUSA)");
+
+    m.def("_gemm_swiglu", &gemm_swiglu, "fused gemm and swish glu");
+    m.def("_gemm_swiglu_w4a16", &gemm_swiglu_w4a16, "fused w4a16 gemm and swish glu");
+    m.def("_fused_moe_gemv", &fused_moe_gemv, "fused w4a16 gemm and swish glu");
+    m.def("_gemm_w4a16", &gemm_w4a16, "w4a16 gemv impl");
+    m.def("_fused_gemv", &fused_gemv, "fp8 gemv and w4a16 gemv per-group impl");
+    m.def("_group_scale_gemm", &group_scale_gemm, "w4a16 gemv per-group impl");
+
+    m.def("flash_attn_fwd", &flash_attn_fwd, "flash_attn_fwd");
+    m.def("mp22_flash_attn_varlen_func", &mp22_flash_attn_varlen_func, "mp22_flash_attn_varlen_func");
+    m.def("decode_mla", &decode_mla, "decode mla impl");
+    m.def("mp31_decode_mha", &mp31_decode_mha, "mp31 decode mha");
+    m.def("varlen_fa_seqlen_pad", &varlen_fa_seqlen_pad, "varlen_fa_seqlen_pad");
+    m.def("varlen_fa_seqlen_unpad", &varlen_fa_seqlen_unpad, "varlen_fa_seqlen_unpad");
+    m.def("varlen_fa_lse_seqlen_unpad", &varlen_fa_lse_seqlen_unpad, "varlen_fa_lse_seqlen_unpad");
+
+    // custom ar(all-reduce)
+    m.def("init_custom_ar", &init_custom_ar, "init_custom_ar");
+    m.def("all_reduce", &all_reduce, "all_reduce");
+    m.def("dispose", &dispose, "dispose");
+    m.def("meta_size", &meta_size, "meta_size");
+    m.def("register_buffer", &register_buffer, "register_buffer");
+    m.def("get_graph_buffer_ipc_meta", &get_graph_buffer_ipc_meta, "get_graph_buffer_ipc_meta");
+    m.def("register_graph_buffers", &register_graph_buffers, "register_graph_buffers");
+
+    // quantization
+    m.def("static_scaled_fp8_quant", &static_scaled_fp8_quant, "static scaled fp8 quant  (MUSA)");
+    m.def("dynamic_scaled_fp8_quant", &dynamic_scaled_fp8_quant, "dynamic scaled fp8 quant (MUSA)");
+    m.def("dynamic_per_token_scaled_fp8_quant", &dynamic_per_token_scaled_fp8_quant,
+          "dynamic per token scaled fp8 quant (MUSA)");
+    m.def("mudnn_w8a8_scaled_mm", &mudnn_w8a8_scaled_mm, "mudnn_w8a8_scaled_mm (MUSA)");
+    m.def("sgl_per_token_group_quant_fp8", &sgl_per_token_group_quant_fp8, "sgl_per_token_group_quant_fp8 (MUSA)");
+}
diff --git a/vllm_musa/csrc/pos_encoding_contiguous_kernels.mu b/vllm_musa/csrc/pos_encoding_contiguous_kernels.mu
new file mode 100644
index 000000000..4bd510615
--- /dev/null
+++ b/vllm_musa/csrc/pos_encoding_contiguous_kernels.mu
@@ -0,0 +1,249 @@
+#include <torch/all.h>
+#include "musa.h"
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+
+namespace vllm_musa {
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_token_rotary_embedding_contiguous(
+    scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim,
+    int64_t head_stride) {
+  int x_index, y_index;
+  scalar_t cos, sin;
+  if (IS_NEOX) {
+    // GPT-NeoX style rotary embedding.
+    x_index = rot_offset;
+    y_index = embed_dim + rot_offset;
+    cos = VLLM_LDG(cos_ptr + x_index);
+    sin = VLLM_LDG(sin_ptr + x_index);
+  } else {
+    // GPT-J style rotary embedding.
+    x_index = 2 * rot_offset;
+    y_index = 2 * rot_offset + 1;
+    cos = VLLM_LDG(cos_ptr + x_index / 2);
+    sin = VLLM_LDG(sin_ptr + x_index / 2);
+  }
+
+  scalar_t* x_ptr = arr + x_index * head_stride;
+  scalar_t* y_ptr = arr + y_index * head_stride;
+
+  const scalar_t x = *x_ptr;
+  const scalar_t y = *y_ptr;
+  *x_ptr = x * cos - y * sin;
+  *y_ptr = y * cos + x * sin;
+}
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_rotary_embedding_contiguous(
+    scalar_t* __restrict__ query,  // [num_tokens, num_heads, head_size]
+    scalar_t* __restrict__ key,    // [num_tokens, num_kv_heads, head_size]
+    const scalar_t* cache_ptr, const int head_size, const int num_heads,
+    const int num_kv_heads, const int rot_dim, const int token_idx,
+    const int64_t query_token_stride, const int64_t query_head_stride,
+    const int64_t query_dim_stride,
+    const int64_t key_token_stride, const int64_t key_head_stride,
+    const int64_t key_dim_stride) {
+  const int embed_dim = rot_dim / 2;
+  const scalar_t* cos_ptr = cache_ptr;
+  const scalar_t* sin_ptr = cache_ptr + embed_dim;
+
+  const int nq = num_heads * embed_dim;
+  for (int i = threadIdx.x; i < nq; i += blockDim.x) {
+    const int head_idx = i / embed_dim;
+    const int rot_offset = i % embed_dim;
+
+    scalar_t* head_query = query +
+                           token_idx * query_token_stride +
+                           head_idx * query_head_stride;
+
+    apply_token_rotary_embedding_contiguous<scalar_t, IS_NEOX>(
+        head_query, cos_ptr, sin_ptr, rot_offset, embed_dim, query_dim_stride);
+  }
+
+  const int nk = num_kv_heads * embed_dim;
+  for (int i = threadIdx.x; i < nk; i += blockDim.x) {
+    const int head_idx = i / embed_dim;
+    const int rot_offset = i % embed_dim;
+
+    scalar_t* head_key = key +
+                         token_idx * key_token_stride +
+                         head_idx * key_head_stride;
+
+    apply_token_rotary_embedding_contiguous<scalar_t, IS_NEOX>(
+        head_key, cos_ptr, sin_ptr, rot_offset, embed_dim, key_dim_stride);
+  }
+}
+
+template <typename scalar_t, bool IS_NEOX>
+__global__ void rotary_embedding_kernel_contiguous(
+    const int64_t* __restrict__ positions,  // [num_tokens]
+    scalar_t* __restrict__ query,           // [num_tokens, num_heads, head_size]
+    scalar_t* __restrict__ key,             // [num_tokens, num_kv_heads, head_size]
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim // 2]
+    const int rot_dim,
+    const int64_t query_token_stride, const int64_t query_head_stride,
+    const int64_t query_dim_stride,
+    const int64_t key_token_stride, const int64_t key_head_stride,
+    const int64_t key_dim_stride,
+    const int num_heads, const int num_kv_heads, const int head_size) {
+  // Each thread block is responsible for one token.
+  const int token_idx = blockIdx.x;
+  int64_t pos = positions[token_idx];
+  const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
+
+  apply_rotary_embedding_contiguous<scalar_t, IS_NEOX>(
+      query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
+      token_idx,
+      query_token_stride, query_head_stride, query_dim_stride,
+      key_token_stride, key_head_stride, key_dim_stride);
+}
+
+template <typename scalar_t, bool IS_NEOX>
+__global__ void batched_rotary_embedding_kernel_contiguous(
+    const int64_t* __restrict__ positions,  // [num_tokens]
+    scalar_t* __restrict__ query,           // [num_tokens, num_heads, head_size]
+    scalar_t* __restrict__ key,             // [num_tokens, num_kv_heads, head_size]
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim // 2]
+    const int64_t* __restrict__ cos_sin_cache_offsets,  // [num_tokens]
+    const int rot_dim,
+    const int64_t query_token_stride, const int64_t query_head_stride,
+    const int64_t query_dim_stride,  // 添加各维度步长
+    const int64_t key_token_stride, const int64_t key_head_stride,
+    const int64_t key_dim_stride,
+    const int num_heads, const int num_kv_heads, const int head_size) {
+  // Each thread block is responsible for one token.
+  const int token_idx = blockIdx.x;
+  int64_t pos = positions[token_idx];
+  int64_t cos_sin_cache_offset = cos_sin_cache_offsets[token_idx];
+  const scalar_t* cache_ptr =
+      cos_sin_cache + (cos_sin_cache_offset + pos) * rot_dim;
+
+  apply_rotary_embedding_contiguous<scalar_t, IS_NEOX>(
+      query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
+      token_idx,
+      query_token_stride, query_head_stride, query_dim_stride,
+      key_token_stride, key_head_stride, key_dim_stride);
+}
+
+}  // namespace vllm_musa
+
+void rotary_embedding_contiguous(
+    torch::Tensor& positions,  // [num_tokens]
+    torch::Tensor& query,      // [num_tokens, num_heads, head_size]
+    torch::Tensor& key,        // [num_tokens, num_kv_heads, head_size]
+    int64_t head_size,
+    torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    bool is_neox) {
+  int64_t num_tokens = positions.size(0);
+
+  TORCH_CHECK(query.dim() == 3, "query must be 3D [num_tokens, num_heads, head_size]");
+  TORCH_CHECK(key.dim() == 3, "key must be 3D [num_tokens, num_kv_heads, head_size]");
+  TORCH_CHECK(query.size(0) == num_tokens && key.size(0) == num_tokens,
+              "query, key and positions must have the same number of tokens");
+
+  int64_t query_token_stride = query.stride(0);
+  int64_t query_head_stride = query.stride(1);
+  int64_t query_dim_stride = query.stride(2);
+
+  int64_t key_token_stride = key.stride(0);
+  int64_t key_head_stride = key.stride(1);
+  int64_t key_dim_stride = key.stride(2);
+
+  int num_heads = query.size(1);
+  int num_kv_heads = key.size(1);
+  int rot_dim = cos_sin_cache.size(1);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding_contiguous", [&] {
+    if (is_neox) {
+      vllm_musa::rotary_embedding_kernel_contiguous<scalar_t, true><<<grid, block, 0, stream>>>(
+          positions.data_ptr<int64_t>(),
+          query.data_ptr<scalar_t>(),
+          key.data_ptr<scalar_t>(),
+          cos_sin_cache.data_ptr<scalar_t>(),
+          rot_dim,
+          query_token_stride, query_head_stride, query_dim_stride,
+          key_token_stride, key_head_stride, key_dim_stride,
+          num_heads, num_kv_heads, head_size);
+    } else {
+      vllm_musa::rotary_embedding_kernel_contiguous<scalar_t, false><<<grid, block, 0, stream>>>(
+          positions.data_ptr<int64_t>(),
+          query.data_ptr<scalar_t>(),
+          key.data_ptr<scalar_t>(),
+          cos_sin_cache.data_ptr<scalar_t>(),
+          rot_dim,
+          query_token_stride, query_head_stride, query_dim_stride,
+          key_token_stride, key_head_stride, key_dim_stride,
+          num_heads, num_kv_heads, head_size);
+    }
+  });
+}
+
+void batched_rotary_embedding_contiguous(
+    torch::Tensor& positions,  // [num_tokens]
+    torch::Tensor& query,      // [num_tokens, num_heads, head_size]
+    torch::Tensor& key,        // [num_tokens, num_kv_heads, head_size]
+    int64_t head_size,
+    torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    bool is_neox, int64_t rot_dim,
+    torch::Tensor& cos_sin_cache_offsets  // [num_tokens]
+) {
+  int64_t num_tokens = cos_sin_cache_offsets.size(0);
+
+
+  TORCH_CHECK(positions.size(0) == num_tokens,
+              "positions must have the same num_tokens as cos_sin_cache_offsets");
+  TORCH_CHECK(query.dim() == 3, "query must be 3D [num_tokens, num_heads, head_size]");
+  TORCH_CHECK(key.dim() == 3, "key must be 3D [num_tokens, num_kv_heads, head_size]");
+
+  int64_t query_token_stride = query.stride(0);
+  int64_t query_head_stride = query.stride(1);
+  int64_t query_dim_stride = query.stride(2);
+
+  int64_t key_token_stride = key.stride(0);
+  int64_t key_head_stride = key.stride(1);
+  int64_t key_dim_stride = key.stride(2);
+
+  int num_heads = query.size(1);
+  int num_kv_heads = key.size(1);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "batched_rotary_embedding_contiguous", [&] {
+    if (is_neox) {
+      vllm_musa::batched_rotary_embedding_kernel_contiguous<scalar_t, true><<<grid, block, 0, stream>>>(
+          positions.data_ptr<int64_t>(),
+          query.data_ptr<scalar_t>(),
+          key.data_ptr<scalar_t>(),
+          cos_sin_cache.data_ptr<scalar_t>(),
+          cos_sin_cache_offsets.data_ptr<int64_t>(),
+          rot_dim,
+          query_token_stride, query_head_stride, query_dim_stride,
+          key_token_stride, key_head_stride, key_dim_stride,
+          num_heads, num_kv_heads, head_size);
+    } else {
+      vllm_musa::batched_rotary_embedding_kernel_contiguous<scalar_t, false><<<grid, block, 0, stream>>>(
+          positions.data_ptr<int64_t>(),
+          query.data_ptr<scalar_t>(),
+          key.data_ptr<scalar_t>(),
+          cos_sin_cache.data_ptr<scalar_t>(),
+          cos_sin_cache_offsets.data_ptr<int64_t>(),
+          rot_dim,
+          query_token_stride, query_head_stride, query_dim_stride,
+          key_token_stride, key_head_stride, key_dim_stride,
+          num_heads, num_kv_heads, head_size);
+    }
+  });
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/pos_encoding_kernels.mu b/vllm_musa/csrc/pos_encoding_kernels.mu
new file mode 100644
index 000000000..916e7b6ec
--- /dev/null
+++ b/vllm_musa/csrc/pos_encoding_kernels.mu
@@ -0,0 +1,277 @@
+#include <torch/all.h>
+#include "musa.h"
+
+#include "cuda_compat.h"
+#include "dispatch_utils.h"
+
+namespace vllm_musa {
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_token_rotary_embedding(
+    scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr,
+    const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim) {
+  int x_index, y_index;
+  scalar_t cos, sin;
+  if (IS_NEOX) {
+    // GPT-NeoX style rotary embedding.
+    x_index = rot_offset;
+    y_index = embed_dim + rot_offset;
+    cos = VLLM_LDG(cos_ptr + x_index);
+    sin = VLLM_LDG(sin_ptr + x_index);
+  } else {
+    // GPT-J style rotary embedding.
+    x_index = 2 * rot_offset;
+    y_index = 2 * rot_offset + 1;
+    cos = VLLM_LDG(cos_ptr + x_index / 2);
+    sin = VLLM_LDG(sin_ptr + x_index / 2);
+  }
+
+  const scalar_t x = arr[x_index];
+  const scalar_t y = arr[y_index];
+  arr[x_index] = x * cos - y * sin;
+  arr[y_index] = y * cos + x * sin;
+}
+
+template <typename scalar_t, bool IS_NEOX>
+inline __device__ void apply_rotary_embedding(
+    scalar_t* __restrict__ query,  // [batch_size, seq_len, num_heads,
+                                   // head_size] or [num_tokens, num_heads,
+                                   // head_size]
+    scalar_t* __restrict__ key,    // [batch_size, seq_len, num_kv_heads,
+                                   // head_size] or [num_tokens, num_kv_heads,
+                                   // head_size]
+    const scalar_t* cache_ptr, const int head_size, const int num_heads,
+    const int num_kv_heads, const int rot_dim, const int token_idx,
+    const int64_t query_stride, const int64_t key_stride) {
+  const int embed_dim = rot_dim / 2;
+  const scalar_t* cos_ptr = cache_ptr;
+  const scalar_t* sin_ptr = cache_ptr + embed_dim;
+
+  const int nq = num_heads * embed_dim;
+  for (int i = threadIdx.x; i < nq; i += blockDim.x) {
+    const int head_idx = i / embed_dim;
+    const int64_t token_head = token_idx * query_stride + head_idx * head_size;
+    const int rot_offset = i % embed_dim;
+    apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+        query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+  }
+
+  const int nk = num_kv_heads * embed_dim;
+  for (int i = threadIdx.x; i < nk; i += blockDim.x) {
+    const int head_idx = i / embed_dim;
+    const int64_t token_head = token_idx * key_stride + head_idx * head_size;
+    const int rot_offset = i % embed_dim;
+    apply_token_rotary_embedding<scalar_t, IS_NEOX>(
+        key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
+  }
+}
+
+template <typename scalar_t, bool IS_NEOX>
+__global__ void rotary_embedding_kernel(
+    const int64_t* __restrict__ positions,  // [batch_size, seq_len] or
+                                            // [num_tokens]
+    scalar_t* __restrict__ query,           // [batch_size, seq_len, num_heads,
+                                   // head_size] or [num_tokens, num_heads,
+                                   // head_size]
+    scalar_t* __restrict__ key,  // [batch_size, seq_len, num_kv_heads,
+                                 // head_size] or [num_tokens, num_kv_heads,
+                                 // head_size]
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
+                                                 // 2]
+    const int rot_dim, const int64_t query_stride, const int64_t key_stride,
+    const int num_heads, const int num_kv_heads, const int head_size) {
+  // Each thread block is responsible for one token.
+  const int token_idx = blockIdx.x;
+  int64_t pos = positions[token_idx];
+  const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
+
+  apply_rotary_embedding<scalar_t, IS_NEOX>(
+      query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
+      token_idx, query_stride, key_stride);
+}
+
+template <typename scalar_t, bool IS_NEOX>
+__global__ void batched_rotary_embedding_kernel(
+    const int64_t* __restrict__ positions,  // [batch_size, seq_len] or
+                                            // [num_tokens]
+    scalar_t* __restrict__ query,           // [batch_size, seq_len, num_heads,
+                                   // head_size] or [num_tokens, num_heads,
+                                   // head_size]
+    scalar_t* __restrict__ key,  // [batch_size, seq_len, num_kv_heads,
+                                 // head_size] or [num_tokens, num_kv_heads,
+                                 // head_size]
+    const scalar_t* __restrict__ cos_sin_cache,  // [max_position, 2, rot_dim //
+                                                 // 2]
+    const int64_t* __restrict__ cos_sin_cache_offsets,  // [batch_size, seq_len]
+                                                        // or [num_tokens]
+    const int rot_dim, const int64_t query_stride, const int64_t key_stride,
+    const int num_heads, const int num_kv_heads, const int head_size) {
+  // Each thread block is responsible for one token.
+  const int token_idx = blockIdx.x;
+  int64_t pos = positions[token_idx];
+  int64_t cos_sin_cache_offset = cos_sin_cache_offsets[token_idx];
+  const scalar_t* cache_ptr =
+      cos_sin_cache + (cos_sin_cache_offset + pos) * rot_dim;
+
+  apply_rotary_embedding<scalar_t, IS_NEOX>(
+      query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim,
+      token_idx, query_stride, key_stride);
+}
+
+}  // namespace vllm_musa
+
+void rotary_embedding(
+    torch::Tensor& positions,  // [batch_size, seq_len] or [num_tokens]
+    torch::Tensor& query,  // [batch_size, seq_len, num_heads * head_size] or
+                           // [num_tokens, num_heads * head_size] or
+                           // [batch_size, seq_len, num_heads, head_size] or
+                           // [num_tokens, num_heads, head_size]
+    torch::Tensor& key,    // [batch_size, seq_len, num_kv_heads * head_size] or
+                           // [num_tokens, num_kv_heads * head_size] or
+                           // [batch_size, seq_len, num_heads, head_size] or
+                           // [num_tokens, num_heads, head_size]
+    int64_t head_size,
+    torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    bool is_neox) {
+  // num_tokens = batch_size * seq_len
+  int64_t num_tokens = positions.numel();
+  int positions_ndim = positions.dim();
+
+  // Make sure num_tokens dim is consistent across positions, query, and key.
+  TORCH_CHECK(
+      positions_ndim == 1 || positions_ndim == 2,
+      "positions must have shape [num_tokens] or [batch_size, seq_len]");
+  if (positions_ndim == 1) {
+    TORCH_CHECK(
+        query.size(0) == positions.size(0) && key.size(0) == positions.size(0),
+        "query, key and positions must have the same number of tokens");
+  }
+  if (positions_ndim == 2) {
+    TORCH_CHECK(
+        query.size(0) == positions.size(0) &&
+            key.size(0) == positions.size(0) &&
+            query.size(1) == positions.size(1) &&
+            key.size(1) == positions.size(1),
+        "query, key and positions must have the same batch_size and seq_len");
+  }
+
+  // Make sure head_size is valid for query and key
+  // hidden_size = num_heads * head_size
+  int query_hidden_size = query.numel() / num_tokens;
+  int key_hidden_size = key.numel() / num_tokens;
+  TORCH_CHECK(query_hidden_size % head_size == 0);
+  TORCH_CHECK(key_hidden_size % head_size == 0);
+
+  // Make sure query and key have consistent number of heads
+  int num_heads = query_hidden_size / head_size;
+  int num_kv_heads = key_hidden_size / head_size;
+  TORCH_CHECK(num_heads % num_kv_heads == 0);
+
+  int rot_dim = cos_sin_cache.size(1);
+  int seq_dim_idx = positions_ndim - 1;
+  int64_t query_stride = query.stride(seq_dim_idx);
+  int64_t key_stride = key.stride(seq_dim_idx);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
+    if (is_neox) {
+      vllm_musa::rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
+          positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
+          key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(), rot_dim,
+          query_stride, key_stride, num_heads, num_kv_heads, head_size);
+    } else {
+      vllm_musa::rotary_embedding_kernel<scalar_t, false>
+          <<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
+              key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(),
+              rot_dim, query_stride, key_stride, num_heads, num_kv_heads,
+              head_size);
+    }
+  });
+}
+
+/*
+Batched version of rotary embedding, pack multiple LoRAs together
+and process in batched manner.
+*/
+void batched_rotary_embedding(
+    torch::Tensor& positions,  // [batch_size, seq_len] or [num_tokens]
+    torch::Tensor& query,  // [batch_size, seq_len, num_heads * head_size] or
+                           // [num_tokens, num_heads * head_size] or
+                           // [batch_size, seq_len, num_heads, head_size] or
+                           // [num_tokens, num_heads, head_size]
+    torch::Tensor& key,    // [batch_size, seq_len, num_kv_heads * head_size] or
+                           // [num_tokens, num_kv_heads * head_size] or
+                           // [batch_size, seq_len, num_heads, head_size] or
+                           // [num_tokens, num_heads, head_size]
+    int64_t head_size,
+    torch::Tensor& cos_sin_cache,  // [max_position, rot_dim]
+    bool is_neox, int64_t rot_dim,
+    torch::Tensor& cos_sin_cache_offsets  // [num_tokens] or [batch_size]
+) {
+  // num_tokens = batch_size * seq_len
+  int64_t num_tokens = cos_sin_cache_offsets.size(0);
+  TORCH_CHECK(
+      positions.size(0) == num_tokens || positions.numel() == num_tokens,
+      "positions must have the same num_tokens or batch_size as "
+      "cos_sin_cache_offsets");
+
+  int positions_ndim = positions.dim();
+  // Make sure num_tokens dim is consistent across positions, query, and key.
+  TORCH_CHECK(
+      positions_ndim == 1 || positions_ndim == 2,
+      "positions must have shape [num_tokens] or [batch_size, seq_len]");
+  if (positions_ndim == 1) {
+    TORCH_CHECK(
+        query.size(0) == positions.size(0) && key.size(0) == positions.size(0),
+        "query, key and positions must have the same number of tokens");
+  }
+  if (positions_ndim == 2) {
+    TORCH_CHECK(
+        query.size(0) == positions.size(0) &&
+            key.size(0) == positions.size(0) &&
+            query.size(1) == positions.size(1) &&
+            key.size(1) == positions.size(1),
+        "query, key and positions must have the same batch_size and seq_len");
+  }
+
+  // Make sure head_size is valid for query and key
+  int query_hidden_size = query.numel() / num_tokens;
+  int key_hidden_size = key.numel() / num_tokens;
+  TORCH_CHECK(query_hidden_size % head_size == 0);
+  TORCH_CHECK(key_hidden_size % head_size == 0);
+
+  // Make sure query and key have concistent number of heads
+  int num_heads = query_hidden_size / head_size;
+  int num_kv_heads = key_hidden_size / head_size;
+  TORCH_CHECK(num_heads % num_kv_heads == 0);
+
+  int seq_dim_idx = positions_ndim - 1;
+  int64_t query_stride = query.stride(seq_dim_idx);
+  int64_t key_stride = key.stride(seq_dim_idx);
+
+  dim3 grid(num_tokens);
+  dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding", [&] {
+    if (is_neox) {
+      vllm_musa::batched_rotary_embedding_kernel<scalar_t, true>
+          <<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
+              key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(),
+              cos_sin_cache_offsets.data_ptr<int64_t>(), rot_dim, query_stride,
+              key_stride, num_heads, num_kv_heads, head_size);
+    } else {
+      vllm_musa::batched_rotary_embedding_kernel<scalar_t, false>
+          <<<grid, block, 0, stream>>>(
+              positions.data_ptr<int64_t>(), query.data_ptr<scalar_t>(),
+              key.data_ptr<scalar_t>(), cos_sin_cache.data_ptr<scalar_t>(),
+              cos_sin_cache_offsets.data_ptr<int64_t>(), rot_dim, query_stride,
+              key_stride, num_heads, num_kv_heads, head_size);
+    }
+  });
+}
diff --git a/vllm_musa/csrc/quantization/fp8/common.mu b/vllm_musa/csrc/quantization/fp8/common.mu
new file mode 100644
index 000000000..9d96fa75d
--- /dev/null
+++ b/vllm_musa/csrc/quantization/fp8/common.mu
@@ -0,0 +1,143 @@
+#include "common.muh"
+#include "../../dispatch_utils.h"
+
+#include "../../musa.h"
+
+namespace vllm_musa {
+
+template <typename scalar_t>
+__global__ void scaled_fp8_quant_kernel(FP8_TYPE* __restrict__ out,
+                                        const scalar_t* __restrict__ input,
+                                        const float* __restrict__ scale,
+                                        int64_t num_elems) {
+  int tid = blockDim.x * blockIdx.x + threadIdx.x;
+
+  // Invert the scale so that we can use multiplications to avoid expensive
+  // division.
+  const float inverted_scale = 1.0f / (*scale);
+  scaled_fp8_conversion_vec<scalar_t, true>(
+      out, input, inverted_scale, num_elems, tid, blockDim.x * gridDim.x);
+}
+
+template <typename scalar_t>
+__global__ void dynamic_per_token_scaled_fp8_quant_kernel(
+    FP8_TYPE* __restrict__ out, float* __restrict__ scale,
+    scalar_t const* __restrict__ input, float const* __restrict__ scale_ub,
+    const int hidden_size) {
+  float const min_scaling_factor = 1.0f / (FP8_E4M3_MAX * 512.f);
+
+  int const tid = threadIdx.x;
+  int const token_idx = blockIdx.x;
+
+  // Use int64 to avoid overflowing an int32 when calculating this offset
+  int64_t offset = static_cast<int64_t>(token_idx) * hidden_size;
+  scalar_t const* __restrict__ token_input = &input[offset];
+  FP8_TYPE* __restrict__ token_output = &out[offset];
+
+  // For vectorization, token_input and token_output pointers need to be
+  // aligned at 8-byte and 4-byte addresses respectively.
+  bool const can_vectorize = hidden_size % 4 == 0;
+
+  float absmax_val = 0.0f;
+  if (can_vectorize) {
+    absmax_val = thread_max_vec(token_input, hidden_size, tid, blockDim.x);
+  } else {
+    for (int i = tid; i < hidden_size; i += blockDim.x) {
+      float const x = static_cast<float>(token_input[i]);
+      absmax_val = max(absmax_val, fabs(x));
+    }
+  }
+
+  using BlockReduce = cub::BlockReduce<float, 1024>;
+  __shared__ typename BlockReduce::TempStorage reduceStorage;
+  float const block_absmax_val_maybe =
+      BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
+  __shared__ float token_scale;
+  if (tid == 0) {
+    if (scale_ub) {
+      token_scale = min(block_absmax_val_maybe, *scale_ub);
+    } else {
+      token_scale = block_absmax_val_maybe;
+    }
+    // token scale computation
+    token_scale = max(token_scale / FP8_E4M3_MAX, min_scaling_factor);
+    scale[token_idx] = token_scale;
+  }
+  __syncthreads();
+
+  // Note that we don't use inverted scales so we can match FBGemm impl.
+  if (can_vectorize) {
+    scaled_fp8_conversion_vec<scalar_t, false>(
+        token_output, token_input, token_scale, hidden_size, tid, blockDim.x);
+  } else {
+    for (int i = tid; i < hidden_size; i += blockDim.x) {
+      token_output[i] = scaled_fp8_conversion<false>(
+          static_cast<float>(token_input[i]), token_scale);
+    }
+  }
+}
+
+}  // namespace vllm_musa
+
+void static_scaled_fp8_quant(torch::Tensor& out,          // [..., d]
+                             torch::Tensor const& input,  // [..., d]
+                             torch::Tensor const& scale)  // [1]
+{
+  int64_t num_tokens = input.numel() / input.size(-1);
+  int64_t num_elems = input.numel();
+  dim3 grid(num_tokens);
+  dim3 block(1024);
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "scaled_fp8_quant_kernel", [&] {
+        vllm_musa::scaled_fp8_quant_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            out.data_ptr<FP8_TYPE>(), input.data_ptr<scalar_t>(),
+            scale.data_ptr<float>(), num_elems);
+      });
+}
+
+void dynamic_scaled_fp8_quant(torch::Tensor& out,          // [..., d]
+                              torch::Tensor const& input,  // [..., d]
+                              torch::Tensor& scale)        // [1]
+{
+  int64_t num_tokens = input.numel() / input.size(-1);
+  int64_t num_elems = input.numel();
+  dim3 grid(num_tokens);
+  dim3 block(1024);
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "scaled_fp8_quant_kernel", [&] {
+        vllm_musa::segmented_max_reduction<scalar_t><<<grid, block, 0, stream>>>(
+            scale.data_ptr<float>(), input.data_ptr<scalar_t>(), num_elems);
+        vllm_musa::scaled_fp8_quant_kernel<scalar_t><<<grid, block, 0, stream>>>(
+            out.data_ptr<FP8_TYPE>(), input.data_ptr<scalar_t>(),
+            scale.data_ptr<float>(), num_elems);
+      });
+}
+
+void dynamic_per_token_scaled_fp8_quant(
+    torch::Tensor& out,          // [..., d]
+    torch::Tensor const& input,  // [..., d]
+    torch::Tensor& scales, std::optional<at::Tensor> const& scale_ub) {
+  TORCH_CHECK(input.is_contiguous());
+  TORCH_CHECK(out.is_contiguous());
+
+  int const hidden_size = input.size(-1);
+  int const num_tokens = input.numel() / hidden_size;
+  dim3 const grid(num_tokens);
+  dim3 const block(std::min(hidden_size, 1024));
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  VLLM_DISPATCH_FLOATING_TYPES(
+      input.scalar_type(), "dynamic_per_token_scaled_fp8_quant_kernel", [&] {
+        vllm_musa::dynamic_per_token_scaled_fp8_quant_kernel<scalar_t>
+            <<<grid, block, 0, stream>>>(
+                out.data_ptr<FP8_TYPE>(), scales.data_ptr<float>(),
+                input.data_ptr<scalar_t>(),
+                scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
+                hidden_size);
+      });
+}
diff --git a/vllm_musa/csrc/quantization/fp8/common.muh b/vllm_musa/csrc/quantization/fp8/common.muh
new file mode 100644
index 000000000..00782cc59
--- /dev/null
+++ b/vllm_musa/csrc/quantization/fp8/common.muh
@@ -0,0 +1,144 @@
+#pragma once
+
+#include "../vectorization.muh"
+
+#include <cmath>
+#include <c10/core/ScalarType.h>
+
+#include <c10/util/Float8_e4m3fn.h>
+using FP8_TYPE = c10::Float8_e4m3fn;
+constexpr auto FP8_E4M3_MAX = std::numeric_limits<FP8_TYPE>::max();
+constexpr static auto kFp8Type = c10::CppTypeToScalarType<FP8_TYPE>::value;
+
+namespace vllm_musa {
+
+__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
+  float old;
+  old = (value >= 0)
+            ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
+            : __uint_as_float(
+                  atomicMin((unsigned int*)addr, __float_as_uint(value)));
+
+  return old;
+}
+
+template <bool is_scale_inverted>
+__device__ __forceinline__ FP8_TYPE scaled_fp8_conversion(float const val,
+                                                          float const scale) {
+  float x = 0.0f;
+  if constexpr (is_scale_inverted) {
+    x = val * scale;
+  } else {
+    x = val / scale;
+  }
+
+  float r = fmax(-FP8_E4M3_MAX, fmin(x, FP8_E4M3_MAX));
+  return static_cast<c10::Float8_e4m3fn>(r);
+}
+
+// Compute the absolute maximum m of the input tensor and store
+// m / float8_e4m3::max() in *scale. Each thread block performs a
+// reduction tree and the memory in scale is atomically updated.
+// So to get the right answer, *scale needs to be initialized to
+// a value <= 0.0 and we need to wait for all thread blocks to
+// finish before consuming *scale.
+template <typename scalar_t>
+__global__ void segmented_max_reduction(float* __restrict__ scale,
+                                        const scalar_t* __restrict__ input,
+                                        int64_t num_elems) {
+  __shared__ float cache[1024];
+  int64_t i = blockDim.x * blockIdx.x + threadIdx.x;
+
+  // First store maximum for all values processes by
+  // the current thread in cache[threadIdx.x]
+  scalar_t tmp = 0.0;
+  while (i < num_elems) {
+    float x = static_cast<float>(input[i]);
+    tmp = max(tmp, fabs(x));
+    i += blockDim.x * gridDim.x;
+  }
+  cache[threadIdx.x] = tmp;
+
+  __syncthreads();
+
+  // Now perform parallel reduction within the thread block
+  int ib = blockDim.x / 2;
+  while (ib != 0) {
+    if (threadIdx.x < ib && cache[threadIdx.x + ib] > cache[threadIdx.x]) {
+      cache[threadIdx.x] = cache[threadIdx.x + ib];
+    }
+    __syncthreads();
+    ib /= 2;
+  }
+  // Finally, since cache[0] contains the maximum for this thread block,
+  // atomically write the max to the target location
+  if (threadIdx.x == 0) {
+    atomicMaxFloat(scale, cache[0] / FP8_E4M3_MAX);
+  }
+}
+
+template <typename scalar_t>
+__device__ float thread_max_vec(scalar_t const* __restrict__ input,
+                                int64_t const num_elems, int const tid,
+                                int const step) {
+  // Vectorized input/output to better utilize memory bandwidth.
+  vec4_t<scalar_t> const* vectorized_in =
+      reinterpret_cast<vec4_t<scalar_t> const*>(input);
+
+  int64_t const num_vec_elems = num_elems >> 2;
+  float absmax_val = 0.0f;
+
+#pragma unroll 4
+  for (int64_t i = tid; i < num_vec_elems; i += step) {
+    vec4_t<scalar_t> in_vec = vectorized_in[i];
+    absmax_val = max(absmax_val, fabs(in_vec.x));
+    absmax_val = max(absmax_val, fabs(in_vec.y));
+    absmax_val = max(absmax_val, fabs(in_vec.z));
+    absmax_val = max(absmax_val, fabs(in_vec.w));
+  }
+
+  // Handle the remaining elements if num_elems is not divisible by 4
+  for (int64_t i = num_vec_elems * 4 + tid; i < num_elems; i += step) {
+    absmax_val = max(absmax_val, fabs(input[i]));
+  }
+
+  return absmax_val;
+}
+
+template <typename scalar_t, bool is_scale_inverted>
+__device__ void scaled_fp8_conversion_vec(FP8_TYPE* __restrict__ out,
+                                          scalar_t const* __restrict__ input,
+                                          float const scale,
+                                          int64_t const num_elems,
+                                          int const tid, int const step) {
+  using float8x4_t = q8x4_t<FP8_TYPE>;
+  // Vectorized input/output to better utilize memory bandwidth.
+  auto const* vectorized_in = reinterpret_cast<vec4_t<scalar_t> const*>(input);
+  auto* vectorized_out = reinterpret_cast<float8x4_t*>(out);
+
+  int64_t const num_vec_elems = num_elems >> 2;
+
+#pragma unroll 4
+  for (int64_t i = tid; i < num_vec_elems; i += step) {
+    vec4_t<scalar_t> in_vec = vectorized_in[i];
+    float8x4_t out_vec;
+
+    out_vec.x = scaled_fp8_conversion<is_scale_inverted>(
+        static_cast<float>(in_vec.x), scale);
+    out_vec.y = scaled_fp8_conversion<is_scale_inverted>(
+        static_cast<float>(in_vec.y), scale);
+    out_vec.z = scaled_fp8_conversion<is_scale_inverted>(
+        static_cast<float>(in_vec.z), scale);
+    out_vec.w = scaled_fp8_conversion<is_scale_inverted>(
+        static_cast<float>(in_vec.w), scale);
+    vectorized_out[i] = out_vec;
+  }
+
+  // Handle the remaining elements if num_elems is not divisible by 4
+  for (int64_t i = num_vec_elems * 4 + tid; i < num_elems; i += step) {
+    out[i] = scaled_fp8_conversion<is_scale_inverted>(
+        static_cast<float>(input[i]), scale);
+  }
+}
+
+}  // namespace vllm_musa
diff --git a/vllm_musa/csrc/quantization/fp8/per_token_group_quant_8bit.mu b/vllm_musa/csrc/quantization/fp8/per_token_group_quant_8bit.mu
new file mode 100644
index 000000000..becf00905
--- /dev/null
+++ b/vllm_musa/csrc/quantization/fp8/per_token_group_quant_8bit.mu
@@ -0,0 +1,279 @@
+#include <torch/all.h>
+#include "torch_musa/csrc/aten/musa/MUSAContext.h"
+#include <c10/util/Float8_e4m3fn.h>
+
+#include <cmath>
+#include "vec_dtypes.cuh"
+
+
+__device__ __forceinline__ float GroupReduceMax(float val, const int tid) {
+  unsigned mask = 0xffff;
+
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 8));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 4));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 2));
+  val = fmaxf(val, __shfl_xor_sync(mask, val, 1));
+  return val;
+}
+
+template <
+    typename T,
+    typename DST_DTYPE,
+    bool IS_COLUMN_MAJOR = false,
+    bool SCALE_UE8M0 = false,
+    typename scale_packed_t = std::conditional_t<SCALE_UE8M0, uint32_t, float>>
+__global__ void per_token_group_quant_8bit_kernel(
+    const T* __restrict__ input,
+    void* __restrict__ output_q,
+    scale_packed_t* __restrict__ output_s,
+    const int group_size,
+    const int num_groups,
+    const int groups_per_block,
+    const float eps,
+    const float min_8bit,
+    const float max_8bit,
+    const int scale_num_rows = 0,
+    const int scale_stride = 0) {
+  const int threads_per_group = 16;
+  const int local_group_id = threadIdx.x / threads_per_group;
+  const int lane_id = threadIdx.x % threads_per_group;
+
+  const int block_group_id = blockIdx.x * groups_per_block;
+  const int global_group_id = block_group_id + local_group_id;
+  const int block_group_offset = global_group_id * group_size;
+
+  float local_absmax = eps;
+
+  using scale_element_t = std::conditional_t<SCALE_UE8M0, uint8_t, float>;
+  static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
+
+  const T* group_input = input + block_group_offset;
+  DST_DTYPE* group_output = static_cast<DST_DTYPE*>(output_q) + block_group_offset;
+  scale_element_t* scale_output;
+
+  if constexpr (IS_COLUMN_MAJOR) {
+    const int num_elems_per_pack = static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
+    const int scale_num_rows_element = scale_num_rows * num_elems_per_pack;
+    const int row_idx = global_group_id / scale_num_rows_element;
+    const int col_idx_raw = global_group_id % scale_num_rows_element;
+    const int col_idx = col_idx_raw / num_elems_per_pack;
+    const int pack_idx = col_idx_raw % num_elems_per_pack;
+    scale_output = reinterpret_cast<scale_element_t*>(output_s) +
+                   (col_idx * scale_stride * num_elems_per_pack + row_idx * num_elems_per_pack + pack_idx);
+  } else {
+    static_assert(!SCALE_UE8M0);
+    scale_output = output_s + global_group_id;
+  }
+
+  constexpr uint32_t vec_size = 16 / sizeof(T);
+  using vec_t = flashinfer::vec_t<T, vec_size>;
+
+  const int32_t num_vec_elems = group_size / vec_size;
+
+  for (int32_t i = lane_id; i < num_vec_elems; i += 16) {
+    vec_t input_vec;
+    input_vec.cast_load(group_input + i * vec_size);
+
+#pragma unroll
+    for (uint32_t j = 0; j < vec_size; ++j) {
+      float val = static_cast<float>(input_vec[j]);
+      float abs_val = fabsf(val);
+      local_absmax = fmaxf(local_absmax, abs_val);
+    }
+  }
+
+  local_absmax = GroupReduceMax(local_absmax, lane_id);
+
+  float y_s = local_absmax / max_8bit;
+  if constexpr (SCALE_UE8M0) {
+    y_s = exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f))));
+  }
+
+  // TODO can optimize
+  scale_element_t y_s_quant;
+  if constexpr (SCALE_UE8M0) {
+    y_s_quant = (uint8_t)(((int)log2f(y_s)) + 127);
+  } else {
+    y_s_quant = y_s;
+  }
+
+  if (lane_id == 0) {
+    *scale_output = y_s_quant;
+  }
+
+  for (int32_t i = lane_id; i < num_vec_elems; i += 16) {
+    vec_t input_vec;
+    input_vec.cast_load(group_input + i * vec_size);
+
+#pragma unroll
+    for (uint32_t j = 0; j < vec_size; ++j) {
+      float val = static_cast<float>(input_vec[j]);
+      float q_val = fminf(fmaxf(val / y_s, min_8bit), max_8bit);
+      group_output[i * vec_size + j] = DST_DTYPE(q_val);
+    }
+  }
+}
+
+#define CHECK_MUSA(x) TORCH_CHECK(x.is_privateuseone(), #x " must be a MUSA tensor")
+
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+
+#define CHECK_INPUT(x) \
+  CHECK_MUSA(x);       \
+  CHECK_CONTIGUOUS(x)
+
+#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b)
+
+#define _DISPATCH_CASE_F16(c_type, ...) \
+  case at::ScalarType::Half: {          \
+    using c_type = __half;             \
+    return __VA_ARGS__();               \
+  }
+
+#define _DISPATCH_CASE_BF16(c_type, ...) \
+  case at::ScalarType::BFloat16: {       \
+    using c_type = mt_bfloat16;          \
+    return __VA_ARGS__();                \
+  }
+
+#define DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(pytorch_dtype, c_type, ...)           \
+  [&]() -> bool {                                                                        \
+    switch (pytorch_dtype) {                                                             \
+      case at::ScalarType::Float: {                                                      \
+        using c_type = float;                                                            \
+        return __VA_ARGS__();                                                            \
+      }                                                                                  \
+        _DISPATCH_CASE_F16(c_type, __VA_ARGS__)                                          \
+        _DISPATCH_CASE_BF16(c_type, __VA_ARGS__)                                         \
+      default:                                                                           \
+        std::ostringstream oss;                                                          \
+        oss << __PRETTY_FUNCTION__ << " failed to dispatch data type " << pytorch_dtype; \
+        TORCH_CHECK(false, oss.str());                                                   \
+        return false;                                                                    \
+    }                                                                                    \
+  }()
+
+void sgl_per_token_group_quant_8bit(
+    torch::Tensor input,
+    torch::Tensor output_q,
+    torch::Tensor output_s,
+    int64_t group_size,
+    double eps,
+    double min_8bit,
+    double max_8bit,
+    bool scale_ue8m0 = false) {
+  CHECK_INPUT(input);
+  CHECK_INPUT(output_q);
+
+  const int num_groups = input.numel() / group_size;
+
+  CHECK_EQ(input.numel() % group_size, 0);
+  CHECK_EQ(output_s.dim(), 2);
+
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  constexpr int THREADS_PER_GROUP = 16;
+
+  int groups_per_block = 1;
+
+  if (num_groups % 16 == 0) {
+    groups_per_block = 16;
+  } else if (num_groups % 8 == 0) {
+    groups_per_block = 8;
+  } else if (num_groups % 4 == 0) {
+    groups_per_block = 4;
+  } else if (num_groups % 2 == 0) {
+    groups_per_block = 2;
+  }
+
+  auto dst_type = output_q.scalar_type();
+  const int num_blocks = num_groups / groups_per_block;
+  const int num_threads = groups_per_block * THREADS_PER_GROUP;
+
+  const bool is_column_major = output_s.stride(0) < output_s.stride(1);
+  const int scale_num_rows = output_s.size(1);
+  const int scale_stride = output_s.stride(1);
+
+#define LAUNCH_KERNEL(T, DST_DTYPE)                                                               \
+  do {                                                                                            \
+    dim3 grid(num_blocks);                                                                        \
+    dim3 block(num_threads);                                                                      \
+    if (is_column_major) {                                                                        \
+      if (scale_ue8m0) {                                                                          \
+        per_token_group_quant_8bit_kernel<T, DST_DTYPE, true, true><<<grid, block, 0, stream>>>(  \
+            static_cast<T*>(input.data_ptr()),                                                    \
+            output_q.data_ptr(),                                                                  \
+            static_cast<uint32_t*>(output_s.data_ptr()),                                          \
+            group_size,                                                                           \
+            num_groups,                                                                           \
+            groups_per_block,                                                                     \
+            (float)eps,                                                                           \
+            (float)min_8bit,                                                                      \
+            (float)max_8bit,                                                                      \
+            scale_num_rows,                                                                       \
+            scale_stride);                                                                        \
+      } else {                                                                                    \
+        per_token_group_quant_8bit_kernel<T, DST_DTYPE, true, false><<<grid, block, 0, stream>>>( \
+            static_cast<T*>(input.data_ptr()),                                                    \
+            output_q.data_ptr(),                                                                  \
+            static_cast<float*>(output_s.data_ptr()),                                             \
+            group_size,                                                                           \
+            num_groups,                                                                           \
+            groups_per_block,                                                                     \
+            (float)eps,                                                                           \
+            (float)min_8bit,                                                                      \
+            (float)max_8bit,                                                                      \
+            scale_num_rows,                                                                       \
+            scale_stride);                                                                        \
+      }                                                                                           \
+    } else {                                                                                      \
+      assert(!scale_ue8m0);                                                                       \
+      per_token_group_quant_8bit_kernel<T, DST_DTYPE, false><<<grid, block, 0, stream>>>(         \
+          static_cast<T*>(input.data_ptr()),                                                      \
+          output_q.data_ptr(),                                                                    \
+          static_cast<float*>(output_s.data_ptr()),                                               \
+          group_size,                                                                             \
+          num_groups,                                                                             \
+          groups_per_block,                                                                       \
+          (float)eps,                                                                             \
+          (float)min_8bit,                                                                        \
+          (float)max_8bit);                                                                       \
+    }                                                                                             \
+  } while (0)
+
+  DISPATCH_PYTORCH_DTYPE_TO_CTYPE_FLOAT_FP16(input.scalar_type(), scalar_t, [&] {
+    if (dst_type == at::ScalarType::Char) {
+      LAUNCH_KERNEL(scalar_t, int8_t);
+      return true;
+    } else if (dst_type == at::ScalarType::Float8_e4m3fn) {
+      LAUNCH_KERNEL(scalar_t, c10::Float8_e4m3fn);
+      return true;
+    }
+    return false;
+  });
+
+#undef LAUNCH_KERNEL
+}
+
+void sgl_per_token_group_quant_int8(
+    torch::Tensor input,
+    torch::Tensor output_q,
+    torch::Tensor output_s,
+    int64_t group_size,
+    double eps,
+    double int8_min,
+    double int8_max) {
+  sgl_per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, int8_min, int8_max);
+}
+
+void sgl_per_token_group_quant_fp8(
+    torch::Tensor input,
+    torch::Tensor output_q,
+    torch::Tensor output_s,
+    int64_t group_size,
+    double eps,
+    double fp8_min,
+    double fp8_max,
+    bool scale_ue8m0) {
+  sgl_per_token_group_quant_8bit(input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0);
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/quantization/fp8/vec_dtypes.cuh b/vllm_musa/csrc/quantization/fp8/vec_dtypes.cuh
new file mode 100644
index 000000000..708f8b367
--- /dev/null
+++ b/vllm_musa/csrc/quantization/fp8/vec_dtypes.cuh
@@ -0,0 +1,1584 @@
+/*
+ * Copyright (c) 2023 by FlashInfer team.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+ #ifndef VEC_DTYPES_CUH_
+ #define VEC_DTYPES_CUH_
+
+ #include <musa_bf16.h>
+ #include <musa_fp16.h>
+ #include <musa_fp8.h>
+ #include <musa_runtime.h>
+
+ #include <type_traits>
+ #include "../../musa.h"
+
+ namespace flashinfer {
+
+ #if (!defined(__MUSA_ARCH__) || (__MUSA_ARCH__ >= 310))
+ #define FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+ #endif
+
+ #define FLASHINFER_INLINE inline __attribute__((always_inline)) __device__
+
+ __device__ __forceinline__ void st_global_release(int4 const& val, int4* addr) {
+   asm volatile("st.release.global.sys.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y),
+                "r"(val.z), "r"(val.w), "l"(addr));
+ }
+
+ __device__ __forceinline__ int4 ld_global_acquire(int4* addr) {
+   int4 val;
+   asm volatile("ld.acquire.global.sys.v4.b32 {%0, %1, %2, %3}, [%4];"
+                : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
+                : "l"(addr));
+   return val;
+ }
+
+ __device__ __forceinline__ void st_global_volatile(int4 const& val, int4* addr) {
+   asm volatile("st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y),
+                "r"(val.z), "r"(val.w), "l"(addr));
+ }
+
+ __device__ __forceinline__ int4 ld_global_volatile(int4* addr) {
+   int4 val;
+   asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];"
+                : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
+                : "l"(addr));
+   return val;
+ }
+
+//  #if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 < 120200) && \
+//      (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 800))
+//  // CUDA version < 12.2 and GPU architecture < 80
+//  FLASHINFER_INLINE __nv_bfloat162 make_bfloat162(const __nv_bfloat16 x, const __nv_bfloat16 y) {
+//    __nv_bfloat162 t;
+//    t.x = x;
+//    t.y = y;
+//    return t;
+//  }
+
+//  FLASHINFER_INLINE __nv_bfloat16 __hmul(const __nv_bfloat16 a, const __nv_bfloat16 b) {
+//    __nv_bfloat16 val;
+//    const float fa = __bfloat162float(a);
+//    const float fb = __bfloat162float(b);
+//    // avoid ftz in device code
+//    val = __float2bfloat16(__fmaf_ieee_rn(fa, fb, -0.0f));
+//    return val;
+//  }
+
+//  FLASHINFER_INLINE __nv_bfloat162 __hmul2(const __nv_bfloat162 a, const __nv_bfloat162 b) {
+//    __nv_bfloat162 val;
+//    val.x = __hmul(a.x, b.x);
+//    val.y = __hmul(a.y, b.y);
+//    return val;
+//  }
+
+//  FLASHINFER_INLINE __nv_bfloat162 __floats2bfloat162_rn(const float a, const float b) {
+//    __nv_bfloat162 val;
+//    val = __nv_bfloat162(__float2bfloat16_rn(a), __float2bfloat16_rn(b));
+//    return val;
+//  }
+
+//  FLASHINFER_INLINE __nv_bfloat162 __float22bfloat162_rn(const float2 a) {
+//    __nv_bfloat162 val = __floats2bfloat162_rn(a.x, a.y);
+//    return val;
+//  }
+//  FLASHINFER_INLINE float2 __bfloat1622float2(const __nv_bfloat162 a) {
+//    float hi_float;
+//    float lo_float;
+//    lo_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).x);
+//    hi_float = __internal_bfloat162float(((__nv_bfloat162_raw)a).y);
+//    return make_float2(lo_float, hi_float);
+//  }
+//  #endif
+
+ /******************* vec_t type cast *******************/
+
+ template <typename dst_t, typename src_t>
+ struct vec_cast {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(dst_t* dst, const src_t* src) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size; ++i) {
+       dst[i] = (dst_t)src[i];
+     }
+   }
+ };
+
+ template <>
+ struct vec_cast<float, half> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(float* dst, const half* src) {
+     if constexpr (vec_size == 1) {
+       dst[0] = (float)src[0];
+     } else {
+ #pragma unroll
+       for (size_t i = 0; i < vec_size / 2; ++i) {
+         ((float2*)dst)[i] = __half22float2(((half2*)src)[i]);
+       }
+     }
+   }
+ };
+
+ template <>
+ struct vec_cast<half, float> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(half* dst, const float* src) {
+     if constexpr (vec_size == 1) {
+       dst[0] = __float2half(src[0]);
+     } else {
+ #pragma unroll
+       for (size_t i = 0; i < vec_size / 2; ++i) {
+         ((half2*)dst)[i] = __float22half2_rn(((float2*)src)[i]);
+       }
+     }
+   }
+ };
+
+ template <typename T>
+ constexpr FLASHINFER_INLINE int get_exponent_bits() {
+   if constexpr (std::is_same_v<T, __nv_fp8_e4m3>) {
+     return 4;
+   } else if constexpr (std::is_same_v<T, __nv_fp8_e5m2>) {
+     return 5;
+   } else if constexpr (std::is_same_v<T, half>) {
+     return 5;
+   } else if constexpr (std::is_same_v<T, nv_bfloat16>) {
+     return 8;
+   }
+ }
+
+ template <typename T>
+ constexpr FLASHINFER_INLINE int get_mantissa_bits() {
+   if constexpr (std::is_same_v<T, __nv_fp8_e4m3>) {
+     return 3;
+   } else if constexpr (std::is_same_v<T, __nv_fp8_e5m2>) {
+     return 2;
+   } else if constexpr (std::is_same_v<T, half>) {
+     return 11;
+   } else if constexpr (std::is_same_v<T, nv_bfloat16>) {
+     return 7;
+   }
+ }
+
+ /*!
+  * \brief Fallback to software fast dequant implementation if hardware dequantization is not
+  * available.
+  * \note Inspired by Marlin's fast dequantization, but here we don't have to permute
+  * weights order.
+  * \ref
+  * https://github.com/vllm-project/vllm/blob/6dffa4b0a6120159ef2fe44d695a46817aff65bc/csrc/quantization/fp8/fp8_marlin.cu#L120
+  */
+ template <typename fp8_dtype, typename fp16_dtype>
+ __device__ void fast_dequant_f8f16x4(uint32_t* input, uint2* output) {
+   uint32_t q = *input;
+   if constexpr (std::is_same_v<fp8_dtype, __nv_fp8_e5m2> && std::is_same_v<fp16_dtype, half>) {
+     output->x = __byte_perm(0U, q, 0x5140);
+     output->y = __byte_perm(0U, q, 0x7362);
+   } else {
+     constexpr int FP8_EXPONENT = get_exponent_bits<fp8_dtype>();
+     constexpr int FP8_MANTISSA = get_mantissa_bits<fp8_dtype>();
+     constexpr int FP16_EXPONENT = get_exponent_bits<fp16_dtype>();
+
+     constexpr int RIGHT_SHIFT = FP16_EXPONENT - FP8_EXPONENT;
+     // Calculate MASK for extracting mantissa and exponent
+     constexpr int MASK1 = 0x80000000;
+     constexpr int MASK2 = MASK1 >> (FP8_EXPONENT + FP8_MANTISSA);
+     constexpr int MASK3 = MASK2 & 0x7fffffff;
+     constexpr int MASK = MASK3 | (MASK3 >> 16);
+     q = __byte_perm(q, q, 0x1302);
+
+     // Extract and shift FP8 values to FP16 format
+     uint32_t Out1 = (q & 0x80008000) | ((q & MASK) >> RIGHT_SHIFT);
+     uint32_t Out2 = ((q << 8) & 0x80008000) | (((q << 8) & MASK) >> RIGHT_SHIFT);
+
+     constexpr int BIAS_OFFSET = (1 << (FP16_EXPONENT - 1)) - (1 << (FP8_EXPONENT - 1));
+     // Construct and apply exponent bias
+     if constexpr (std::is_same_v<fp16_dtype, half>) {
+       const half2 bias_reg = __float2half2_rn(float(1 << BIAS_OFFSET));
+
+       // Convert to half2 and apply bias
+       *(half2*)&(output->x) = __hmul2(*reinterpret_cast<const half2*>(&Out1), bias_reg);
+       *(half2*)&(output->y) = __hmul2(*reinterpret_cast<const half2*>(&Out2), bias_reg);
+     } else {
+       constexpr uint32_t BIAS = (BIAS_OFFSET + 127) << 23;
+       const nv_bfloat162 bias_reg = __float2bfloat162_rn(*reinterpret_cast<const float*>(&BIAS));
+       // Convert to bfloat162 and apply bias
+       *(nv_bfloat162*)&(output->x) =
+           __hmul2(*reinterpret_cast<const nv_bfloat162*>(&Out1), bias_reg);
+       *(nv_bfloat162*)&(output->y) =
+           __hmul2(*reinterpret_cast<const nv_bfloat162*>(&Out2), bias_reg);
+     }
+   }
+ }
+
+ template <>
+ struct vec_cast<nv_bfloat16, __nv_fp8_e4m3> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e4m3* src) {
+     if constexpr (vec_size == 1) {
+       dst[0] = nv_bfloat16(src[0]);
+     } else if constexpr (vec_size == 2) {
+       dst[0] = nv_bfloat16(src[0]);
+       dst[1] = nv_bfloat16(src[1]);
+     } else {
+       static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4");
+ #pragma unroll
+       for (uint32_t i = 0; i < vec_size / 4; ++i) {
+         fast_dequant_f8f16x4<__nv_fp8_e4m3, nv_bfloat16>((uint32_t*)&src[i * 4],
+                                                          (uint2*)&dst[i * 4]);
+       }
+     }
+   }
+ };
+
+ template <>
+ struct vec_cast<nv_bfloat16, __nv_fp8_e5m2> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const __nv_fp8_e5m2* src) {
+     if constexpr (vec_size == 1) {
+       dst[0] = nv_bfloat16(src[0]);
+     } else if constexpr (vec_size == 2) {
+       dst[0] = nv_bfloat16(src[0]);
+       dst[1] = nv_bfloat16(src[1]);
+     } else {
+       static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4");
+ #pragma unroll
+       for (uint32_t i = 0; i < vec_size / 4; ++i) {
+         fast_dequant_f8f16x4<__nv_fp8_e5m2, nv_bfloat16>((uint32_t*)&src[i * 4],
+                                                          (uint2*)&dst[i * 4]);
+       }
+     }
+   }
+ };
+
+ template <>
+ struct vec_cast<__nv_fp8_e4m3, half> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(__nv_fp8_e4m3* dst, const half* src) {
+ #ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+     if constexpr (vec_size == 1) {
+       dst[0] = __nv_fp8_e4m3(src[0]);
+     } else {
+ #pragma unroll
+       for (size_t i = 0; i < vec_size / 2; ++i) {
+         uint16_t y;
+         uint32_t x = *(uint32_t*)&src[i * 2];
+         asm volatile("cvt.rn.satfinite.e4m3x2.f16x2 %0, %1;" : "=h"(y) : "r"(x));
+         *(uint16_t*)&dst[i * 2] = y;
+       }
+     }
+ #else
+ #pragma unroll
+     for (size_t i = 0; i < vec_size; ++i) {
+       dst[i] = __nv_fp8_e4m3(src[i]);
+     }
+ #endif  // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+   }
+ };
+
+ template <>
+ struct vec_cast<__nv_fp8_e5m2, half> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(__nv_fp8_e5m2* dst, const half* src) {
+ #ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+     if constexpr (vec_size == 1) {
+       dst[0] = __nv_fp8_e5m2(src[0]);
+     } else {
+ #pragma unroll
+       for (size_t i = 0; i < vec_size / 2; ++i) {
+         uint16_t y;
+         uint32_t x = *(uint32_t*)&src[i * 2];
+         asm volatile("cvt.rn.satfinite.e5m2x2.f16x2 %0, %1;" : "=h"(y) : "r"(x));
+         *(uint16_t*)&dst[i * 2] = y;
+       }
+     }
+ #else
+ #pragma unroll
+     for (size_t i = 0; i < vec_size; ++i) {
+       dst[i] = __nv_fp8_e5m2(src[i]);
+     }
+ #endif  // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+   }
+ };
+
+ template <>
+ struct vec_cast<half, __nv_fp8_e4m3> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e4m3* src) {
+ #ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+     if constexpr (vec_size == 1) {
+       dst[0] = half(src[0]);
+     } else {
+ #pragma unroll
+       for (size_t i = 0; i < vec_size / 2; ++i) {
+         uint32_t y;
+         uint16_t x = *(uint16_t*)&src[i * 2];
+         asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" : "=r"(y) : "h"(x));
+         *(uint32_t*)&dst[i * 2] = y;
+       }
+     }
+ #else
+     if constexpr (vec_size == 1) {
+       dst[0] = half(src[0]);
+     } else if constexpr (vec_size == 2) {
+       dst[0] = half(src[0]);
+       dst[1] = half(src[1]);
+     } else {
+       static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4");
+ #pragma unroll
+       for (uint32_t i = 0; i < vec_size / 4; ++i) {
+         fast_dequant_f8f16x4<__nv_fp8_e4m3, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]);
+       }
+     }
+ #endif  // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+   }
+ };
+
+ template <>
+ struct vec_cast<half, __nv_fp8_e5m2> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(half* dst, const __nv_fp8_e5m2* src) {
+ #ifdef FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+     if constexpr (vec_size == 1) {
+       dst[0] = half(src[0]);
+     } else {
+ #pragma unroll
+       for (size_t i = 0; i < vec_size / 2; ++i) {
+         uint32_t y;
+         uint16_t x = *(uint16_t*)&src[i * 2];
+         asm volatile("cvt.rn.f16x2.e5m2x2 %0, %1;" : "=r"(y) : "h"(x));
+         *(uint32_t*)&dst[i * 2] = y;
+       }
+     }
+ #else
+     if constexpr (vec_size == 1) {
+       dst[0] = half(src[0]);
+     } else if constexpr (vec_size == 2) {
+       dst[0] = half(src[0]);
+       dst[1] = half(src[1]);
+     } else {
+       static_assert(vec_size % 4 == 0, "vec_size must be a multiple of 4");
+ #pragma unroll
+       for (uint32_t i = 0; i < vec_size / 4; ++i) {
+         fast_dequant_f8f16x4<__nv_fp8_e5m2, half>((uint32_t*)&src[i * 4], (uint2*)&dst[i * 4]);
+       }
+     }
+ #endif  // FLASHINFER_HARDWARE_FP8_CONVERSION_ENABLED
+   }
+ };
+
+ template <>
+ struct vec_cast<float, nv_bfloat16> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(float* dst, const nv_bfloat16* src) {
+     if constexpr (vec_size == 1) {
+       dst[0] = (float)src[0];
+     } else {
+ #pragma unroll
+       for (size_t i = 0; i < vec_size / 2; ++i) {
+         ((float2*)dst)[i] = __bfloat1622float2(((nv_bfloat162*)src)[i]);
+       }
+     }
+   }
+ };
+
+ template <>
+ struct vec_cast<nv_bfloat16, float> {
+   template <size_t vec_size>
+   FLASHINFER_INLINE static void cast(nv_bfloat16* dst, const float* src) {
+     if constexpr (vec_size == 1) {
+       dst[0] = nv_bfloat16(src[0]);
+     } else {
+ #pragma unroll
+       for (size_t i = 0; i < vec_size / 2; ++i) {
+         ((nv_bfloat162*)dst)[i] = __float22bfloat162_rn(((float2*)src)[i]);
+       }
+     }
+   }
+ };
+
+ template <typename float_t, size_t vec_size>
+ struct vec_t {
+   FLASHINFER_INLINE float_t& operator[](size_t i);
+   FLASHINFER_INLINE const float_t& operator[](size_t i) const;
+   FLASHINFER_INLINE void fill(float_t val);
+   FLASHINFER_INLINE void load(const float_t* ptr);
+   FLASHINFER_INLINE void store(float_t* ptr) const;
+   FLASHINFER_INLINE void load_global_acquire(float* addr);
+   FLASHINFER_INLINE void store_global_release(float* addr) const;
+   FLASHINFER_INLINE void load_global_volatile(float* addr);
+   FLASHINFER_INLINE void store_global_volatile(float* addr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, vec_size>& src);
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr);
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const;
+   FLASHINFER_INLINE static void memcpy(float_t* dst, const float_t* src);
+   FLASHINFER_INLINE float_t* ptr();
+ };
+
+ template <typename src_float_t, typename tgt_float_t, size_t vec_size>
+ FLASHINFER_INLINE void cast_from_impl(vec_t<tgt_float_t, vec_size>& dst,
+                                       const vec_t<src_float_t, vec_size>& src) {
+   vec_cast<tgt_float_t, src_float_t>::cast<vec_size>(
+       dst.ptr(), const_cast<vec_t<src_float_t, vec_size>*>(&src)->ptr());
+ }
+
+ template <typename src_float_t, typename tgt_float_t, size_t vec_size>
+ FLASHINFER_INLINE void cast_load_impl(vec_t<tgt_float_t, vec_size>& dst,
+                                       const src_float_t* src_ptr) {
+   if constexpr (std::is_same_v<src_float_t, tgt_float_t>) {
+     dst.load(src_ptr);
+   } else {
+     vec_t<src_float_t, vec_size> tmp;
+     tmp.load(src_ptr);
+     dst.cast_from(tmp);
+   }
+ }
+
+ template <typename src_float_t, typename tgt_float_t, size_t vec_size>
+ FLASHINFER_INLINE void cast_store_impl(tgt_float_t* dst_ptr,
+                                        const vec_t<src_float_t, vec_size>& src) {
+   if constexpr (std::is_same_v<src_float_t, tgt_float_t>) {
+     src.store(dst_ptr);
+   } else {
+     vec_t<tgt_float_t, vec_size> tmp;
+     tmp.cast_from(src);
+     tmp.store(dst_ptr);
+   }
+ }
+
+ /******************* vec_t<__nv_fp8_e4m3> *******************/
+
+ // __nv_fp8_e4m3 x 1
+ template <>
+ struct vec_t<__nv_fp8_e4m3, 1> {
+   __nv_fp8_e4m3 data;
+
+   FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const {
+     return ((const __nv_fp8_e4m3*)(&data))[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val);
+   FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr);
+   FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 1>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::fill(__nv_fp8_e4m3 val) { data = val; }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::load(const __nv_fp8_e4m3* ptr) { data = *ptr; }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::store(__nv_fp8_e4m3* ptr) const { *ptr = data; }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 1>::memcpy(__nv_fp8_e4m3* dst,
+                                                        const __nv_fp8_e4m3* src) {
+   *dst = *src;
+ }
+
+ // __nv_fp8_e4m3 x 2
+ template <>
+ struct vec_t<__nv_fp8_e4m3, 2> {
+   __nv_fp8x2_e4m3 data;
+
+   FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const {
+     return ((const __nv_fp8_e4m3*)(&data))[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val);
+   FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr);
+   FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 2>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::fill(__nv_fp8_e4m3 val) {
+   data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::load(const __nv_fp8_e4m3* ptr) {
+   data = *((__nv_fp8x2_e4m3*)ptr);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::store(__nv_fp8_e4m3* ptr) const {
+   *((__nv_fp8x2_e4m3*)ptr) = data;
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 2>::memcpy(__nv_fp8_e4m3* dst,
+                                                        const __nv_fp8_e4m3* src) {
+   *((__nv_fp8x2_e4m3*)dst) = *((__nv_fp8x2_e4m3*)src);
+ }
+
+ // __nv_fp8_e4m3 x 4
+
+ template <>
+ struct vec_t<__nv_fp8_e4m3, 4> {
+   __nv_fp8x4_e4m3 data;
+
+   FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const {
+     return ((const __nv_fp8_e4m3*)(&data))[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val);
+   FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr);
+   FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 4>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::fill(__nv_fp8_e4m3 val) {
+   data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+              (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::load(const __nv_fp8_e4m3* ptr) {
+   data = *((__nv_fp8x4_e4m3*)ptr);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::store(__nv_fp8_e4m3* ptr) const {
+   *((__nv_fp8x4_e4m3*)ptr) = data;
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 4>::memcpy(__nv_fp8_e4m3* dst,
+                                                        const __nv_fp8_e4m3* src) {
+   *((__nv_fp8x4_e4m3*)dst) = *((__nv_fp8x4_e4m3*)src);
+ }
+
+ // __nv_fp8_e4m3 x 8
+
+ template <>
+ struct vec_t<__nv_fp8_e4m3, 8> {
+   uint2 data;
+
+   FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)(&data))[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const {
+     return ((const __nv_fp8_e4m3*)(&data))[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val);
+   FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr);
+   FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 8>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::fill(__nv_fp8_e4m3 val) {
+   ((__nv_fp8x4_e4m3*)(&data.x))->__x =
+       (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+       (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+   ((__nv_fp8x4_e4m3*)(&data.y))->__x =
+       (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+       (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::load(const __nv_fp8_e4m3* ptr) {
+   data = *((uint2*)ptr);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::store(__nv_fp8_e4m3* ptr) const {
+   *((uint2*)ptr) = data;
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e4m3, 8>::memcpy(__nv_fp8_e4m3* dst,
+                                                        const __nv_fp8_e4m3* src) {
+   *((uint2*)dst) = *((uint2*)src);
+ }
+
+ // __nv_fp8_e4m3 x 16 or more
+ template <size_t vec_size>
+ struct vec_t<__nv_fp8_e4m3, vec_size> {
+   static_assert(vec_size % 16 == 0, "Invalid vector size");
+   int4 data[vec_size / 16];
+
+   FLASHINFER_INLINE __nv_fp8_e4m3& operator[](size_t i) { return ((__nv_fp8_e4m3*)data)[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e4m3& operator[](size_t i) const {
+     return ((const __nv_fp8_e4m3*)data)[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e4m3* ptr() { return reinterpret_cast<__nv_fp8_e4m3*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e4m3 val) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       ((__nv_fp8x4_e4m3*)(&(data[i].x)))->__x =
+           (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+           (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+       ((__nv_fp8x4_e4m3*)(&(data[i].y)))->__x =
+           (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+           (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+       ((__nv_fp8x4_e4m3*)(&(data[i].z)))->__x =
+           (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+           (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+       ((__nv_fp8x4_e4m3*)(&(data[i].w)))->__x =
+           (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+           (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+     }
+   }
+   FLASHINFER_INLINE void load(const __nv_fp8_e4m3* ptr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       data[i] = ((int4*)ptr)[i];
+     }
+   }
+   FLASHINFER_INLINE void store(__nv_fp8_e4m3* ptr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       ((int4*)ptr)[i] = data[i];
+     }
+   }
+   FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e4m3* addr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       *((int4*)(data + i)) = ld_global_acquire((int4*)(addr + i * 16));
+     }
+   }
+   FLASHINFER_INLINE void store_global_release(__nv_fp8_e4m3* addr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       st_global_release(data[i], (int4*)(addr + i * 16));
+     }
+   }
+   FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e4m3* addr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       data[i] = ld_global_volatile((int4*)(addr + i * 16));
+     }
+   }
+   FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e4m3* addr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       st_global_volatile(data[i], (int4*)(addr + i * 16));
+     }
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, vec_size>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e4m3* dst, const __nv_fp8_e4m3* src) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       ((int4*)dst)[i] = ((int4*)src)[i];
+     }
+   }
+ };
+
+ /******************* vec_t<__nv_fp8_e5m2> *******************/
+
+ // __nv_fp8_e5m2 x 1
+ template <>
+ struct vec_t<__nv_fp8_e5m2, 1> {
+   __nv_fp8_e5m2 data;
+
+   FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const {
+     return ((const __nv_fp8_e5m2*)(&data))[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val);
+   FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr);
+   FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 1>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::fill(__nv_fp8_e5m2 val) { data = val; }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::load(const __nv_fp8_e5m2* ptr) { data = *ptr; }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::store(__nv_fp8_e5m2* ptr) const { *ptr = data; }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 1>::memcpy(__nv_fp8_e5m2* dst,
+                                                        const __nv_fp8_e5m2* src) {
+   *dst = *src;
+ }
+
+ // __nv_fp8_e5m2 x 2
+ template <>
+ struct vec_t<__nv_fp8_e5m2, 2> {
+   __nv_fp8x2_e5m2 data;
+
+   FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const {
+     return ((const __nv_fp8_e5m2*)(&data))[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val);
+   FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr);
+   FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 2>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::fill(__nv_fp8_e5m2 val) {
+   data.__x = (__nv_fp8x2_storage_t(val.__x) << 8) | __nv_fp8x2_storage_t(val.__x);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::load(const __nv_fp8_e5m2* ptr) {
+   data = *((__nv_fp8x2_e5m2*)ptr);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::store(__nv_fp8_e5m2* ptr) const {
+   *((__nv_fp8x2_e5m2*)ptr) = data;
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 2>::memcpy(__nv_fp8_e5m2* dst,
+                                                        const __nv_fp8_e5m2* src) {
+   *((__nv_fp8x2_e5m2*)dst) = *((__nv_fp8x2_e5m2*)src);
+ }
+
+ // __nv_fp8_e5m2 x 4
+
+ template <>
+ struct vec_t<__nv_fp8_e5m2, 4> {
+   __nv_fp8x4_e5m2 data;
+
+   FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const {
+     return ((const __nv_fp8_e5m2*)(&data))[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val);
+   FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr);
+   FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 4>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::fill(__nv_fp8_e5m2 val) {
+   data.__x = (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+              (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::load(const __nv_fp8_e5m2* ptr) {
+   data = *((__nv_fp8x4_e5m2*)ptr);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::store(__nv_fp8_e5m2* ptr) const {
+   *((__nv_fp8x4_e5m2*)ptr) = data;
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 4>::memcpy(__nv_fp8_e5m2* dst,
+                                                        const __nv_fp8_e5m2* src) {
+   *((__nv_fp8x4_e5m2*)dst) = *((__nv_fp8x4_e5m2*)src);
+ }
+
+ // __nv_fp8_e5m2 x 8
+
+ template <>
+ struct vec_t<__nv_fp8_e5m2, 8> {
+   uint2 data;
+
+   FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)(&data))[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const {
+     return ((const __nv_fp8_e5m2*)(&data))[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val);
+   FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr);
+   FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 8>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::fill(__nv_fp8_e5m2 val) {
+   ((__nv_fp8x4_e5m2*)(&data.x))->__x =
+       (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+       (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+   ((__nv_fp8x4_e5m2*)(&data.y))->__x =
+       (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+       (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::load(const __nv_fp8_e5m2* ptr) {
+   data = *((uint2*)ptr);
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::store(__nv_fp8_e5m2* ptr) const {
+   *((uint2*)ptr) = data;
+ }
+
+ FLASHINFER_INLINE void vec_t<__nv_fp8_e5m2, 8>::memcpy(__nv_fp8_e5m2* dst,
+                                                        const __nv_fp8_e5m2* src) {
+   *((uint2*)dst) = *((uint2*)src);
+ }
+
+ // __nv_fp8_e5m2 x 16 or more
+
+ template <size_t vec_size>
+ struct vec_t<__nv_fp8_e5m2, vec_size> {
+   static_assert(vec_size % 16 == 0, "Invalid vector size");
+   int4 data[vec_size / 16];
+
+   FLASHINFER_INLINE __nv_fp8_e5m2& operator[](size_t i) { return ((__nv_fp8_e5m2*)data)[i]; }
+   FLASHINFER_INLINE const __nv_fp8_e5m2& operator[](size_t i) const {
+     return ((const __nv_fp8_e5m2*)data)[i];
+   }
+   FLASHINFER_INLINE __nv_fp8_e5m2* ptr() { return reinterpret_cast<__nv_fp8_e5m2*>(&data); }
+   FLASHINFER_INLINE void fill(__nv_fp8_e5m2 val) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       ((__nv_fp8x4_e5m2*)(&(data[i].x)))->__x =
+           (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+           (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+       ((__nv_fp8x4_e5m2*)(&(data[i].y)))->__x =
+           (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+           (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+       ((__nv_fp8x4_e5m2*)(&(data[i].z)))->__x =
+           (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+           (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+       ((__nv_fp8x4_e5m2*)(&(data[i].w)))->__x =
+           (__nv_fp8x4_storage_t(val.__x) << 24) | (__nv_fp8x4_storage_t(val.__x) << 16) |
+           (__nv_fp8x4_storage_t(val.__x) << 8) | __nv_fp8x4_storage_t(val.__x);
+     }
+   }
+   FLASHINFER_INLINE void load(const __nv_fp8_e5m2* ptr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       data[i] = ((int4*)ptr)[i];
+     }
+   }
+   FLASHINFER_INLINE void store(__nv_fp8_e5m2* ptr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       ((int4*)ptr)[i] = data[i];
+     }
+   }
+   FLASHINFER_INLINE void store_global_release(__nv_fp8_e5m2* addr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       st_global_release(data[i], (int4*)(addr + i * 16));
+     }
+   }
+   FLASHINFER_INLINE void load_global_acquire(__nv_fp8_e5m2* addr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       data[i] = ld_global_acquire((int4*)(addr + i * 16));
+     }
+   }
+   FLASHINFER_INLINE void store_global_volatile(__nv_fp8_e5m2* addr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       st_global_volatile(data[i], (int4*)(addr + i * 16));
+     }
+   }
+   FLASHINFER_INLINE void load_global_volatile(__nv_fp8_e5m2* addr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       data[i] = ld_global_volatile((int4*)(addr + i * 16));
+     }
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, vec_size>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(__nv_fp8_e5m2* dst, const __nv_fp8_e5m2* src) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 16; ++i) {
+       ((int4*)dst)[i] = ((int4*)src)[i];
+     }
+   }
+ };
+
+ /******************* vec_t<half> *******************/
+
+ // half x 1
+ template <>
+ struct vec_t<half, 1> {
+   half data;
+
+   FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; }
+   FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; }
+   FLASHINFER_INLINE half* ptr() { return reinterpret_cast<half*>(&data); }
+   FLASHINFER_INLINE void fill(half val);
+   FLASHINFER_INLINE void load(const half* ptr);
+   FLASHINFER_INLINE void store(half* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 1>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(half* dst, const half* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<half, 1>::fill(half val) { data = val; }
+
+ FLASHINFER_INLINE void vec_t<half, 1>::load(const half* ptr) { data = *ptr; }
+
+ FLASHINFER_INLINE void vec_t<half, 1>::store(half* ptr) const { *ptr = data; }
+
+ FLASHINFER_INLINE void vec_t<half, 1>::memcpy(half* dst, const half* src) { *dst = *src; }
+
+ // half x 2
+ template <>
+ struct vec_t<half, 2> {
+   half2 data;
+
+   FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; }
+   FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; }
+   FLASHINFER_INLINE half* ptr() { return reinterpret_cast<half*>(&data); }
+   FLASHINFER_INLINE void fill(half val);
+   FLASHINFER_INLINE void load(const half* ptr);
+   FLASHINFER_INLINE void store(half* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 2>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+
+   FLASHINFER_INLINE static void memcpy(half* dst, const half* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<half, 2>::fill(half val) { data = make_half2(val, val); }
+
+ FLASHINFER_INLINE void vec_t<half, 2>::load(const half* ptr) { data = *((half2*)ptr); }
+
+ FLASHINFER_INLINE void vec_t<half, 2>::store(half* ptr) const { *((half2*)ptr) = data; }
+
+ FLASHINFER_INLINE void vec_t<half, 2>::memcpy(half* dst, const half* src) {
+   *((half2*)dst) = *((half2*)src);
+ }
+
+ // half x 4
+
+ template <>
+ struct vec_t<half, 4> {
+   uint2 data;
+
+   FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)(&data))[i]; }
+   FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)(&data))[i]; }
+   FLASHINFER_INLINE half* ptr() { return reinterpret_cast<half*>(&data); }
+   FLASHINFER_INLINE void fill(half val);
+   FLASHINFER_INLINE void load(const half* ptr);
+   FLASHINFER_INLINE void store(half* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 4>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(half* dst, const half* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<half, 4>::fill(half val) {
+   *(half2*)(&data.x) = make_half2(val, val);
+   *(half2*)(&data.y) = make_half2(val, val);
+ }
+
+ FLASHINFER_INLINE void vec_t<half, 4>::load(const half* ptr) { data = *((uint2*)ptr); }
+
+ FLASHINFER_INLINE void vec_t<half, 4>::store(half* ptr) const { *((uint2*)ptr) = data; }
+
+ FLASHINFER_INLINE void vec_t<half, 4>::memcpy(half* dst, const half* src) {
+   *((uint2*)dst) = *((uint2*)src);
+ }
+
+ // half x 8 or more
+
+ template <size_t vec_size>
+ struct vec_t<half, vec_size> {
+   static_assert(vec_size % 8 == 0, "Invalid vector size");
+   int4 data[vec_size / 8];
+   FLASHINFER_INLINE half& operator[](size_t i) { return ((half*)data)[i]; }
+   FLASHINFER_INLINE const half& operator[](size_t i) const { return ((const half*)data)[i]; }
+   FLASHINFER_INLINE half* ptr() { return reinterpret_cast<half*>(&data); }
+   FLASHINFER_INLINE void fill(half val) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       *(half2*)(&(data[i].x)) = make_half2(val, val);
+       *(half2*)(&(data[i].y)) = make_half2(val, val);
+       *(half2*)(&(data[i].z)) = make_half2(val, val);
+       *(half2*)(&(data[i].w)) = make_half2(val, val);
+     }
+   }
+   FLASHINFER_INLINE void load(const half* ptr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       data[i] = ((int4*)ptr)[i];
+     }
+   }
+   FLASHINFER_INLINE void store(half* ptr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       ((int4*)ptr)[i] = data[i];
+     }
+   }
+   FLASHINFER_INLINE void load_global_acquire(half* addr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       data[i] = ld_global_acquire((int4*)(addr + i * 8));
+     }
+   }
+   FLASHINFER_INLINE void store_global_release(half* addr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       st_global_release(data[i], (int4*)(addr + i * 8));
+     }
+   }
+   FLASHINFER_INLINE void store_global_volatile(half* addr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       st_global_volatile(data[i], (int4*)(addr + i * 8));
+     }
+   }
+   FLASHINFER_INLINE void load_global_volatile(half* addr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       data[i] = ld_global_volatile((int4*)(addr + i * 8));
+     }
+   }
+
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, vec_size>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(half* dst, const half* src) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       ((int4*)dst)[i] = ((int4*)src)[i];
+     }
+   }
+ };
+
+ /******************* vec_t<nv_bfloat16> *******************/
+
+ // nv_bfloat16 x 1
+ template <>
+ struct vec_t<nv_bfloat16, 1> {
+   nv_bfloat16 data;
+   FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; }
+   FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const {
+     return ((const nv_bfloat16*)(&data))[i];
+   }
+   FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast<nv_bfloat16*>(&data); }
+   FLASHINFER_INLINE void fill(nv_bfloat16 val);
+   FLASHINFER_INLINE void load(const nv_bfloat16* ptr);
+   FLASHINFER_INLINE void store(nv_bfloat16* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 1>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 1>::fill(nv_bfloat16 val) { data = val; }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 1>::load(const nv_bfloat16* ptr) { data = *ptr; }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 1>::store(nv_bfloat16* ptr) const { *ptr = data; }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 1>::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) {
+   *dst = *src;
+ }
+
+ // nv_bfloat16 x 2
+ template <>
+ struct vec_t<nv_bfloat16, 2> {
+   nv_bfloat162 data;
+
+   FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; }
+   FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const {
+     return ((const nv_bfloat16*)(&data))[i];
+   }
+   FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast<nv_bfloat16*>(&data); }
+   FLASHINFER_INLINE void fill(nv_bfloat16 val);
+   FLASHINFER_INLINE void load(const nv_bfloat16* ptr);
+   FLASHINFER_INLINE void store(nv_bfloat16* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 2>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 2>::fill(nv_bfloat16 val) {
+   data = make_bfloat162(val, val);
+ }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 2>::load(const nv_bfloat16* ptr) {
+   data = *((nv_bfloat162*)ptr);
+ }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 2>::store(nv_bfloat16* ptr) const {
+   *((nv_bfloat162*)ptr) = data;
+ }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 2>::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) {
+   *((nv_bfloat162*)dst) = *((nv_bfloat162*)src);
+ }
+
+ // nv_bfloat16 x 4
+
+ template <>
+ struct vec_t<nv_bfloat16, 4> {
+   uint2 data;
+
+   FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)(&data))[i]; }
+   FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const {
+     return ((const nv_bfloat16*)(&data))[i];
+   }
+   FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast<nv_bfloat16*>(&data); }
+   FLASHINFER_INLINE void fill(nv_bfloat16 val);
+   FLASHINFER_INLINE void load(const nv_bfloat16* ptr);
+   FLASHINFER_INLINE void store(nv_bfloat16* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 4>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 4>::fill(nv_bfloat16 val) {
+   *(nv_bfloat162*)(&data.x) = make_bfloat162(val, val);
+   *(nv_bfloat162*)(&data.y) = make_bfloat162(val, val);
+ }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 4>::load(const nv_bfloat16* ptr) {
+   data = *((uint2*)ptr);
+ }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 4>::store(nv_bfloat16* ptr) const {
+   *((uint2*)ptr) = data;
+ }
+
+ FLASHINFER_INLINE void vec_t<nv_bfloat16, 4>::memcpy(nv_bfloat16* dst, const nv_bfloat16* src) {
+   *((uint2*)dst) = *((uint2*)src);
+ }
+
+ // nv_bfloat16 x 8 or more
+
+ template <size_t vec_size>
+ struct vec_t<nv_bfloat16, vec_size> {
+   static_assert(vec_size % 8 == 0, "Invalid vector size");
+   int4 data[vec_size / 8];
+
+   FLASHINFER_INLINE nv_bfloat16& operator[](size_t i) { return ((nv_bfloat16*)data)[i]; }
+   FLASHINFER_INLINE const nv_bfloat16& operator[](size_t i) const {
+     return ((const nv_bfloat16*)data)[i];
+   }
+   FLASHINFER_INLINE nv_bfloat16* ptr() { return reinterpret_cast<nv_bfloat16*>(&data); }
+   FLASHINFER_INLINE void fill(nv_bfloat16 val) {
+ #pragma unoll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       *(nv_bfloat162*)(&(data[i].x)) = make_bfloat162(val, val);
+       *(nv_bfloat162*)(&(data[i].y)) = make_bfloat162(val, val);
+       *(nv_bfloat162*)(&(data[i].z)) = make_bfloat162(val, val);
+       *(nv_bfloat162*)(&(data[i].w)) = make_bfloat162(val, val);
+     }
+   }
+   FLASHINFER_INLINE void load(const nv_bfloat16* ptr) {
+ #pragma unoll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       data[i] = ((int4*)ptr)[i];
+     }
+   }
+   FLASHINFER_INLINE void store(nv_bfloat16* ptr) const {
+ #pragma unoll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       ((int4*)ptr)[i] = data[i];
+     }
+   }
+   FLASHINFER_INLINE void store_global_release(nv_bfloat16* addr) const {
+ #pragma unoll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       st_global_release(data[i], (int4*)(addr + i * 8));
+     }
+   }
+   FLASHINFER_INLINE void load_global_acquire(nv_bfloat16* addr) {
+ #pragma unoll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       data[i] = ld_global_acquire((int4*)(addr + i * 8));
+     }
+   }
+   FLASHINFER_INLINE void store_global_volatile(nv_bfloat16* addr) const {
+ #pragma unoll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       st_global_volatile(data[i], (int4*)(addr + i * 8));
+     }
+   }
+   FLASHINFER_INLINE void load_global_volatile(nv_bfloat16* addr) {
+ #pragma unoll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       data[i] = ld_global_volatile((int4*)(addr + i * 8));
+     }
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, vec_size>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(nv_bfloat16* dst, const nv_bfloat16* src) {
+ #pragma unoll
+     for (size_t i = 0; i < vec_size / 8; ++i) {
+       ((int4*)dst)[i] = ((int4*)src)[i];
+     }
+   }
+ };
+
+ /******************* vec_t<float> *******************/
+
+ // float x 1
+
+ template <>
+ struct vec_t<float, 1> {
+   float data;
+
+   FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(&data))[i]; }
+   FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(&data))[i]; }
+   FLASHINFER_INLINE float* ptr() { return reinterpret_cast<float*>(&data); }
+   FLASHINFER_INLINE void fill(float val);
+   FLASHINFER_INLINE void load(const float* ptr);
+   FLASHINFER_INLINE void store(float* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 1>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(float* dst, const float* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<float, 1>::fill(float val) { data = val; }
+
+ FLASHINFER_INLINE void vec_t<float, 1>::load(const float* ptr) { data = *ptr; }
+
+ FLASHINFER_INLINE void vec_t<float, 1>::store(float* ptr) const { *ptr = data; }
+
+ FLASHINFER_INLINE void vec_t<float, 1>::memcpy(float* dst, const float* src) { *dst = *src; }
+
+ // float x 2
+
+ template <>
+ struct vec_t<float, 2> {
+   float2 data;
+
+   FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(&data))[i]; }
+   FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(&data))[i]; }
+   FLASHINFER_INLINE float* ptr() { return reinterpret_cast<float*>(&data); }
+   FLASHINFER_INLINE void fill(float val);
+   FLASHINFER_INLINE void load(const float* ptr);
+   FLASHINFER_INLINE void store(float* ptr) const;
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, 2>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(float* dst, const float* src);
+ };
+
+ FLASHINFER_INLINE void vec_t<float, 2>::fill(float val) { data = make_float2(val, val); }
+
+ FLASHINFER_INLINE void vec_t<float, 2>::load(const float* ptr) { data = *((float2*)ptr); }
+
+ FLASHINFER_INLINE void vec_t<float, 2>::store(float* ptr) const { *((float2*)ptr) = data; }
+
+ FLASHINFER_INLINE void vec_t<float, 2>::memcpy(float* dst, const float* src) {
+   *((float2*)dst) = *((float2*)src);
+ }
+
+ // float x 4 or more
+ template <size_t vec_size>
+ struct vec_t<float, vec_size> {
+   static_assert(vec_size % 4 == 0, "Invalid vector size");
+   float4 data[vec_size / 4];
+
+   FLASHINFER_INLINE float& operator[](size_t i) { return ((float*)(data))[i]; }
+   FLASHINFER_INLINE const float& operator[](size_t i) const { return ((const float*)(data))[i]; }
+   FLASHINFER_INLINE float* ptr() { return reinterpret_cast<float*>(&data); }
+   FLASHINFER_INLINE void fill(float val) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 4; ++i) {
+       data[i] = make_float4(val, val, val, val);
+     }
+   }
+   FLASHINFER_INLINE void load(const float* ptr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 4; ++i) {
+       data[i] = ((float4*)ptr)[i];
+     }
+   }
+   FLASHINFER_INLINE void store(float* ptr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 4; ++i) {
+       ((float4*)ptr)[i] = data[i];
+     }
+   }
+   FLASHINFER_INLINE void store_global_release(float* addr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 4; ++i) {
+       st_global_release(*(int4*)(data + i), (int4*)(addr + i * 4));
+     }
+   }
+   FLASHINFER_INLINE void load_global_acquire(float* addr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 4; ++i) {
+       *((int4*)(data + i)) = ld_global_acquire((int4*)(addr + i * 4));
+     }
+   }
+   FLASHINFER_INLINE void store_global_volatile(float* addr) const {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 4; ++i) {
+       st_global_volatile(*(int4*)(data + i), (int4*)(addr + i * 4));
+     }
+   }
+   FLASHINFER_INLINE void load_global_volatile(float* addr) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 4; ++i) {
+       *((int4*)(data + i)) = ld_global_volatile((int4*)(addr + i * 4));
+     }
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_from(const vec_t<T, vec_size>& src) {
+     cast_from_impl(*this, src);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_load(const T* ptr) {
+     cast_load_impl(*this, ptr);
+   }
+   template <typename T>
+   FLASHINFER_INLINE void cast_store(T* ptr) const {
+     cast_store_impl(ptr, *this);
+   }
+   FLASHINFER_INLINE static void memcpy(float* dst, const float* src) {
+ #pragma unroll
+     for (size_t i = 0; i < vec_size / 4; ++i) {
+       ((float4*)dst)[i] = ((float4*)src)[i];
+     }
+   }
+ };
+
+ template <typename T>
+ struct vec2_dtype {
+   using type = T;
+ };
+
+ template <>
+ struct vec2_dtype<half> {
+   using type = half2;
+ };
+
+ template <>
+ struct vec2_dtype<__nv_bfloat16> {
+   using type = __nv_bfloat162;
+ };
+
+ template <>
+ struct vec2_dtype<__nv_fp8_e4m3> {
+   using type = __nv_fp8x2_e4m3;
+ };
+
+ template <>
+ struct vec2_dtype<__nv_fp8_e5m2> {
+   using type = __nv_fp8x2_e5m2;
+ };
+
+ template <typename T>
+ using vec2_dtype_t = typename vec2_dtype<T>::type;
+
+ template <typename T, size_t VEC_SIZE>
+ FLASHINFER_INLINE vec2_dtype_t<T> get_vec2_element(vec_t<T, VEC_SIZE>& vec, int i) {
+   static_assert(VEC_SIZE % 2 == 0, "VEC_SIZE must be a multiple of 2");
+   return ((vec2_dtype_t<T>*)&(vec[0]))[i];
+ }
+
+ }  // namespace flashinfer
+
+ #endif  // VEC_DTYPES_CUH_
\ No newline at end of file
diff --git a/vllm_musa/csrc/quantization/matmul_mudnn.cpp b/vllm_musa/csrc/quantization/matmul_mudnn.cpp
new file mode 100644
index 000000000..7ff6a22b2
--- /dev/null
+++ b/vllm_musa/csrc/quantization/matmul_mudnn.cpp
@@ -0,0 +1,218 @@
+#include <torch/all.h>
+#include <torch/extension.h>
+#include <torch_musa/csrc/core/MUSAGuard.h>
+#include <torch_musa/csrc/aten/utils/Utils.h>
+#include <torch_musa/csrc/aten/utils/Context.h>
+#include <torch/csrc/utils/pycfunction_helpers.h>
+
+#include <optional>
+#include <vector>
+#include <mudnn.h>
+
+using at::musa::muTensor;
+using at::musa::GetComputeModeFromCtx;
+
+void inline W4A16KernelCheck(torch::Tensor &out, torch::Tensor &in,
+                             torch::Tensor &weight, torch::Tensor &scales,
+                             std::optional<torch::Tensor> &bias, const bool trans_w) {
+  // input shape: [M, K], fp16
+  // weight shape: [K // 8, N], int32
+  // out shape: [M, N], fp16
+  // scale shape: [1, N], float
+  TORCH_CHECK(out.dim() == in.dim() && in.dim() == weight.dim() &&
+                  weight.dim() == 2,
+              "output, input and weight should all be 2D tensors",
+              "got output: ", out.dim(), ", input: ", in.dim(),
+              ", weight: ", weight.dim());
+  TORCH_CHECK(
+      out.scalar_type() == at::ScalarType::Half,
+      "dtype of output tensor of mudnn w4a16 kernel should be half, got ",
+      out.scalar_type());
+  TORCH_CHECK(
+      in.scalar_type() == at::ScalarType::Half,
+      "dtype of input tensor of mudnn w4a16 kernel should be half, got ",
+      in.scalar_type());
+  TORCH_CHECK(
+      weight.scalar_type() == at::ScalarType::Int,
+      "dtype of weight tensor of mudnn w4a16 kernel should be int32, got ",
+      weight.scalar_type());
+  TORCH_CHECK(out.size(0) == in.size(0) && out.size(1) == weight.size(1),
+              "unmatched shape of out, input and weight");
+  TORCH_CHECK(in.size(1) == weight.size(0) * 8,
+              "unmatched shape of input and weight");
+  TORCH_CHECK(scales.scalar_type() == at::ScalarType::Float,
+              "scales dtype should be float");
+  TORCH_CHECK(scales.dim() == 2,
+              "scales should all be 2D tensors");
+  TORCH_CHECK(scales.size(1) == weight.size(1),
+              "scales shape should be equal to weight channels");
+  if (bias.has_value()) {
+    TORCH_CHECK(bias->scalar_type() == out.scalar_type(), "bias should have same dtype as in&we");
+    TORCH_CHECK(bias->size(0) == out.size(1), "bias shape should match out-channel");
+  }
+}
+
+muTensor CreateW4MUTensor(torch::Tensor &weight, torch::Tensor &scales,
+                          const bool trans) {
+  muTensor w_m;
+  // weight: (K/8, N) -> w_m: (K, M) with INT4
+  CHECK_MUDNN_STATUS(w_m.SetType(muTensor::Type::QINT4), "Set QINT4 dtype");
+  CHECK_MUDNN_STATUS(w_m.SetAddr(weight.data_ptr()), "SetAddr");
+  CHECK_MUDNN_STATUS(w_m.SetFormat(muTensor::Format::NCHW), "SetFormat");
+  if (trans) {
+    // set transposed shape and stride
+    CHECK_MUDNN_STATUS(w_m.SetNdInfo({weight.size(1), weight.size(0) * 8},
+                                     {weight.stride(1) * 8, weight.stride(0)}),
+                       "SetNdInfo");
+  } else {
+    CHECK_MUDNN_STATUS(w_m.SetNdInfo({weight.size(0) * 8, weight.size(1)},
+                                     {weight.stride(0), weight.stride(1)}),
+                       "SetNdInfo");
+  }
+  return w_m;
+}
+
+void mudnn_w4a16_matmul(torch::Tensor out, torch::Tensor in,
+                        torch::Tensor weight, torch::Tensor scales,
+                        std::optional<torch::Tensor> bias,
+                        const bool split_k = false) {
+  bool trans_l = at::musa::IsTranspose(in);
+  bool trans_r = at::musa::IsTranspose(weight);
+  W4A16KernelCheck(out, in, weight, scales, bias, trans_r);
+
+  at::musa::muHandle &h = at::GetMudnnHandle();
+  muTensor in_mt = at::musa::CreateMUTensor(in);
+  muTensor o_mt = at::musa::CreateMUTensor(out);
+  muTensor s_mt = at::musa::CreateMUTensor(scales);
+  muTensor w_mt = CreateW4MUTensor(weight, scales, trans_r);
+  torch::Tensor bias_ = bias.value_or(torch::Tensor());
+  muTensor b_mt = at::musa::CreateMUTensor(bias_);
+
+  ::musa::dnn::BatchMatMul op;
+  CHECK_MUDNN_STATUS(op.SetAlpha(1.0), "SetAlpha");
+  CHECK_MUDNN_STATUS(op.SetBeta(0.0), "SetBeta");
+  CHECK_MUDNN_STATUS(op.SetGamma(0.0), "SetGamma");
+  CHECK_MUDNN_STATUS(op.SetTranspose(trans_l, trans_r), "SetTranspose");
+  CHECK_MUDNN_STATUS(op.SetSplitK(split_k), "SetSplitK");
+  CHECK_MUDNN_STATUS(
+      op.SetComputeMode(static_cast<::musa::dnn::MatMul::ComputeMode>(0)),
+      "SetComputeMode");
+
+  ::musa::dnn::MatMulLtParam param;
+  CHECK_MUDNN_STATUS(param.SetScale(muTensor(), s_mt, muTensor(), muTensor()),
+                     "SetScale");
+
+  CHECK_MUDNN_STATUS(op.RunLt(h, o_mt, in_mt, w_mt, o_mt, b_mt, param,
+                              at::musa::InternalMemAlloc),
+                     "RunLt");
+}
+
+void mudnn_w8a8_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
+  torch::Tensor const& b, torch::Tensor const& a_scales,
+  torch::Tensor const& b_scales,
+  std::optional<torch::Tensor> const& bias){
+  TORCH_CHECK(a.dim() == 2 && b.dim() == 2 && c.dim() == 2);
+  TORCH_CHECK(c.size(0) == a.size(0) && a.size(1) == b.size(1) &&
+  b.size(0) == c.size(1));
+
+  if (bias) {
+    TORCH_CHECK(bias->numel() == b.size(1) && bias->is_contiguous() &&
+      bias->dim() == 1);
+  }
+
+  at::musa::OptionalMUSAGuard const device_guard(device_of(a));
+
+  const auto a_type = a.scalar_type();
+
+  muTensor a_mu = at::musa::CreateMUTensor(a);
+  muTensor b_mu = at::musa::CreateMUTensor(b);
+  muTensor out_mu = at::musa::CreateMUTensor(c);
+  torch::Tensor bias_ = bias.value_or(torch::Tensor());
+  muTensor bias_mu = at::musa::CreateMUTensor(bias_);
+
+  muTensor a_scales_mu = at::musa::CreateMUTensor(a_scales);
+  muTensor b_scales_mu = at::musa::CreateMUTensor(b_scales);
+
+  auto& handle = at::GetMudnnHandle();
+  ::musa::dnn::BatchMatMul op;
+  ::musa::dnn::MatMulLtParam param;
+
+  CHECK_MUDNN_STATUS(param.SetScale(a_scales_mu, b_scales_mu, muTensor(), muTensor(), 128), "SetScale");
+  CHECK_MUDNN_STATUS(op.SetTranspose(false, true), "SetTranspose");
+  CHECK_MUDNN_STATUS(op.SetSplitK(false), "SetSplitK");
+  CHECK_MUDNN_STATUS(op.SetComputeMode(GetComputeModeFromCtx(a_type)), "SetComputeMode");
+  CHECK_MUDNN_STATUS(op.SetBeta(0.0), "SetBeta");
+
+  CHECK_MUDNN_STATUS(
+    op.RunLt(handle,
+      out_mu,
+      a_mu,
+      b_mu,
+      out_mu,
+      bias_mu,
+      param,
+      at::musa::InternalMemAlloc),
+  "RunLt");
+}
+
+
+void mla_mudnn_bmm(
+  torch::Tensor& c,
+  torch::Tensor const& a,
+  torch::Tensor const& b) {
+  TORCH_CHECK(a.dim() == 3 && b.dim() == 3 && c.dim() == 3);
+
+  at::musa::OptionalMUSAGuard const device_guard(device_of(a));
+
+  const auto a_type = a.scalar_type();
+
+  muTensor a_mu = at::musa::CreateMUTensor(a);
+  muTensor b_mu = at::musa::CreateMUTensor(b);
+  muTensor out_mu = at::musa::CreateMUTensor(c);
+
+  bool trans_a = false;
+  bool trans_b = false;
+
+  if (b.stride(1) == 1) {
+      trans_b = true;
+      // (B, N, P) * (N, P, L) -> (B, N, L)
+      a_mu.SetNdInfo({a.size(1), a.size(0), a.size(2)},
+      {a.stride(1), a.stride(0), a.stride(2)});
+      b_mu.SetNdInfo({b.size(0), b.size(2), b.size(1)},
+        {b.stride(0), b.stride(2), b.stride(1)});
+      out_mu.SetNdInfo({c.size(1), c.size(0), c.size(2)},
+        {c.stride(1), c.stride(0), c.stride(2)});
+        // (N, B, P) * (N, L, P) -> (N, B, L)
+  } else if (b.stride(2) == 1) {
+      // (B, N, P) * (N, P, L) -> (B, N, L)
+      a_mu.SetNdInfo({a.size(1), a.size(0), a.size(2)},
+      {a.stride(1), a.stride(0), a.stride(2)});
+      b_mu.SetNdInfo({b.size(0), b.size(1), b.size(2)},
+        {b.stride(0), b.stride(1), b.stride(2)});
+      out_mu.SetNdInfo({c.size(1), c.size(0), c.size(2)},
+        {c.stride(1), c.stride(0), c.stride(2)});
+      // (N, B, P) * (N, P, L) -> (N, B, L)
+  } else {
+    TORCH_CHECK(false, "mla_mudnn_bmm not support current stide, export MUDNN_LOG_LEVEL=INFO");
+  }
+
+  auto& handle = at::GetMudnnHandle();
+  ::musa::dnn::BatchMatMul op;
+  ::musa::dnn::MatMulLtParam param;
+
+  CHECK_MUDNN_STATUS(op.SetTranspose(trans_a, trans_b), "SetTranspose");
+  CHECK_MUDNN_STATUS(op.SetSplitK(false), "SetSplitK");
+  CHECK_MUDNN_STATUS(op.SetComputeMode(GetComputeModeFromCtx(a_type)), "SetComputeMode");
+  CHECK_MUDNN_STATUS(op.SetBeta(0.0), "SetBeta");
+
+  CHECK_MUDNN_STATUS(
+    op.RunLt(handle,
+      out_mu,
+      a_mu,
+      b_mu,
+      out_mu,
+      muTensor(),
+      param,
+      at::musa::InternalMemAlloc),
+  "RunLt");
+}
\ No newline at end of file
diff --git a/vllm_musa/csrc/quantization/quant_utils.muh b/vllm_musa/csrc/quantization/quant_utils.muh
new file mode 100644
index 000000000..306ed86a9
--- /dev/null
+++ b/vllm_musa/csrc/quantization/quant_utils.muh
@@ -0,0 +1,30 @@
+#pragma once
+
+#include "../attention/attention_dtypes.h"
+#include <assert.h>
+#include <float.h>
+#include <stdint.h>
+#include <type_traits>
+
+namespace vllm_musa {
+namespace fp8 {
+// The following macro is used to dispatch the conversion function based on
+// the data type of the key and value cache. The FN is a macro that calls a
+// function with template<typename scalar_t, typename cache_t,
+// Fp8KVCacheDataType kv_dt>.
+#define DISPATCH_BY_KV_CACHE_DTYPE(SRC_DTYPE, KV_DTYPE, FN)                      \
+    if (KV_DTYPE == "auto") {                                                    \
+        if (SRC_DTYPE == at::ScalarType::Float) {                                \
+          FN(float, float, vllm_musa::Fp8KVCacheDataType::kAuto);                     \
+        } else if (SRC_DTYPE == at::ScalarType::Half) {                          \
+          FN(uint16_t, uint16_t, vllm_musa::Fp8KVCacheDataType::kAuto);               \
+        } else if (SRC_DTYPE == at::ScalarType::BFloat16) {                      \
+          FN(__nv_bfloat16, __nv_bfloat16, vllm_musa::Fp8KVCacheDataType::kAuto);     \
+        } else {                                                                 \
+          TORCH_CHECK(false, "Unsupported input type of kv cache: ", SRC_DTYPE); \
+        }                                                                        \
+    } else {                                                                     \
+        TORCH_CHECK(false, "Unsupported data type of kv cache: ", KV_DTYPE);     \
+    }
+} // namespace fp8
+} // namespace vllm_musa
\ No newline at end of file
diff --git a/vllm_musa/csrc/quantization/vectorization.muh b/vllm_musa/csrc/quantization/vectorization.muh
new file mode 100644
index 000000000..eef71e4b0
--- /dev/null
+++ b/vllm_musa/csrc/quantization/vectorization.muh
@@ -0,0 +1,33 @@
+#pragma once
+/**
+ * __device__ datatypes vectorized by 4
+ */
+
+// Include both AMD and NVIDIA fp8 types to avoid circular import
+// TODO(luka/varun) use FP8_TYPE instead after refactoring
+#include <c10/util/Float8_e4m3fnuz.h>
+#include <c10/util/Float8_e4m3fn.h>
+
+namespace vllm_musa {
+
+// Vectorization containers
+template <typename scalar_t>
+struct __align__(8) vec4_t {
+  scalar_t x;
+  scalar_t y;
+  scalar_t z;
+  scalar_t w;
+};
+
+template <typename quant_type_t>
+struct __align__(4) q8x4_t {
+  static_assert(std::is_same_v<quant_type_t, int8_t> ||
+                std::is_same_v<quant_type_t, c10::Float8_e4m3fn> ||
+                std::is_same_v<quant_type_t, c10::Float8_e4m3fnuz>);
+  quant_type_t x;
+  quant_type_t y;
+  quant_type_t z;
+  quant_type_t w;
+};
+
+}  // namespace vllm_musa
diff --git a/vllm_musa/csrc/rmsnorm_cast/dtype.muh b/vllm_musa/csrc/rmsnorm_cast/dtype.muh
new file mode 100644
index 000000000..4a4ba1553
--- /dev/null
+++ b/vllm_musa/csrc/rmsnorm_cast/dtype.muh
@@ -0,0 +1,409 @@
+#pragma once
+#include <musa_runtime.h>
+#include <stdint.h>
+
+#ifdef __MTGPU__
+#include <musa_robust.h>
+#endif
+
+#include "musa_bf16.h"
+#include "musa_fp16.h"
+#include "musa_fp8.h"
+
+namespace vllm_musa {
+
+typedef __half float16_t;
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+typedef __mt_bfloat16 bfloat16_t;
+#else
+typedef int16_t bfloat16_t;
+#endif
+
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+#define __MUSA_HOSTDEVICE__ __host__ __device__
+template <typename T> struct FDtypeWrapper {
+  T _x;
+  __MUSA_HOSTDEVICE__ FDtypeWrapper() = default;
+  __MUSA_HOSTDEVICE__ FDtypeWrapper(const float x) {
+    if constexpr (std::is_same<__mt_fp8_e4m3, T>::value) {
+      // hw instruction only support satfinity mode, add if to support normal
+      // mode
+      uint32_t f_abs_bits = (*(uint32_t *)&x) & (uint32_t)0x7fffffff;
+      // https://github.com/pytorch/pytorch/blob/16de25b1dc3af7cc3efdb24e5637a62b8c329881/c10/util/Float8_e4m3fn.h#L183-L186
+      constexpr uint32_t fp8_max = 0x43e80001;
+      if (f_abs_bits >= fp8_max) {
+        _x.__x = 0x7f;
+        _x.__x |= UINT8_C(((*(uint32_t *)&x) & UINT32_C(0x80000000)) >> 24); //
+      } else {
+        _x = __mt_fp8_e4m3(x);
+      }
+    } else if constexpr (std::is_same<__mt_fp8_e5m2, T>::value) {
+      // hw instruction only support satfinity mode, add if to support normal
+      // mode
+      uint32_t f_abs_bits = (*(uint32_t *)&x) & (uint32_t)0x7fffffff;
+      // https://github.com/pytorch/pytorch/blob/16de25b1dc3af7cc3efdb24e5637a62b8c329881/c10/util/Float8_e5m2.h#L64
+      constexpr uint32_t fp8_max = 0x47700000;
+      constexpr uint32_t fp32_inf = UINT32_C(255) << 23;
+      if (f_abs_bits >= fp8_max) {
+        _x.__x = f_abs_bits > fp32_inf ? UINT8_C(0x7F) : UINT8_C(0x7C);
+        _x.__x |= UINT8_C(((*(uint32_t *)&x) & UINT32_C(0x80000000)) >> 24);
+      } else {
+        _x = __mt_fp8_e5m2(x);
+      }
+    } else {
+      _x = T(x);
+    }
+  }
+  __MUSA_HOSTDEVICE__ operator float() const {
+    if constexpr (std::is_same<__mt_fp8_e4m3, T>::value) {
+      __half temp;
+      asm("CVT.RTE.DENORM.BST1 %0.f16, %1.e4m3\n" : "=R"(temp) : "R"(_x));
+      return (float)temp;
+    } else if constexpr (std::is_same<__mt_fp8_e5m2, T>::value) {
+      __half temp;
+      asm("CVT.RTE.DENORM.BST1 %0.f16, %1.e5m2\n" : "=R"(temp) : "R"(_x));
+      return (float)temp;
+    } else {
+      return (float)_x;
+    }
+  }
+  __MUSA_HOSTDEVICE__ FDtypeWrapper &operator+=(const float x) {
+    _x = T((float)_x + x);
+    return *this;
+  }
+  __MUSA_HOSTDEVICE__ FDtypeWrapper &operator-=(const float x) {
+    _x = T((float)_x - x);
+    return *this;
+  }
+  __MUSA_HOSTDEVICE__ FDtypeWrapper &operator*=(const float x) {
+    _x = T((float)_x * x);
+    return *this;
+  }
+  __MUSA_HOSTDEVICE__ FDtypeWrapper &operator/=(const float x) {
+    _x = T((float)_x / x);
+    return *this;
+  }
+};
+#undef __MUSA_HOSTDEVICE__
+typedef FDtypeWrapper<__mt_fp8_e4m3> fp8_e4m3_t;
+typedef FDtypeWrapper<__mt_fp8_e5m2> fp8_e5m2_t;
+#else
+typedef int8_t fp8_e4m3_t;
+typedef uint8_t fp8_e5m2_t;
+#endif
+
+#define MACRO_UNROLL _Pragma("unroll")
+
+#define ATTR_ALIGNED(v) __attribute__((aligned(v)))
+#define SELF_VEC_DEF(BASE_TYPE, VEC_TYPE_V2, VEC_TYPE_V4)                      \
+  struct ATTR_ALIGNED(sizeof(BASE_TYPE) * 2) VEC_TYPE_V2 {                     \
+    __device__ VEC_TYPE_V2() {}                                                \
+    __device__ VEC_TYPE_V2(const VEC_TYPE_V2 &t) {                             \
+      this->x = t.x;                                                           \
+      this->y = t.y;                                                           \
+    }                                                                          \
+    BASE_TYPE x, y;                                                            \
+  };                                                                           \
+                                                                               \
+  __device__ __forceinline__ VEC_TYPE_V2 make_##VEC_TYPE_V2(BASE_TYPE x,       \
+                                                            BASE_TYPE y) {     \
+    VEC_TYPE_V2 t;                                                             \
+    t.x = x, t.y = y;                                                          \
+    return t;                                                                  \
+  }                                                                            \
+                                                                               \
+  struct ATTR_ALIGNED(sizeof(BASE_TYPE) * 4) VEC_TYPE_V4 {                     \
+    __device__ VEC_TYPE_V4() {}                                                \
+    __device__ VEC_TYPE_V4(const VEC_TYPE_V4 &t) {                             \
+      this->x = t.x;                                                           \
+      this->y = t.y;                                                           \
+      this->z = t.z;                                                           \
+      this->w = t.w;                                                           \
+    }                                                                          \
+    BASE_TYPE x, y, z, w;                                                      \
+  };                                                                           \
+                                                                               \
+  __device__ __forceinline__ VEC_TYPE_V4 make_##VEC_TYPE_V4(                   \
+      BASE_TYPE x, BASE_TYPE y, BASE_TYPE z, BASE_TYPE w) {                    \
+    VEC_TYPE_V4 t;                                                             \
+    t.x = x, t.y = y, t.z = z, t.w = w;                                        \
+    return t;                                                                  \
+  }
+
+SELF_VEC_DEF(float16_t, Half2, Half4)
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+SELF_VEC_DEF(bfloat16_t, Bhalf2, Bhalf4)
+#endif
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+SELF_VEC_DEF(fp8_e5m2_t, E5m2x2, E5m2x4)
+SELF_VEC_DEF(fp8_e4m3_t, E4m3x2, E4m3x4)
+#endif
+
+#define GEN_VECTYPE(_CTYPE, _VECTYPE, _BYTES, _VLEN)                           \
+  struct ATTR_ALIGNED(_BYTES) _VECTYPE {                                       \
+    __device__ _VECTYPE() {}                                                   \
+    __device__ _VECTYPE(const _VECTYPE &t) {                                   \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < _VLEN; i++) {                                        \
+        this->arr[i] = t.arr[i];                                               \
+      }                                                                        \
+    }                                                                          \
+    _CTYPE arr[_VLEN];                                                         \
+  }
+GEN_VECTYPE(float16_t, Half8, 16, 8);
+GEN_VECTYPE(float, Float8, 32, 8);
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+GEN_VECTYPE(bfloat16_t, Bhalf8, 16, 8);
+GEN_VECTYPE(bfloat16_t, Bhalf16, 32, 16);
+GEN_VECTYPE(bfloat16_t, Bhalf32, 64, 32);
+#endif
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+GEN_VECTYPE(fp8_e4m3_t, E4m3x8, 8, 8);
+GEN_VECTYPE(fp8_e4m3_t, E4m3x16, 16, 16);
+GEN_VECTYPE(fp8_e4m3_t, E4m3x32, 32, 32);
+GEN_VECTYPE(fp8_e5m2_t, E5m2x8, 8, 8);
+GEN_VECTYPE(fp8_e5m2_t, E5m2x16, 16, 16);
+GEN_VECTYPE(fp8_e5m2_t, E5m2x32, 32, 32);
+#endif
+GEN_VECTYPE(float16_t, Half16, 32, 16);
+GEN_VECTYPE(float, Float16, 64, 16);
+
+template <typename type> class Dtype;
+
+#define INST(_type, _vec2, _vec4)                                              \
+  template <> class Dtype<_type> {                                             \
+  public:                                                                      \
+    using Scalar = _type;                                                      \
+    using Vec2 = _vec2;                                                        \
+    using Vec4 = _vec4;                                                        \
+    static __device__ __forceinline__ Vec2 make_vec2(_type x, _type y) {       \
+      return make_##_vec2(x, y);                                               \
+    }                                                                          \
+    static __device__ __forceinline__ Vec4 make_vec4(_type x, _type y,         \
+                                                     _type z, _type w) {       \
+      return make_##_vec4(x, y, z, w);                                         \
+    }                                                                          \
+  }
+
+INST(float, float2, float4);
+INST(float16_t, Half2, Half4);
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+INST(bfloat16_t, Bhalf2, Bhalf4);
+#endif
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+INST(fp8_e5m2_t, E5m2x2, E5m2x4);
+INST(fp8_e4m3_t, E4m3x2, E4m3x4);
+#endif
+
+#undef INST
+
+template <typename T> struct DeduceVectorizedType { using Type = T; };
+
+template <> struct DeduceVectorizedType<bool> { using Type = int8_t; };
+
+template <> struct DeduceVectorizedType<half> { using Type = _Float16; };
+
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+template <> struct DeduceVectorizedType<bfloat16_t> { using Type = _Float16; };
+#endif
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+template <> struct DeduceVectorizedType<fp8_e4m3_t> { using Type = int8_t; };
+template <> struct DeduceVectorizedType<fp8_e5m2_t> { using Type = int8_t; };
+#endif
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ == 310 // MUSIFY_EXCL_LINE
+#define LD_BYP_SLC(_BITS, _BYTES)                                              \
+  VecType dst;                                                                 \
+  const BaseType *addr = ptr + idx;                                            \
+  asm volatile("LSU.LD.B" #_BITS " %0, %1, _, " #_BYTES                        \
+               ", 1, 1, inner_persist=0, "                                     \
+               "outer_persist=2, chrnt=l2_l3, slc=byp, persist=0, "            \
+               "stride_add_first=0"                                            \
+               : "=R"(dst)                                                     \
+               : "R"(addr));                                                   \
+  return dst;
+#else
+#define LD_BYP_SLC(_BITS, _BYTES) return *(VecType *)(ptr + idx);
+#endif
+
+template <typename T, int bits = 16 * 8> struct VecType;
+
+#define DEF_VECT(_CTYPE, _VECTYPE)                                             \
+  template <> struct VecType<_CTYPE, sizeof(_VECTYPE) * 8> {                   \
+    static constexpr int vec_bytes = sizeof(_VECTYPE);                         \
+    static constexpr int bit_per_byte = 8;                                     \
+    using BaseType = _CTYPE;                                                   \
+    using RobustTypePtr = __musa::robust_ptr<_CTYPE>;                          \
+    using Ttype = _VECTYPE;                                                    \
+    static constexpr int bits = vec_bytes * bit_per_byte;                      \
+    static constexpr int vlen = bits / (sizeof(BaseType) * bit_per_byte);      \
+    using VectorizedType = typename DeduceVectorizedType<BaseType>::Type;      \
+    typedef VectorizedType VxTy __attribute__((vector_size(vec_bytes)));       \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ VecType load(const BaseType *ptr,        \
+                                                   OffsetType idx) {           \
+      return *(VecType *)(ptr + idx);                                          \
+    }                                                                          \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ VecType                                  \
+    load_byp_slc(const BaseType *ptr, OffsetType idx) {                        \
+      if constexpr (vec_bytes == 16) {                                         \
+        LD_BYP_SLC(128, 16);                                                   \
+      } else if constexpr (vec_bytes == 8) {                                   \
+        LD_BYP_SLC(64, 8);                                                     \
+      } else if constexpr (vec_bytes == 4) {                                   \
+        LD_BYP_SLC(32, 4);                                                     \
+      } else if constexpr (vec_bytes == 2) {                                   \
+        LD_BYP_SLC(16, 2);                                                     \
+      } else {                                                                 \
+        LD_BYP_SLC(8, 1);                                                      \
+      }                                                                        \
+    }                                                                          \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ VecType                                  \
+    robust_load(const RobustTypePtr ptr, OffsetType idx) {                     \
+      return __musa::robust_load<VecType, BaseType>(ptr, idx);                 \
+    }                                                                          \
+                                                                               \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ void                                     \
+    store(BaseType *ptr, OffsetType idx, const VecType &dst) {                 \
+      *(VecType *)(ptr + idx) = dst;                                           \
+    }                                                                          \
+    template <typename OffsetType>                                             \
+    static __device__ __forceinline__ void                                     \
+    robust_store(RobustTypePtr ptr, OffsetType idx, const VecType &dst) {      \
+      __musa::robust_store<VecType, BaseType>(dst, ptr, idx);                  \
+    }                                                                          \
+                                                                               \
+    __device__ VecType() {                                                     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = 0;                                                \
+      }                                                                        \
+    }                                                                          \
+    __device__ VecType(const VecType &t) {                                     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = t.val_.elem[i];                                   \
+      }                                                                        \
+    }                                                                          \
+    __device__ VecType &operator=(const VecType &t) {                          \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = t.val_.elem[i];                                   \
+      }                                                                        \
+      return *this;                                                            \
+    }                                                                          \
+    __device__ VecType(_CTYPE val) {                                           \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = val;                                              \
+      }                                                                        \
+    }                                                                          \
+    template <typename SrcVecType>                                             \
+    friend __device__ VecType operator+(VecType lhs, const SrcVecType &rhs) {  \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        lhs.val_.elem[i] += static_cast<BaseType>(rhs.val_.elem[i]);           \
+      }                                                                        \
+      return lhs;                                                              \
+    }                                                                          \
+    friend __device__ VecType operator+(VecType lhs, const _CTYPE &rhs) {      \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        lhs.val_.elem[i] += rhs;                                               \
+      }                                                                        \
+      return lhs;                                                              \
+    }                                                                          \
+    friend __device__ VecType operator-(VecType lhs, const VecType &rhs) {     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        lhs.val_.elem[i] -= rhs.val_.elem[i];                                  \
+      }                                                                        \
+      return lhs;                                                              \
+    }                                                                          \
+    friend __device__ VecType operator*(VecType lhs, const VecType &rhs) {     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        lhs.val_.elem[i] *= rhs.val_.elem[i];                                  \
+      }                                                                        \
+      return lhs;                                                              \
+    }                                                                          \
+    template <typename Func> __device__ VecType &apply() {                     \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        this->val_.elem[i] = Func::apply(this->val_.elem[i]);                  \
+      }                                                                        \
+      return *this;                                                            \
+    }                                                                          \
+    template <typename SrcVecType>                                             \
+    static __device__ VecType cvt(const SrcVecType &src) {                     \
+      VecType dst;                                                             \
+      MACRO_UNROLL                                                             \
+      for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {             \
+        dst.val_.elem[i] = (BaseType)(src.val_.elem[i]);                       \
+      }                                                                        \
+      return dst;                                                              \
+    }                                                                          \
+    union U {                                                                  \
+      __device__ U() {                                                         \
+        MACRO_UNROLL                                                           \
+        for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) {           \
+          this->elem[i] = 0;                                                   \
+        }                                                                      \
+      }                                                                        \
+      Ttype storage;                                                           \
+      BaseType elem[sizeof(Ttype) / sizeof(BaseType)];                         \
+      VxTy vt_elem;                                                            \
+    };                                                                         \
+    U val_{};                                                                  \
+  }
+DEF_VECT(float16_t, float16_t);
+DEF_VECT(float16_t, Half2);
+DEF_VECT(float16_t, Half4);
+DEF_VECT(float16_t, Half8);
+DEF_VECT(float16_t, Half16);
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
+DEF_VECT(bfloat16_t, bfloat16_t);
+DEF_VECT(bfloat16_t, Bhalf2);
+DEF_VECT(bfloat16_t, Bhalf4);
+DEF_VECT(bfloat16_t, Bhalf8);
+DEF_VECT(bfloat16_t, Bhalf16);
+DEF_VECT(bfloat16_t, Bhalf32);
+#endif
+#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 310 || !defined(__MUSA_ARCH__))
+DEF_VECT(fp8_e4m3_t, fp8_e4m3_t);
+DEF_VECT(fp8_e4m3_t, E4m3x2);
+DEF_VECT(fp8_e4m3_t, E4m3x4);
+DEF_VECT(fp8_e4m3_t, E4m3x8);
+DEF_VECT(fp8_e4m3_t, E4m3x16);
+DEF_VECT(fp8_e4m3_t, E4m3x32);
+DEF_VECT(fp8_e5m2_t, fp8_e5m2_t);
+DEF_VECT(fp8_e5m2_t, E5m2x2);
+DEF_VECT(fp8_e5m2_t, E5m2x4);
+DEF_VECT(fp8_e5m2_t, E5m2x8);
+DEF_VECT(fp8_e5m2_t, E5m2x16);
+DEF_VECT(fp8_e5m2_t, E5m2x32);
+#endif
+DEF_VECT(float, float);
+DEF_VECT(float, float2);
+DEF_VECT(float, float3);
+DEF_VECT(float, float4);
+DEF_VECT(float, Float8);
+DEF_VECT(float, Float16);
+
+#undef DEF_VECT
+#undef MACRO_UNROLL
+template <typename T> struct ComputeDType {
+  using Type = typename std::conditional<
+      (sizeof(T) >= 4), T,
+      typename std::conditional<
+          std::is_integral<T>::value,
+          typename std::conditional<std::is_unsigned<T>::value, uint32_t,
+                                    int32_t>::type,
+          float>::type>::type;
+};
+template <> struct ComputeDType<bool> { using Type = bool; };
+
+} // namespace vllm_musa
diff --git a/vllm_musa/csrc/rmsnorm_cast/math.muh b/vllm_musa/csrc/rmsnorm_cast/math.muh
new file mode 100644
index 000000000..5e14180a1
--- /dev/null
+++ b/vllm_musa/csrc/rmsnorm_cast/math.muh
@@ -0,0 +1,26 @@
+
+#pragma once
+#include "dtype.muh"
+#include <math.h>
+#include <type_traits>
+
+namespace vllm_musa {
+
+static __device__ __forceinline__ float fast_rcpf(float x) {
+  float y = __frcp_rn(x);
+  y = y * (2.f - x * y);
+  return y;
+}
+
+static __device__ __forceinline__ float fast_rsqrtf(float a) {
+  float x = 0.5f * a;
+  float y = __frsqrt_rn(a);
+  y = y * (1.5f - x * y * y);
+  return y;
+}
+
+static __device__ __forceinline__ float fast_divf(float a, float b) {
+  return a * fast_rcpf(b);
+}
+
+} // namespace vllm_musa
diff --git a/vllm_musa/csrc/rmsnorm_cast/rmsnorm_cast.mu b/vllm_musa/csrc/rmsnorm_cast/rmsnorm_cast.mu
new file mode 100644
index 000000000..3798d287b
--- /dev/null
+++ b/vllm_musa/csrc/rmsnorm_cast/rmsnorm_cast.mu
@@ -0,0 +1,299 @@
+#include "../dispatch_utils.h"
+#include "../type_convert.muh"
+#include "dtype.muh"
+#include "math.muh"
+#include "update_helper.muh"
+
+#include <torch/all.h>
+
+namespace vllm_musa {
+
+template <typename ComputeType, int width>
+__device__ ComputeType warp_max_reduce(ComputeType val) {
+#pragma unroll
+  for (int offset = width / 2; offset > 0; offset /= 2) {
+    val = fmaxf(val, mudnn_shfl_down_sync<ComputeType, width>(val, offset));
+  }
+  int lane_id = threadIdx.x & (warpSize - 1);
+  int subwarp_id = lane_id / width;
+  int subwarp_start = subwarp_id * width;
+  val = mudnn_shfl_sync<ComputeType, width>(val, subwarp_start);
+  return val;
+}
+
+template <typename SrcDtype, typename ComputeType, typename DstDtype,
+          typename StatDtype, int BLOCK_X, int BLOCK_Y, int vlen,
+          int quant_size>
+__global__ void rmsnorm_cast_kernel(
+                                    fp8_e4m3_t *output_fp8, float *scale,
+                                    const SrcDtype *input,
+                                    const DstDtype *gamma, const size_t M,
+                                    const size_t N, const ComputeType eps) {
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 310)
+  size_t tx = threadIdx.x;
+  size_t ty = threadIdx.y;
+  size_t m_idx = blockIdx.x * blockDim.y + ty;
+  size_t n_idx = tx * vlen;
+  size_t n_step = (size_t)blockDim.x * vlen;
+
+  using SrcVec = VecType<SrcDtype, vlen * sizeof(SrcDtype) * 8>;
+  using DstVec = VecType<DstDtype, vlen * sizeof(DstDtype) * 8>;
+  using Fp8Vec = VecType<fp8_e4m3_t, vlen * sizeof(fp8_e4m3_t) * 8>;
+
+  ComputeType mu = 0;
+  ComputeType m2 = 0;
+  ComputeType cnt = 0;
+
+  const SrcDtype *__restrict p_src = input + m_idx * N;
+
+  // TODO(wuke): use robust_load, robust_store
+  bool m_valid = m_idx < M;
+  if (m_valid) {
+    for (size_t j = n_idx; j < N; j += n_step) {
+      SrcVec curr = *(const SrcVec *)(p_src + j);
+#pragma unroll
+      for (int k = 0; k < vlen; k++) {
+        VarUpdate<ComputeType, VarUpdateMode::WELFORD_ONLY_MEAN> welford_op;
+        welford_op.apply((ComputeType)curr.val_.elem[k] *
+                             (ComputeType)curr.val_.elem[k],
+                         &mu, &cnt);
+      }
+    }
+  }
+
+  AllReduceOp<ComputeType, BLOCK_X, BLOCK_Y, ReduceMode::WELFORD_ONLY_MEAN,
+              ReduceAxis::COLUMN>
+      all_reduce_op;
+  all_reduce_op.apply(&mu, &cnt, tx, ty);
+
+  if (m_valid) {
+    ComputeType mean = mu;
+    ComputeType inv_var;
+    inv_var = fast_rsqrtf(mean + eps);
+    (void)m2;
+
+    // DstDtype *__restrict p_dst = output + m_idx * N;
+
+    // max allreduce
+    constexpr int group_sz = quant_size / vlen;
+
+    for (size_t j = n_idx; j < N; j += n_step) {
+      SrcVec val = *(const SrcVec *)(p_src + j);
+      DstVec gamma_val = *(const DstVec *)(gamma + j);
+      DstVec dst;
+      Fp8Vec dst_fp8;
+      ComputeType local_max = eps;
+#pragma unroll
+      for (int k = 0; k < vlen; k++) {
+        dst.val_.elem[k] = (DstDtype)((ComputeType)val.val_.elem[k] * inv_var *
+                                      (ComputeType)gamma_val.val_.elem[k]);
+        local_max = fmaxf(local_max, fabsf((ComputeType)dst.val_.elem[k]));
+      }
+      // per group_sz compute max
+      // *(DstVec*)(p_dst + j) = dst;
+      ComputeType amax = warp_max_reduce<ComputeType, group_sz>(local_max);
+      // printf("thread id:[%d, %d], amax = %f.\n", tx, ty, amax);
+      // compute amax & scale
+      ComputeType scale_val = amax / 448.;
+      scale[m_idx * (N / quant_size) + j / quant_size] = (float)scale_val;
+// compute fp8 output
+#pragma unroll
+      for (int k = 0; k < vlen; k++) {
+        dst_fp8.val_.elem[k] =
+            (fp8_e4m3_t)((ComputeType)(dst.val_.elem[k]) / scale_val);
+      }
+      *(Fp8Vec *)(output_fp8 + m_idx * N + j) = dst_fp8;
+    }
+  }
+#endif
+}
+
+template <typename SrcDtype, typename ComputeType, typename DstDtype,
+          typename StatDtype, int BLOCK_X, int BLOCK_Y, int vlen,
+          int quant_size>
+__global__ void rmsnorm_add_cast_kernel(
+                                        fp8_e4m3_t *output_fp8, float *scale,
+                                        const SrcDtype *input,
+                                        const SrcDtype *resdual,
+                                        const DstDtype *gamma, const size_t M,
+                                        const size_t N, const ComputeType eps) {
+#if defined(__MUSA_ARCH__) && (__MUSA_ARCH__ == 310)
+  size_t tx = threadIdx.x;
+  size_t ty = threadIdx.y;
+  size_t m_idx = blockIdx.x * blockDim.y + ty;
+  size_t n_idx = tx * vlen;
+  size_t n_step = (size_t)blockDim.x * vlen;
+
+  using SrcVec = VecType<SrcDtype, vlen * sizeof(SrcDtype) * 8>;
+  using DstVec = VecType<DstDtype, vlen * sizeof(DstDtype) * 8>;
+  using Fp8Vec = VecType<fp8_e4m3_t, vlen * sizeof(fp8_e4m3_t) * 8>;
+
+  ComputeType mu = 0;
+  ComputeType m2 = 0;
+  ComputeType cnt = 0;
+
+  const SrcDtype *__restrict p_src = input + m_idx * N;
+  const SrcDtype *__restrict r_src = resdual + m_idx * N;
+  // TODO(wuke): use robust_load, robust_store
+  bool m_valid = m_idx < M;
+  if (m_valid) {
+    for (size_t j = n_idx; j < N; j += n_step) {
+      SrcVec curr = *(const SrcVec *)(p_src + j);
+      SrcVec resd = *(const SrcVec *)(r_src + j);
+#pragma unroll
+      for (int k = 0; k < vlen; k++) {
+        curr.val_.elem[k] = curr.val_.elem[k] + resd.val_.elem[k];
+        VarUpdate<ComputeType, VarUpdateMode::WELFORD_ONLY_MEAN> welford_op;
+        welford_op.apply((ComputeType)curr.val_.elem[k] *
+                             (ComputeType)curr.val_.elem[k],
+                         &mu, &cnt);
+      }
+      *(SrcVec*)(p_src + j) = curr;
+    }
+  }
+
+  AllReduceOp<ComputeType, BLOCK_X, BLOCK_Y, ReduceMode::WELFORD_ONLY_MEAN,
+              ReduceAxis::COLUMN>
+      all_reduce_op;
+  all_reduce_op.apply(&mu, &cnt, tx, ty);
+
+  if (m_valid) {
+    ComputeType mean = mu;
+    ComputeType inv_var;
+    inv_var = fast_rsqrtf(mean + eps);
+    (void)m2;
+
+    // DstDtype *__restrict p_dst = output + m_idx * N;
+    // max allreduce
+    constexpr int group_sz = quant_size / vlen;
+
+    for (size_t j = n_idx; j < N; j += n_step) {
+      SrcVec val = *(const SrcVec *)(p_src + j);
+      DstVec gamma_val = *(const DstVec *)(gamma + j);
+      DstVec dst;
+      Fp8Vec dst_fp8;
+      ComputeType local_max = eps;
+#pragma unroll
+      for (int k = 0; k < vlen; k++) {
+        dst.val_.elem[k] = (DstDtype)((ComputeType)val.val_.elem[k] * inv_var *
+                                      (ComputeType)gamma_val.val_.elem[k]);
+        local_max = fmaxf(local_max, fabsf((ComputeType)dst.val_.elem[k]));
+      }
+      // per group_sz compute max
+      // *(DstVec*)(p_dst + j) = dst;
+      ComputeType amax = warp_max_reduce<ComputeType, group_sz>(local_max);
+      // printf("thread id:[%d, %d], amax = %f.\n", tx, ty, amax);
+      // compute amax & scale
+      ComputeType scale_val = amax / 448.;
+      scale[m_idx * (N / quant_size) + j / quant_size] = (float)scale_val;
+// compute fp8 output
+#pragma unroll
+      for (int k = 0; k < vlen; k++) {
+        dst_fp8.val_.elem[k] =
+            (fp8_e4m3_t)((ComputeType)(dst.val_.elem[k]) / scale_val);
+      }
+      *(Fp8Vec *)(output_fp8 + m_idx * N + j) = dst_fp8;
+    }
+  }
+#endif
+}
+} // namespace vllm_musa
+
+void rmsnorm_cast(
+                  torch::Tensor &output_fp8, torch::Tensor &inv_scale,
+                  const torch::Tensor &input, const torch::Tensor &gamma,
+                  const float eps, const int quant_size) {
+  size_t M = input.size(0);
+  size_t N = input.size(1);
+
+  TORCH_CHECK(quant_size == 128, "quant block size only support 128.");
+  constexpr int vlen = 8;
+  int64_t tune_n = N % vlen == 0 ? N / vlen : N;
+#define LAUNCH_RMSNORM(BX, BY)                                                 \
+  do {                                                                         \
+    constexpr int BLOCK_X = (BX), BLOCK_Y = (BY);                              \
+    dim3 block(BLOCK_X, BLOCK_Y);                                              \
+    dim3 grid((M + BLOCK_Y - 1) / BLOCK_Y);                                    \
+    VLLM_DISPATCH_FLOATING_TYPES(                                              \
+        input.scalar_type(), "rmsnorm_cast_kernel", [&] {                      \
+          using scalar_torch = scalar_t;                                       \
+          using scalar_internal = std::conditional_t<                          \
+              std::is_same<scalar_torch, c10::Half>::value,                    \
+              vllm_musa::float16_t,                                            \
+              std::conditional_t<                                              \
+                  std::is_same<scalar_torch, c10::BFloat16>::value,            \
+                  vllm_musa::bfloat16_t, scalar_torch>>;                       \
+          vllm_musa::rmsnorm_cast_kernel<scalar_internal, float,               \
+                                         scalar_internal, float, BLOCK_X,      \
+                                         BLOCK_Y, vlen, 128>                   \
+              <<<grid, block, 0, stream>>>(                                    \
+                  static_cast<vllm_musa::fp8_e4m3_t *>(output_fp8.data_ptr()), \
+                  static_cast<float *>(inv_scale.data_ptr()),                  \
+                  static_cast<scalar_internal *>(input.data_ptr()),            \
+                  static_cast<scalar_internal *>(gamma.data_ptr()), M, N,      \
+                  static_cast<float>(eps));                                    \
+        });                                                                    \
+  } while (0)
+
+  const musaStream_t stream = at::musa::getCurrentMUSAStream();
+  if (M / tune_n > 16) {
+    LAUNCH_RMSNORM(64, 8);
+  } else if (tune_n > 256 && tune_n <= 1024) {
+    LAUNCH_RMSNORM(128, 4);
+  } else if (tune_n <= 256) {
+    LAUNCH_RMSNORM(64, 1);
+  } else {
+    LAUNCH_RMSNORM(256, 1);
+  }
+}
+
+void rmsnorm_add_cast(
+                      torch::Tensor &output_fp8, torch::Tensor &inv_scale,
+                      const torch::Tensor &input, const torch::Tensor &resdual,
+                      const torch::Tensor &gamma, const float eps,
+                      const int quant_size) {
+  size_t M = input.size(0);
+  size_t N = input.size(1);
+
+  TORCH_CHECK(quant_size == 128, "quant block size only support 128.");
+  constexpr int vlen = 8;
+  int64_t tune_n = N % vlen == 0 ? N / vlen : N;
+
+#define LAUNCH_RMSNORM_ADD(BX, BY)                                             \
+  do {                                                                         \
+    constexpr int BLOCK_X = (BX), BLOCK_Y = (BY);                              \
+    dim3 block(BLOCK_X, BLOCK_Y);                                              \
+    dim3 grid((M + BLOCK_Y - 1) / BLOCK_Y);                                    \
+    VLLM_DISPATCH_FLOATING_TYPES(                                              \
+        input.scalar_type(), "rmsnorm_add_cast_kernel", [&] {                  \
+          using scalar_torch = scalar_t;                                       \
+          using scalar_internal = std::conditional_t<                          \
+              std::is_same<scalar_torch, c10::Half>::value,                    \
+              vllm_musa::float16_t,                                            \
+              std::conditional_t<                                              \
+                  std::is_same<scalar_torch, c10::BFloat16>::value,            \
+                  vllm_musa::bfloat16_t, scalar_torch>>;                       \
+          vllm_musa::rmsnorm_add_cast_kernel<scalar_internal, float,           \
+                                             scalar_internal, float, BLOCK_X,  \
+                                             BLOCK_Y, vlen, 128>               \
+              <<<grid, block, 0, stream>>>(                                    \
+                  static_cast<vllm_musa::fp8_e4m3_t *>(output_fp8.data_ptr()), \
+                  static_cast<float *>(inv_scale.data_ptr()),                  \
+                  static_cast<scalar_internal *>(input.data_ptr()),            \
+                  static_cast<scalar_internal *>(resdual.data_ptr()),          \
+                  static_cast<scalar_internal *>(gamma.data_ptr()), M, N,      \
+                  static_cast<float>(eps));                                    \
+        });                                                                    \
+  } while (0)
+
+  const musaStream_t stream = at::musa::getCurrentMUSAStream();
+  if (M / tune_n > 16) {
+    LAUNCH_RMSNORM_ADD(64, 8);
+  } else if (tune_n > 256 && tune_n <= 1024) {
+    LAUNCH_RMSNORM_ADD(128, 4);
+  } else if (tune_n <= 256) {
+    LAUNCH_RMSNORM_ADD(64, 1);
+  } else {
+    LAUNCH_RMSNORM_ADD(256, 1);
+  }
+}
diff --git a/vllm_musa/csrc/rmsnorm_cast/shuffle.muh b/vllm_musa/csrc/rmsnorm_cast/shuffle.muh
new file mode 100644
index 000000000..3b8a739bc
--- /dev/null
+++ b/vllm_musa/csrc/rmsnorm_cast/shuffle.muh
@@ -0,0 +1,223 @@
+#pragma once
+
+#include <musa_runtime.h>
+
+namespace vllm_musa {
+
+// FULL_MASK
+#define MASK_SHFL_128 ((~(128 - 1)) & 0x7f) << 7 | (128 - 1)
+#define MASK_SHFL_64 ((~(64 - 1)) & 0x7f) << 7 | (64 - 1)
+#define MASK_SHFL_32 ((~(32 - 1)) & 0x7f) << 7 | (32 - 1)
+#define MASK_SHFL_16 ((~(16 - 1)) & 0x7f) << 7 | (16 - 1)
+#define MASK_SHFL_8 ((~(8 - 1)) & 0x7f) << 7 | (8 - 1)
+#define MASK_SHFL_4 ((~(4 - 1)) & 0x7f) << 7 | (4 - 1)
+#define MASK_SHFL_2 ((~(2 - 1)) & 0x7f) << 7 | (2 - 1)
+
+#define MASK_SHFL_UP_128 ((~(128 - 1)) & 0x7f) << 7
+#define MASK_SHFL_UP_64 ((~(64 - 1)) & 0x7f) << 7
+#define MASK_SHFL_UP_32 ((~(32 - 1)) & 0x7f) << 7
+#define MASK_SHFL_UP_16 ((~(16 - 1)) & 0x7f) << 7
+#define MASK_SHFL_UP_8 ((~(8 - 1)) & 0x7f) << 7
+#define MASK_SHFL_UP_4 ((~(4 - 1)) & 0x7f) << 7
+#define MASK_SHFL_UP_2 ((~(2 - 1)) & 0x7f) << 7
+
+template <typename T, int width>
+__device__ __forceinline__ T mudnn_shfl_sync(T val, int src_lane) {
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 220)) // MUSIFY_EXCL_LINE
+  return __shfl_sync(0xffffffff, val, src_lane, width);
+#elif (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 210)) // MUSIFY_EXCL_LINE
+  // width must have a value which is a power of two in the range [1, warpSize]
+  static_assert((width >= 1) && (width <= 128) && ((width & (width - 1)) == 0));
+  if constexpr (width == 1) {
+    return val;
+  } else {
+    auto shfl_func = [&](int &var) {
+      if constexpr (width == 128) {
+        return __musa_shfl_idx_sync_i32(var, src_lane & 0x7f, MASK_SHFL_128);
+      } else if constexpr (width == 64) {
+        return __musa_shfl_idx_sync_i32(var, src_lane & 0x3f, MASK_SHFL_64);
+      } else if constexpr (width == 32) {
+        return __musa_shfl_idx_sync_i32(var, src_lane & 0x1f, MASK_SHFL_32);
+      } else if constexpr (width == 16) {
+        return __musa_shfl_idx_sync_i32(var, src_lane & 0xf, MASK_SHFL_16);
+      } else if constexpr (width == 8) {
+        return __musa_shfl_idx_sync_i32(var, src_lane & 0x7, MASK_SHFL_8);
+      } else if constexpr (width == 4) {
+        return __musa_shfl_idx_sync_i32(var, src_lane & 0x3, MASK_SHFL_4);
+      } else if constexpr (width == 2) {
+        return __musa_shfl_idx_sync_i32(var, src_lane & 0x1, MASK_SHFL_2);
+      }
+    };
+    if constexpr (sizeof(T) == 4) {
+      int var = *(reinterpret_cast<int32_t *>(&val));
+      int ret = shfl_func(var);
+      return *(reinterpret_cast<T *>(&ret));
+    } else {
+      struct __Bits {
+        int __a, __b;
+      };
+      __Bits __tmp;
+      memcpy(&__tmp, &val, sizeof(val));
+      __tmp.__a = shfl_func(__tmp.__a);
+      __tmp.__b = shfl_func(__tmp.__b);
+      int64_t ret = *(reinterpret_cast<int64_t *>(&__tmp));
+      return *(reinterpret_cast<T *>(&ret));
+    }
+  }
+#else
+  return 0;
+#endif
+}
+
+template <typename T, int width>
+__device__ __forceinline__ T mudnn_shfl_down_sync(T val, unsigned int delta) {
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 220)) // MUSIFY_EXCL_LINE
+  return __shfl_down_sync(0xffffffff, val, delta, width);
+#elif (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 210)) // MUSIFY_EXCL_LINE
+  static_assert((width >= 2) && (width <= 128) && ((width & (width - 1)) == 0));
+
+  auto shfl_func = [&](int &var) {
+    if constexpr (width == 128) {
+      return __musa_shfl_down_sync_i32(var, delta & 0x7f, MASK_SHFL_128);
+    } else if constexpr (width == 64) {
+      return __musa_shfl_down_sync_i32(var, delta & 0x3f, MASK_SHFL_64);
+    } else if constexpr (width == 32) {
+      return __musa_shfl_down_sync_i32(var, delta & 0x1f, MASK_SHFL_32);
+    } else if constexpr (width == 16) {
+      return __musa_shfl_down_sync_i32(var, delta & 0xf, MASK_SHFL_16);
+    } else if constexpr (width == 8) {
+      return __musa_shfl_down_sync_i32(var, delta & 0x7, MASK_SHFL_8);
+    } else if constexpr (width == 4) {
+      return __musa_shfl_down_sync_i32(var, delta & 0x3, MASK_SHFL_4);
+    } else if constexpr (width == 2) {
+      return __musa_shfl_down_sync_i32(var, delta & 0x1, MASK_SHFL_2);
+    }
+  };
+
+  if constexpr (sizeof(T) == 4) {
+    int var = *(reinterpret_cast<int32_t *>(&val));
+    int ret = shfl_func(var);
+    return *(reinterpret_cast<T *>(&ret));
+  } else {
+    struct __Bits {
+      int __a, __b;
+    };
+    __Bits __tmp;
+    memcpy(&__tmp, &val, sizeof(val));
+    __tmp.__a = shfl_func(__tmp.__a);
+    __tmp.__b = shfl_func(__tmp.__b);
+    int64_t ret = *(reinterpret_cast<int64_t *>(&__tmp));
+    return *(reinterpret_cast<T *>(&ret));
+  }
+#else
+  return 0;
+#endif
+}
+
+template <typename T, int width>
+__device__ __forceinline__ T mudnn_shfl_xor_sync(T val, int lane_mask) {
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 220)) // MUSIFY_EXCL_LINE
+  return __shfl_xor_sync(0xffffffff, val, lane_mask, width);
+#elif (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 210)) // MUSIFY_EXCL_LINE
+  static_assert((width >= 2) && (width <= 128) && ((width & (width - 1)) == 0));
+
+  auto shfl_func = [&](int &var) {
+    if constexpr (width == 128) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x7f, MASK_SHFL_128);
+    } else if constexpr (width == 64) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x3f, MASK_SHFL_64);
+    } else if constexpr (width == 32) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x1f, MASK_SHFL_32);
+    } else if constexpr (width == 16) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0xf, MASK_SHFL_16);
+    } else if constexpr (width == 8) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x7, MASK_SHFL_8);
+    } else if constexpr (width == 4) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x3, MASK_SHFL_4);
+    } else if constexpr (width == 2) {
+      return __musa_shfl_xor_sync_i32(var, lane_mask & 0x1, MASK_SHFL_2);
+    }
+  };
+
+  if constexpr (sizeof(T) == 4) {
+    int var = *(reinterpret_cast<int32_t *>(&val));
+    int ret = shfl_func(var);
+    return *(reinterpret_cast<T *>(&ret));
+  } else {
+    struct __Bits {
+      int __a, __b;
+    };
+    __Bits __tmp;
+    memcpy(&__tmp, &val, sizeof(val));
+    __tmp.__a = shfl_func(__tmp.__a);
+    __tmp.__b = shfl_func(__tmp.__b);
+    int64_t ret = *(reinterpret_cast<int64_t *>(&__tmp));
+    return *(reinterpret_cast<T *>(&ret));
+  }
+#else
+  return 0;
+#endif
+}
+
+template <typename T, int width>
+__device__ __forceinline__ T mudnn_shfl_up_sync(T val, unsigned int delta) {
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 220)) // MUSIFY_EXCL_LINE
+  return __shfl_up_sync(0xffffffff, val, delta, width);
+#elif (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ > 210)) // MUSIFY_EXCL_LINE
+  static_assert((width >= 2) && (width <= 128) && ((width & (width - 1)) == 0));
+
+  auto shfl_func = [&](int &var) {
+    if constexpr (width == 128) {
+      return __musa_shfl_up_sync_i32(var, delta & 0x7f, MASK_SHFL_UP_128);
+    } else if constexpr (width == 64) {
+      return __musa_shfl_up_sync_i32(var, delta & 0x3f, MASK_SHFL_UP_64);
+    } else if constexpr (width == 32) {
+      return __musa_shfl_up_sync_i32(var, delta & 0x1f, MASK_SHFL_UP_32);
+    } else if constexpr (width == 16) {
+      return __musa_shfl_up_sync_i32(var, delta & 0xf, MASK_SHFL_UP_16);
+    } else if constexpr (width == 8) {
+      return __musa_shfl_up_sync_i32(var, delta & 0x7, MASK_SHFL_UP_8);
+    } else if constexpr (width == 4) {
+      return __musa_shfl_up_sync_i32(var, delta & 0x3, MASK_SHFL_UP_4);
+    } else if constexpr (width == 2) {
+      return __musa_shfl_up_sync_i32(var, delta & 0x1, MASK_SHFL_UP_2);
+    }
+  };
+
+  if constexpr (sizeof(T) == 4) {
+    int var = *(reinterpret_cast<int32_t *>(&val));
+    int ret = shfl_func(var);
+    return *(reinterpret_cast<T *>(&ret));
+  } else {
+    struct __Bits {
+      int __a, __b;
+    };
+    __Bits __tmp;
+    memcpy(&__tmp, &val, sizeof(val));
+    __tmp.__a = shfl_func(__tmp.__a);
+    __tmp.__b = shfl_func(__tmp.__b);
+    int64_t ret = *(reinterpret_cast<int64_t *>(&__tmp));
+    return *(reinterpret_cast<T *>(&ret));
+  }
+#else
+  return 0;
+#endif
+}
+
+#undef MASK_SHFL_128
+#undef MASK_SHFL_64
+#undef MASK_SHFL_32
+#undef MASK_SHFL_16
+#undef MASK_SHFL_8
+#undef MASK_SHFL_4
+#undef MASK_SHFL_2
+
+#undef MASK_SHFL_UP_128
+#undef MASK_SHFL_UP_64
+#undef MASK_SHFL_UP_32
+#undef MASK_SHFL_UP_16
+#undef MASK_SHFL_UP_8
+#undef MASK_SHFL_UP_4
+#undef MASK_SHFL_UP_2
+
+} // namespace vllm_musa
diff --git a/vllm_musa/csrc/rmsnorm_cast/update_helper.muh b/vllm_musa/csrc/rmsnorm_cast/update_helper.muh
new file mode 100644
index 000000000..8b9bf090a
--- /dev/null
+++ b/vllm_musa/csrc/rmsnorm_cast/update_helper.muh
@@ -0,0 +1,230 @@
+
+#pragma once
+
+#include "dtype.muh"
+#include "math.muh"
+#include "shuffle.muh"
+#include <stdint.h>
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ <= 220 // MUSIFY_EXCL_LINE
+#define WARP_THREAD 128
+#else
+#define WARP_THREAD 32
+#endif
+
+#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ == 310 // MUSIFY_EXCL_LINE
+#define __SYNCTHREADS_LM __syncthreads_lm()
+#else
+#define __SYNCTHREADS_LM __syncthreads()
+#endif
+
+namespace vllm_musa {
+#define DEVICE_INLINE __device__ __forceinline__
+
+enum class VarUpdateMode { WELFORD, WELFORD_ONLY_MEAN, CHAN, CHAN_ONLY_MEAN };
+
+template <typename T, VarUpdateMode Mode> struct VarUpdate;
+
+template <typename T> struct VarUpdate<T, VarUpdateMode::WELFORD> {
+  DEVICE_INLINE void apply(T curr, T *mu, T *m2, T *cnt) {
+    *cnt += 1;
+    T delta = curr - *mu;
+    *mu += fast_divf(delta, *cnt);
+    T delta2 = curr - *mu;
+    *m2 += delta * delta2;
+  }
+};
+
+template <typename T> struct VarUpdate<T, VarUpdateMode::WELFORD_ONLY_MEAN> {
+  DEVICE_INLINE void apply(T curr, T *mu, T *cnt) {
+    *cnt += 1;
+    T delta = curr - *mu;
+    *mu += fast_divf(delta, *cnt);
+  }
+};
+
+template <typename T> struct VarUpdate<T, VarUpdateMode::CHAN> {
+  DEVICE_INLINE void apply(T mu_B, T m2_B, T cnt_B, T *mu, T *m2, T *cnt) {
+    if (cnt_B > 0) {
+      T n_AB = cnt_B + (*cnt);
+      T delta = mu_B - (*mu);
+      *mu += delta * fast_divf(cnt_B, n_AB);
+      *m2 += m2_B + fast_divf(delta * delta * cnt_B * (*cnt), n_AB);
+      *cnt = n_AB;
+    }
+  }
+};
+
+template <typename T> struct VarUpdate<T, VarUpdateMode::CHAN_ONLY_MEAN> {
+  DEVICE_INLINE void apply(T mu_B, T cnt_B, T *mu, T *cnt) {
+    if (cnt_B > 0) {
+      T n_AB = cnt_B + (*cnt);
+      T delta = mu_B - (*mu);
+      *mu += delta * fast_divf(cnt_B, n_AB);
+      *cnt = n_AB;
+    }
+  }
+};
+
+enum class ReduceMode { WELFORD, WELFORD_ONLY_MEAN, SUM, MAX };
+enum class ReduceAxis { COLUMN, ABC };
+
+template <typename ComputeType, int BLOCK_X, int BLOCK_Y, ReduceMode Mode,
+          ReduceAxis Axis, int Vlen = 1>
+struct AllReduceOp;
+
+#if (defined(__MUSA_ARCH__) && (__MUSA_ARCH__ >= 220)) // >=QY2
+#define SMEM_STOP (WARP_THREAD / 2)
+#define SHFL_START min(WARP_THREAD / 2, BLOCK_X / 2)
+#else // <= QY1
+#define SMEM_STOP 0
+#define SHFL_START 0
+#endif
+
+template <typename ComputeType, int BLOCK_X, int BLOCK_Y, ReduceAxis Axis,
+          int Vlen>
+struct AllReduceOp<ComputeType, BLOCK_X, BLOCK_Y, ReduceMode::WELFORD, Axis,
+                   Vlen> {
+  DEVICE_INLINE void apply(ComputeType *mu, ComputeType *m2, ComputeType *cnt,
+                           int tx, int ty) {
+    __shared__ ComputeType __attribute__((aligned(16)))
+    smem[BLOCK_X * BLOCK_Y * 3];
+    ComputeType *smem_mu = &smem[0];
+    ComputeType *smem_m2 = &smem[BLOCK_X * BLOCK_Y];
+    ComputeType *smem_cnt = &smem[BLOCK_X * BLOCK_Y * 2];
+
+    VarUpdate<ComputeType, VarUpdateMode::CHAN> chan_op;
+
+    if constexpr (Axis == ReduceAxis::COLUMN) {
+      ComputeType mu_B;
+      ComputeType m2_B;
+      ComputeType cnt_B;
+#pragma unroll
+      for (int offset = BLOCK_X / 2; offset > SMEM_STOP; offset /= 2) {
+        if ((tx < offset * 2) && (tx >= offset)) {
+          smem_mu[ty * BLOCK_X + tx] = *mu;
+          smem_m2[ty * BLOCK_X + tx] = *m2;
+          smem_cnt[ty * BLOCK_X + tx] = *cnt;
+        }
+        __SYNCTHREADS_LM;
+        if (tx < offset) {
+          mu_B = smem_mu[ty * BLOCK_X + tx + offset];
+          m2_B = smem_m2[ty * BLOCK_X + tx + offset];
+          cnt_B = smem_cnt[ty * BLOCK_X + tx + offset];
+          chan_op.apply(mu_B, m2_B, cnt_B, mu, m2, cnt);
+        }
+      }
+#if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ >= 220))
+#pragma unroll
+      for (int offset = SHFL_START; offset > 0; offset /= 2) {
+        mu_B = mudnn_shfl_down_sync<ComputeType, WARP_THREAD>(*mu, offset);
+        m2_B = mudnn_shfl_down_sync<ComputeType, WARP_THREAD>(*m2, offset);
+        cnt_B = mudnn_shfl_down_sync<ComputeType, WARP_THREAD>(*cnt, offset);
+        chan_op.apply(mu_B, m2_B, cnt_B, mu, m2, cnt);
+      }
+#endif
+      if (tx == 0) {
+        smem_mu[ty * BLOCK_X + tx] = *mu;
+        smem_m2[ty * BLOCK_X + tx] = *m2;
+        smem_cnt[ty * BLOCK_X + tx] = *cnt;
+      }
+      __SYNCTHREADS_LM;
+      *mu = smem_mu[ty * BLOCK_X];
+      *m2 = smem_m2[ty * BLOCK_X];
+      *cnt = smem_cnt[ty * BLOCK_X];
+    } else {
+      smem_mu[ty * BLOCK_X + tx] = *mu;
+      smem_m2[ty * BLOCK_X + tx] = *m2;
+      smem_cnt[ty * BLOCK_X + tx] = *cnt;
+      __SYNCTHREADS_LM;
+
+#pragma unroll
+      for (int offset = BLOCK_Y / 2; offset > 0; offset /= 2) {
+        if (ty < offset) {
+          ComputeType mu_B = smem_mu[(ty + offset) * BLOCK_X + tx];
+          ComputeType m2_B = smem_m2[(ty + offset) * BLOCK_X + tx];
+          ComputeType cnt_B = smem_cnt[(ty + offset) * BLOCK_X + tx];
+          chan_op.apply(mu_B, m2_B, cnt_B, mu, m2, cnt);
+          smem_mu[ty * BLOCK_X + tx] = *mu;
+          smem_m2[ty * BLOCK_X + tx] = *m2;
+          smem_cnt[ty * BLOCK_X + tx] = *cnt;
+        }
+        __SYNCTHREADS_LM;
+      }
+
+      *mu = smem_mu[tx];
+      *m2 = smem_m2[tx];
+      *cnt = smem_cnt[tx];
+    }
+  }
+};
+
+template <typename ComputeType, int BLOCK_X, int BLOCK_Y, ReduceAxis Axis,
+          int Vlen>
+struct AllReduceOp<ComputeType, BLOCK_X, BLOCK_Y, ReduceMode::WELFORD_ONLY_MEAN,
+                   Axis, Vlen> {
+  DEVICE_INLINE void apply(ComputeType *mu, ComputeType *cnt, int tx, int ty) {
+    __shared__ ComputeType __attribute__((aligned(16)))
+    smem[BLOCK_X * BLOCK_Y * 2];
+    ComputeType *smem_mu = &smem[0];
+    ComputeType *smem_cnt = &smem[BLOCK_X * BLOCK_Y];
+
+    VarUpdate<ComputeType, VarUpdateMode::CHAN_ONLY_MEAN> chan_op;
+
+    if constexpr (Axis == ReduceAxis::COLUMN) {
+      ComputeType mu_B;
+      ComputeType cnt_B;
+#pragma unroll
+      for (int offset = BLOCK_X / 2; offset > SMEM_STOP; offset /= 2) {
+        if ((tx < offset * 2) && (tx >= offset)) {
+          smem_mu[ty * BLOCK_X + tx] = *mu;
+          smem_cnt[ty * BLOCK_X + tx] = *cnt;
+        }
+        __SYNCTHREADS_LM;
+        if (tx < offset) {
+          mu_B = smem_mu[ty * BLOCK_X + tx + offset];
+          cnt_B = smem_cnt[ty * BLOCK_X + tx + offset];
+          chan_op.apply(mu_B, cnt_B, mu, cnt);
+        }
+      }
+#if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ >= 220))
+#pragma unroll
+      for (int offset = SHFL_START; offset > 0; offset /= 2) {
+        mu_B = mudnn_shfl_down_sync<ComputeType, WARP_THREAD>(*mu, offset);
+        cnt_B = mudnn_shfl_down_sync<ComputeType, WARP_THREAD>(*cnt, offset);
+        chan_op.apply(mu_B, cnt_B, mu, cnt);
+      }
+#endif
+      if (tx == 0) {
+        smem_mu[ty * BLOCK_X + tx] = *mu;
+        smem_cnt[ty * BLOCK_X + tx] = *cnt;
+      }
+      __SYNCTHREADS_LM;
+      *mu = smem_mu[ty * BLOCK_X];
+      *cnt = smem_cnt[ty * BLOCK_X];
+    } else {
+      smem_mu[ty * BLOCK_X + tx] = *mu;
+      smem_cnt[ty * BLOCK_X + tx] = *cnt;
+      __SYNCTHREADS_LM;
+
+#pragma unroll
+      for (int offset = BLOCK_Y / 2; offset > 0; offset /= 2) {
+        if (ty < offset) {
+          ComputeType mu_B = smem_mu[(ty + offset) * BLOCK_X + tx];
+          ComputeType cnt_B = smem_cnt[(ty + offset) * BLOCK_X + tx];
+          chan_op.apply(mu_B, cnt_B, mu, cnt);
+          smem_mu[ty * BLOCK_X + tx] = *mu;
+          smem_cnt[ty * BLOCK_X + tx] = *cnt;
+        }
+        __SYNCTHREADS_LM;
+      }
+
+      *mu = smem_mu[tx];
+      *cnt = smem_cnt[tx];
+    }
+  }
+};
+
+#undef SMEM_STOP
+#undef SHFL_START
+} // namespace vllm_musa
diff --git a/vllm_musa/csrc/type_convert.muh b/vllm_musa/csrc/type_convert.muh
new file mode 100644
index 000000000..aa71c3565
--- /dev/null
+++ b/vllm_musa/csrc/type_convert.muh
@@ -0,0 +1,148 @@
+#pragma once
+
+#include <torch/all.h>
+
+#include "musa.h"
+
+namespace vllm_musa {
+/* Converter structs for the conversion from torch types to HIP/CUDA types,
+   and the associated type conversions within HIP/CUDA. These helpers need
+   to be implemented for now because the relevant type conversion
+   operators/constructors are not consistently implemented by HIP/CUDA, so
+   a generic conversion via type casts cannot be implemented.
+
+   Each struct should have the member static constexpr bool `exists`:
+   If false, the optimized kernel is not used for the corresponding torch type.
+   If true, the struct should be fully defined as shown in the examples below.
+ */
+template <typename torch_type>
+struct _typeConvert {
+  static constexpr bool exists = false;
+};
+
+template <>
+struct _typeConvert<c10::Half> {
+  static constexpr bool exists = true;
+  using hip_type = __half;
+  using packed_hip_type = __half2;
+
+  __device__ static inline float convert(hip_type x) { return __half2float(x); }
+  __device__ static inline float2 convert(packed_hip_type x) {
+    return __half22float2(x);
+  }
+  __device__ static inline hip_type convert(float x) {
+    return __float2half_rn(x);
+  }
+  __device__ static inline packed_hip_type convert(float2 x) {
+    return __float22half2_rn(x);
+  }
+};
+
+template <>
+struct _typeConvert<c10::BFloat16> {
+  static constexpr bool exists = true;
+  using hip_type = __nv_bfloat16;
+  using packed_hip_type = __nv_bfloat162;
+
+  __device__ static inline float convert(hip_type x) {
+    return __bfloat162float(x);
+  }
+  __device__ static inline float2 convert(packed_hip_type x) {
+    return __bfloat1622float2(x);
+  }
+  __device__ static inline hip_type convert(float x) {
+    return __float2bfloat16(x);
+  }
+  __device__ static inline packed_hip_type convert(float2 x) {
+    return __float22bfloat162_rn(x);
+  }
+};
+
+/* Vector POD struct to generate vectorized and packed FP16/BF16 ops
+   for appropriate specializations of fused_add_rms_norm_kernel.
+   Only functions that are necessary in that kernel are implemented.
+   Alignment to 16 bytes is required to use 128-bit global memory ops.
+ */
+template <typename scalar_t, int width>
+struct alignas(16) _f16Vec {
+  /* Not theoretically necessary that width is a power of 2 but should
+     almost always be the case for optimization purposes */
+  static_assert(width > 0 && (width & (width - 1)) == 0,
+                "Width is not a positive power of 2!");
+  using Converter = _typeConvert<scalar_t>;
+  using T1 = typename Converter::hip_type;
+  using T2 = typename Converter::packed_hip_type;
+  T1 data[width];
+
+  __device__ _f16Vec& operator+=(const _f16Vec<scalar_t, width>& other) {
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        T2 temp{data[i], data[i + 1]};
+        temp += T2{other.data[i], other.data[i + 1]};
+        data[i] = temp.x;
+        data[i + 1] = temp.y;
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < width; ++i) data[i] += other.data[i];
+    }
+    return *this;
+  }
+
+  __device__ _f16Vec& operator*=(const _f16Vec<scalar_t, width>& other) {
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        T2 temp{data[i], data[i + 1]};
+        temp *= T2{other.data[i], other.data[i + 1]};
+        data[i] = temp.x;
+        data[i + 1] = temp.y;
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < width; ++i) data[i] *= other.data[i];
+    }
+    return *this;
+  }
+
+  __device__ _f16Vec& operator*=(const float scale) {
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        float2 temp_f = Converter::convert(T2{data[i], data[i + 1]});
+        temp_f.x *= scale;
+        temp_f.y *= scale;
+        T2 temp = Converter::convert(temp_f);
+        data[i] = temp.x;
+        data[i + 1] = temp.y;
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < width; ++i) {
+        float temp = Converter::convert(data[i]) * scale;
+        data[i] = Converter::convert(temp);
+      }
+    }
+    return *this;
+  }
+
+  __device__ float sum_squares() const {
+    float result = 0.0f;
+    if constexpr (width % 2 == 0) {
+#pragma unroll
+      for (int i = 0; i < width; i += 2) {
+        float2 z = Converter::convert(T2{data[i], data[i + 1]});
+        result += z.x * z.x + z.y * z.y;
+      }
+    } else {
+#pragma unroll
+      for (int i = 0; i < width; ++i) {
+        float x = Converter::convert(data[i]);
+        result += x * x;
+      }
+    }
+    return result;
+  }
+};
+}  // namespace vllm_musa
\ No newline at end of file
diff --git a/vllm_musa/format.sh b/vllm_musa/format.sh
new file mode 100755
index 000000000..413cb1f86
--- /dev/null
+++ b/vllm_musa/format.sh
@@ -0,0 +1,341 @@
+#!/usr/bin/env bash
+
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+# Adapted from https://github.com/vllm-project/vllm/tree/main/tools
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+# YAPF formatter, adapted from ray and skypilot.
+#
+# Usage:
+#    # Do work and commit your work.
+
+#    # Format files that differ from origin/main.
+#    bash format.sh
+
+#    # Commit changed files with message 'Run yapf and ruff'
+#
+#
+# YAPF + Clang formatter (if installed). This script formats all changed files from the last mergebase.
+# You are encouraged to run this locally before pushing changes for review.
+
+# Cause the script to exit if a single command fails
+set -eo pipefail
+
+# this stops git rev-parse from failing if we run this from the .git directory
+builtin cd "$(dirname "${BASH_SOURCE:-$0}")"
+ROOT="$(git rev-parse --show-toplevel)"
+builtin cd "$ROOT" || exit 1
+
+check_command() {
+    if ! command -v "$1" &> /dev/null; then
+        echo "❓❓$1 is not installed, please run \`pip install -r requirements-lint.txt\`"
+        exit 1
+    fi
+}
+
+check_command yapf
+check_command ruff
+check_command mypy
+check_command codespell
+check_command isort
+check_command clang-format
+
+YAPF_VERSION=$(yapf --version | awk '{print $2}')
+RUFF_VERSION=$(ruff --version | awk '{print $2}')
+MYPY_VERSION=$(mypy --version | awk '{print $2}')
+CODESPELL_VERSION=$(codespell --version)
+ISORT_VERSION=$(isort --vn)
+CLANGFORMAT_VERSION=$(clang-format --version | awk '{print $3}')
+SPHINX_LINT_VERSION=$(sphinx-lint --version | awk '{print $2}')
+
+# params: tool name, tool version, required version
+tool_version_check() {
+    expected=$(grep "$1" requirements-lint.txt | cut -d'=' -f3)
+    if [[ "$2" != "$expected" ]]; then
+        echo "❓❓Wrong $1 version installed: $expected is required, not $2."
+        exit 1
+    fi
+}
+
+tool_version_check "yapf" "$YAPF_VERSION"
+tool_version_check "ruff" "$RUFF_VERSION"
+tool_version_check "mypy" "$MYPY_VERSION"
+tool_version_check "isort" "$ISORT_VERSION"
+tool_version_check "codespell" "$CODESPELL_VERSION"
+tool_version_check "clang-format" "$CLANGFORMAT_VERSION"
+tool_version_check "sphinx-lint" "$SPHINX_LINT_VERSION"
+
+YAPF_FLAGS=(
+    '--recursive'
+    '--parallel'
+)
+
+YAPF_EXCLUDES=(
+    '--exclude' 'build/**'
+)
+
+# Format specified files
+format() {
+    yapf --in-place "${YAPF_FLAGS[@]}" "$@"
+}
+
+# Format files that differ from main branch. Ignores dirs that are not slated
+# for autoformat yet.
+format_changed() {
+    # The `if` guard ensures that the list of filenames is not empty, which
+    # could cause yapf to receive 0 positional arguments, making it hang
+    # waiting for STDIN.
+    #
+    # `diff-filter=ACM` and $MERGEBASE is to ensure we only format files that
+    # exist on both branches.
+    MERGEBASE="$(git merge-base origin/main HEAD)"
+
+    if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
+        git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs -P 5 \
+             yapf --in-place "${YAPF_EXCLUDES[@]}" "${YAPF_FLAGS[@]}"
+    fi
+
+}
+
+# Format all files
+format_all() {
+    yapf --in-place "${YAPF_FLAGS[@]}" "${YAPF_EXCLUDES[@]}" .
+}
+
+## This flag formats individual files. --files *must* be the first command line
+## arg to use this option.
+if [[ "$1" == '--files' ]]; then
+   format "${@:2}"
+   # If `--all` is passed, then any further arguments are ignored and the
+   # entire python directory is formatted.
+elif [[ "$1" == '--all' ]]; then
+   format_all
+else
+   # Format only the files that changed in last commit.
+   format_changed
+fi
+echo 'vLLM yapf: Done'
+
+# Run mypy
+echo 'vLLM mypy:'
+tools/mypy.sh
+echo 'vLLM mypy: Done'
+
+
+# If git diff returns a file that is in the skip list, the file may be checked anyway:
+# https://github.com/codespell-project/codespell/issues/1915
+# Avoiding the "./" prefix and using "/**" globs for directories appears to solve the problem
+CODESPELL_EXCLUDES=(
+    '--skip' 'tests/prompts/**,./benchmarks/sonnet.txt,*tests/lora/data/**,build/**'
+)
+
+# check spelling of specified files
+spell_check() {
+    codespell "$@"
+}
+
+spell_check_all(){
+  codespell --toml pyproject.toml "${CODESPELL_EXCLUDES[@]}"
+}
+
+# Spelling check of files that differ from main branch.
+spell_check_changed() {
+    # The `if` guard ensures that the list of filenames is not empty, which
+    # could cause ruff to receive 0 positional arguments, making it hang
+    # waiting for STDIN.
+    #
+    # `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that
+    # exist on both branches.
+    MERGEBASE="$(git merge-base origin/main HEAD)"
+    if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
+        git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \
+            codespell "${CODESPELL_EXCLUDES[@]}"
+    fi
+}
+
+# Run Codespell
+## This flag runs spell check of individual files. --files *must* be the first command line
+## arg to use this option.
+if [[ "$1" == '--files' ]]; then
+   spell_check "${@:2}"
+   # If `--all` is passed, then any further arguments are ignored and the
+   # entire python directory is linted.
+elif [[ "$1" == '--all' ]]; then
+   spell_check_all
+else
+   # Check spelling only of the files that changed in last commit.
+   spell_check_changed
+fi
+echo 'vLLM codespell: Done'
+
+
+# Lint specified files
+lint() {
+    ruff check "$@"
+}
+
+# Lint files that differ from main branch. Ignores dirs that are not slated
+# for autolint yet.
+lint_changed() {
+    # The `if` guard ensures that the list of filenames is not empty, which
+    # could cause ruff to receive 0 positional arguments, making it hang
+    # waiting for STDIN.
+    #
+    # `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that
+    # exist on both branches.
+    MERGEBASE="$(git merge-base origin/main HEAD)"
+
+    if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
+        git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \
+             ruff check
+    fi
+
+}
+
+# Run Ruff
+### This flag lints individual files. --files *must* be the first command line
+### arg to use this option.
+if [[ "$1" == '--files' ]]; then
+   lint "${@:2}"
+   # If `--all` is passed, then any further arguments are ignored and the
+   # entire python directory is linted.
+elif [[ "$1" == '--all' ]]; then
+   lint vllm tests
+else
+   # Format only the files that changed in last commit.
+   lint_changed
+fi
+echo 'vLLM ruff: Done'
+
+# check spelling of specified files
+isort_check() {
+    isort "$@"
+}
+
+isort_check_all(){
+  isort .
+}
+
+# Spelling  check of files that differ from main branch.
+isort_check_changed() {
+    # The `if` guard ensures that the list of filenames is not empty, which
+    # could cause ruff to receive 0 positional arguments, making it hang
+    # waiting for STDIN.
+    #
+    # `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that
+    # exist on both branches.
+    MERGEBASE="$(git merge-base origin/main HEAD)"
+
+    if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then
+        git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \
+             isort
+    fi
+}
+
+# Run Isort
+# This flag runs spell check of individual files. --files *must* be the first command line
+# arg to use this option.
+if [[ "$1" == '--files' ]]; then
+   isort_check "${@:2}"
+   # If `--all` is passed, then any further arguments are ignored and the
+   # entire python directory is linted.
+elif [[ "$1" == '--all' ]]; then
+   isort_check_all
+else
+   # Check spelling only of the files that changed in last commit.
+   isort_check_changed
+fi
+echo 'vLLM isort: Done'
+
+# Clang-format section
+# Exclude some files for formatting because they are vendored
+# NOTE: Keep up to date with .github/workflows/clang-format.yml
+CLANG_FORMAT_EXCLUDES=(
+    'csrc/moe/topk_softmax_kernels.cu'
+    'csrc/quantization/gguf/ggml-common.h'
+    'csrc/quantization/gguf/dequantize.cuh'
+    'csrc/quantization/gguf/vecdotq.cuh'
+    'csrc/quantization/gguf/mmq.cuh'
+    'csrc/quantization/gguf/mmvq.cuh'
+)
+
+# Format specified files with clang-format
+clang_format() {
+    clang-format -i "$@"
+}
+
+# Format files that differ from main branch with clang-format.
+clang_format_changed() {
+    # The `if` guard ensures that the list of filenames is not empty, which
+    # could cause clang-format to receive 0 positional arguments, making it hang
+    # waiting for STDIN.
+    #
+    # `diff-filter=ACM` and $MERGEBASE is to ensure we only format files that
+    # exist on both branches.
+    MERGEBASE="$(git merge-base origin/main HEAD)"
+
+    # Get the list of changed files, excluding the specified ones
+    changed_files=$(git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.h' '*.cpp' '*.cu' '*.cuh' | (grep -vFf <(printf "%s\n" "${CLANG_FORMAT_EXCLUDES[@]}") || echo -e))
+    if [ -n "$changed_files" ]; then
+        echo "$changed_files" | xargs -P 5 clang-format -i
+    fi
+}
+
+# Format all files with clang-format
+clang_format_all() {
+    find csrc/ \( -name '*.h' -o -name '*.cpp' -o -name '*.cu' -o -name '*.cuh' \) -print \
+        | grep -vFf <(printf "%s\n" "${CLANG_FORMAT_EXCLUDES[@]}") \
+        | xargs clang-format -i
+}
+
+# Run clang-format
+if [[ "$1" == '--files' ]]; then
+   clang_format "${@:2}"
+elif [[ "$1" == '--all' ]]; then
+   clang_format_all
+else
+   clang_format_changed
+fi
+echo 'vLLM clang-format: Done'
+
+echo 'vLLM actionlint:'
+tools/actionlint.sh -color
+echo 'vLLM actionlint: Done'
+
+echo 'vLLM shellcheck:'
+tools/shellcheck.sh
+echo 'vLLM shellcheck: Done'
+
+echo 'excalidraw png check:'
+tools/png-lint.sh
+echo 'excalidraw png check: Done'
+
+if ! git diff --quiet &>/dev/null; then
+    echo
+    echo "🔍🔍There are files changed by the format checker or by you that are not added and committed:"
+    git --no-pager diff --name-only
+    echo "🔍🔍Format checker passed, but please add, commit and push all the files above to include changes made by the format checker."
+
+    exit 1
+else
+    echo "✨🎉 Format check passed! Congratulations! 🎉✨"
+fi
+
+# echo 'vLLM sphinx-lint:'
+# tools/sphinx-lint.sh
+# echo 'vLLM sphinx-lint: Done'
diff --git a/vllm_musa/hack/README b/vllm_musa/hack/README
new file mode 100644
index 000000000..f5a20d2a6
--- /dev/null
+++ b/vllm_musa/hack/README
@@ -0,0 +1,23 @@
+# torch & torch_musa
+## ONLY FOR Torch==2.2.0
+运行 vllm 需要注册 torch.ops 由于 torch 版本较低, 需要拷贝 hack/impl.py 到 /opt/conda/envs/py310/lib/python3.10/site-packages/torch/_custom_op/impl.py
+
+```shell
+cp hack/impl.py /opt/conda/envs/py310/lib/python3.10/site-packages/torch/_custom_op/impl.py
+```
+
+# musa graph
+针对musa graph的需求:
+1. musa runtime patch包（基于kuae0603的ddk和软件栈）， 作者： 范伟
+0603 ddk和软件栈下载地址： [here](https://oss.mthreads.com:9001/buckets/release-ci/browse/Y29tcHV0ZVFBL2N1ZGFfY29tcGF0aWJsZS9DSS9yZWxlYXNlX0tVQUVfMi4wX2Zvcl9QSDFfTTNELzIwMjUtMDYtMDMv)
+musa runtime patch包下载地址： `https://oss.mthreads.com/dependency/mcc-ai-serving/musa_graph_patch/kuae_0603_fp8_musagraph_patch_to_zhiguo.0710.tar.gz`
+安装方式：
+```shell
+tar vxf kuae_0603_fp8_musagraph_patch_to_zhiguo.0710.tar.gz
+cd kuae_0603_fp8_musagraph_patch_to_zhiguo.0710
+./install.sh
+info: musa will be installed to "/usr/local/musa"
+cp: cannot stat './src/shared_include/mupti*': No such file or directory
+cp: cannot stat './src/mupti/mupti_*_cbid.h': No such file or directory
+#上边No such file or directory的报错，不影响
+```
diff --git a/vllm_musa/hack/compiled_dag_node.py b/vllm_musa/hack/compiled_dag_node.py
new file mode 100644
index 000000000..d7438eb51
--- /dev/null
+++ b/vllm_musa/hack/compiled_dag_node.py
@@ -0,0 +1,3318 @@
+import weakref
+import asyncio
+from collections import defaultdict
+from contextlib import nullcontext
+from dataclasses import dataclass, asdict
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    List,
+    Tuple,
+    Union,
+    Optional,
+    Set,
+)
+import logging
+import threading
+import time
+import uuid
+import traceback
+
+from ray.experimental.channel.auto_transport_type import (
+    AutoTransportType,
+    TypeHintResolver,
+)
+import ray.exceptions
+from ray.dag.dag_operation_future import GPUFuture, DAGOperationFuture, ResolvedFuture
+from ray.experimental.channel.cached_channel import CachedChannel
+from ray.experimental.channel.communicator import Communicator
+from ray.dag.constants import (
+    RAY_CGRAPH_ENABLE_NVTX_PROFILING,
+    RAY_CGRAPH_ENABLE_TORCH_PROFILING,
+    RAY_CGRAPH_VISUALIZE_SCHEDULE,
+)
+import ray
+from ray.exceptions import (
+    RayCgraphCapacityExceeded,
+    RayTaskError,
+    RayChannelError,
+    RayChannelTimeoutError,
+)
+from ray.experimental.compiled_dag_ref import (
+    CompiledDAGRef,
+    CompiledDAGFuture,
+    _process_return_vals,
+)
+from ray.experimental.channel import (
+    ChannelContext,
+    ChannelInterface,
+    ChannelOutputType,
+    ReaderInterface,
+    SynchronousReader,
+    WriterInterface,
+    SynchronousWriter,
+    AwaitableBackgroundReader,
+    AwaitableBackgroundWriter,
+    CompiledDAGArgs,
+    CompositeChannel,
+    IntraProcessChannel,
+)
+from ray.util.annotations import DeveloperAPI
+
+from ray.experimental.channel.shared_memory_channel import (
+    SharedMemoryType,
+)
+from ray.experimental.channel.torch_tensor_type import TorchTensorType
+
+from ray.experimental.channel.torch_tensor_nccl_channel import (
+    _init_communicator,
+    _destroy_communicator,
+)
+
+from ray.dag.dag_node_operation import (
+    _DAGNodeOperation,
+    _DAGNodeOperationType,
+    _DAGOperationGraphNode,
+    _build_dag_node_operation_graph,
+    _extract_execution_schedule,
+    _generate_actor_to_execution_schedule,
+    _generate_overlapped_execution_schedule,
+    _visualize_execution_schedule,
+)
+
+from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
+
+if TYPE_CHECKING:
+    import cupy as cp
+
+logger = logging.getLogger(__name__)
+
+# Keep tracking of every compiled dag created during the lifetime of
+# this process. It tracks them as weakref meaning when the compiled dag
+# is GC'ed, it is automatically removed from here. It is used to teardown
+# compiled dags at interpreter shutdown time.
+_compiled_dags = weakref.WeakValueDictionary()
+
+
+# Relying on __del__ doesn't work well upon shutdown because
+# the destructor order is not guaranteed. We call this function
+# upon `ray.worker.shutdown` which is registered to atexit handler
+# so that teardown is properly called before objects are destructed.
+def _shutdown_all_compiled_dags():
+    global _compiled_dags
+    for _, compiled_dag in _compiled_dags.items():
+        # Kill DAG actors to avoid hanging during shutdown if the actor tasks
+        # cannot be cancelled.
+        compiled_dag.teardown(kill_actors=True)
+    _compiled_dags = weakref.WeakValueDictionary()
+
+
+def _check_unused_dag_input_attributes(
+    output_node: "ray.dag.MultiOutputNode", input_attributes: Set[str]
+) -> Set[str]:
+    """
+    Helper function to check that all input attributes are used in the DAG.
+    For example, if the user creates an input attribute by calling
+    InputNode()["x"], we ensure that there is a path from the
+    InputAttributeNode corresponding to "x" to the DAG's output. If an
+    input attribute is not used, throw an error.
+
+    Args:
+        output_node: The starting node for the traversal.
+        input_attributes: A set of attributes accessed by the InputNode.
+    """
+    from ray.dag import InputAttributeNode
+
+    used_attributes = set()
+    visited_nodes = set()
+    stack: List["ray.dag.DAGNode"] = [output_node]
+
+    while stack:
+        current_node = stack.pop()
+        if current_node in visited_nodes:
+            continue
+        visited_nodes.add(current_node)
+
+        if isinstance(current_node, InputAttributeNode):
+            used_attributes.add(current_node.key)
+
+        stack.extend(current_node._upstream_nodes)
+
+    unused_attributes = input_attributes - used_attributes
+    if unused_attributes:
+        unused_attributes_str = ", ".join(str(key) for key in unused_attributes)
+        input_attributes_str = ", ".join(str(key) for key in input_attributes)
+        unused_phrase = "is unused" if len(unused_attributes) == 1 else "are unused"
+
+        raise ValueError(
+            "Compiled Graph expects input to be accessed "
+            f"using all of attributes {input_attributes_str}, "
+            f"but {unused_attributes_str} {unused_phrase}. "
+            "Ensure all input attributes are used and contribute "
+            "to the computation of the Compiled Graph output."
+        )
+
+
+@DeveloperAPI
+def do_allocate_channel(
+    self,
+    reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]],
+    typ: ChannelOutputType,
+    driver_actor_id: Optional[str] = None,
+) -> ChannelInterface:
+    """Generic actor method to allocate an output channel.
+
+    Args:
+        reader_and_node_list: A list of tuples, where each tuple contains a reader
+            actor handle and the node ID where the actor is located.
+        typ: The output type hint for the channel.
+        driver_actor_id: If this channel is read by a driver and that driver is an
+            actual actor, this will be the actor ID of that driver actor.
+
+    Returns:
+        The allocated channel.
+    """
+    # None means it is called from a driver.
+    writer: Optional["ray.actor.ActorHandle"] = None
+    try:
+        writer = ray.get_runtime_context().current_actor
+    except RuntimeError:
+        # This is the driver so there is no current actor handle.
+        pass
+
+    output_channel = typ.create_channel(
+        writer,
+        reader_and_node_list,
+        driver_actor_id,
+    )
+    return output_channel
+
+
+@DeveloperAPI
+def do_exec_tasks(
+    self,
+    tasks: List["ExecutableTask"],
+    schedule: List[_DAGNodeOperation],
+    overlap_gpu_communication: bool = False,
+) -> None:
+    """A generic actor method to begin executing the operations belonging to an
+    actor. This runs an infinite loop to execute each _DAGNodeOperation in the
+    order specified by the schedule. It exits only if the actor dies or an
+    exception is thrown.
+
+    Args:
+        tasks: the executable tasks corresponding to the actor methods.
+        schedule: A list of _DAGNodeOperation that should be executed in order.
+        overlap_gpu_communication: Whether to overlap GPU communication with
+            computation during DAG execution to improve performance.
+    """
+    try:
+        for task in tasks:
+            task.prepare(overlap_gpu_communication=overlap_gpu_communication)
+
+        if RAY_CGRAPH_ENABLE_NVTX_PROFILING:
+            assert (
+                not RAY_CGRAPH_ENABLE_TORCH_PROFILING
+            ), "NVTX and torch profiling cannot be enabled at the same time."
+            try:
+                import nvtx
+            except ImportError:
+                raise ImportError(
+                    "Please install nvtx to enable nsight profiling. "
+                    "You can install it by running `pip install nvtx`."
+                )
+            nvtx_profile = nvtx.Profile()
+            nvtx_profile.enable()
+
+        if RAY_CGRAPH_ENABLE_TORCH_PROFILING:
+            assert (
+                not RAY_CGRAPH_ENABLE_NVTX_PROFILING
+            ), "NVTX and torch profiling cannot be enabled at the same time."
+
+            import torch
+
+            torch_profile = torch.profiler.profile(
+                activities=[
+                    torch.profiler.ProfilerActivity.CPU,
+                    torch.profiler.ProfilerActivity.CUDA,
+                ],
+                with_stack=True,
+                on_trace_ready=torch.profiler.tensorboard_trace_handler(
+                    "compiled_graph_torch_profiles"
+                ),
+            )
+            torch_profile.start()
+            logger.info("Torch profiling started")
+
+        done = False
+        while True:
+            if done:
+                break
+            for operation in schedule:
+                done = tasks[operation.exec_task_idx].exec_operation(
+                    self, operation.type, overlap_gpu_communication
+                )
+                if done:
+                    break
+
+        if RAY_CGRAPH_ENABLE_NVTX_PROFILING:
+            nvtx_profile.disable()
+
+        if RAY_CGRAPH_ENABLE_TORCH_PROFILING:
+            torch_profile.stop()
+            logger.info("Torch profiling stopped")
+    except Exception:
+        logging.exception("Compiled DAG task exited with exception")
+        raise
+
+
+@DeveloperAPI
+def do_profile_tasks(
+    self,
+    tasks: List["ExecutableTask"],
+    schedule: List[_DAGNodeOperation],
+    overlap_gpu_communication: bool = False,
+) -> None:
+    """A generic actor method similar to `do_exec_tasks`, but with profiling enabled.
+
+    Args:
+        tasks: the executable tasks corresponding to the actor methods.
+        schedule: A list of _DAGNodeOperation that should be executed in order.
+        overlap_gpu_communication: Whether to overlap GPU communication with
+            computation during DAG execution to improve performance.
+    """
+    try:
+        for task in tasks:
+            task.prepare(overlap_gpu_communication=overlap_gpu_communication)
+
+        if not hasattr(self, "__ray_cgraph_events"):
+            self.__ray_cgraph_events = []
+
+        done = False
+        while True:
+            if done:
+                break
+            for operation in schedule:
+                start_t = time.perf_counter()
+                task = tasks[operation.exec_task_idx]
+                done = task.exec_operation(
+                    self, operation.type, overlap_gpu_communication
+                )
+                end_t = time.perf_counter()
+
+                self.__ray_cgraph_events.append(
+                    _ExecutableTaskRecord(
+                        actor_classname=self.__class__.__name__,
+                        actor_name=ray.get_runtime_context().get_actor_name(),
+                        actor_id=ray.get_runtime_context().get_actor_id(),
+                        method_name=task.method_name,
+                        bind_index=task.bind_index,
+                        operation=operation.type.value,
+                        start_t=start_t,
+                        end_t=end_t,
+                    )
+                )
+
+                if done:
+                    break
+    except Exception:
+        logging.exception("Compiled DAG task exited with exception")
+        raise
+
+
+@DeveloperAPI
+def do_cancel_executable_tasks(self, tasks: List["ExecutableTask"]) -> None:
+    # CUDA events should be destroyed before other CUDA resources.
+    for task in tasks:
+        task.destroy_cuda_event()
+    for task in tasks:
+        task.cancel()
+
+
+def _wrap_exception(exc):
+    backtrace = ray._private.utils.format_error_message(
+        "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)),
+        task_exception=True,
+    )
+    wrapped = RayTaskError(
+        function_name="do_exec_tasks",
+        traceback_str=backtrace,
+        cause=exc,
+    )
+    return wrapped
+
+
+def _get_nccl_group_id(type_hint: ChannelOutputType) -> Optional[str]:
+    """
+    Get the NCCL group ID from the type hint. If the type hint does not
+    require NCCL, return None.
+
+    Args:
+        type_hint: The type hint of the channel.
+
+    Returns:
+        The NCCL group ID if the type hint requires NCCL, otherwise None.
+    """
+    if type_hint.requires_nccl():
+        assert isinstance(type_hint, TorchTensorType)
+        return type_hint.communicator_id
+    return None
+
+
+def _device_context_manager():
+    """
+    Return a context manager for executing communication operations
+    (i.e., READ and WRITE). For NCCL operations, the context manager
+    uses the proper cuda device from channel context, otherwise,
+    nullcontext will be returned.
+    """
+    if not ChannelContext.get_current().torch_available:
+        return nullcontext()
+
+    import torch
+
+    device = ChannelContext.get_current().torch_device
+
+    #if device.type == "cuda" and torch.cuda.is_available():
+        # In the case of mocked NCCL, we may get a device with type "cuda"
+        # but CUDA is not available. We return nullcontext() in that case,
+        # otherwise torch raises a runtime error if the cuda device context
+        # manager is used.
+        # TODO(rui): consider better mocking NCCL to support device context.
+    #    return torch.cuda.device(device)
+    return nullcontext()
+
+
+@DeveloperAPI
+class CompiledTask:
+    """Wraps the normal Ray DAGNode with some metadata."""
+
+    def __init__(self, idx: int, dag_node: "ray.dag.DAGNode"):
+        """
+        Args:
+            idx: A unique index into the original DAG.
+            dag_node: The original DAG node created by the user.
+        """
+        self.idx = idx
+        self.dag_node = dag_node
+
+        # Dict from task index to actor handle for immediate downstream tasks.
+        self.downstream_task_idxs: Dict[int, "ray.actor.ActorHandle"] = {}
+        # Case 1: The task represents a ClassMethodNode.
+        #
+        # Multiple return values are written to separate `output_channels`.
+        # `output_idxs` represents the tuple index of the output value for
+        # multiple returns in a tuple. If an output index is None, it means
+        # the complete return value is written to the output channel.
+        # Otherwise, the return value is a tuple and the index is used
+        # to extract the value to be written to the output channel.
+        #
+        # Case 2: The task represents an InputNode.
+        #
+        # `output_idxs` can be an integer or a string to retrieve the
+        # corresponding value from `args` or `kwargs` in the DAG's input.
+        self.output_channels: List[ChannelInterface] = []
+        self.output_idxs: List[Optional[Union[int, str]]] = []
+        # The DAGNodes that are arguments to this task.
+        # This is used for lazy resolution of the arguments' type hints.
+        self.arg_nodes: List["ray.dag.DAGNode"] = []
+        # idxs of possible ClassMethodOutputNodes if they exist, used for visualization
+        self.output_node_idxs: List[int] = []
+
+    @property
+    def args(self) -> Tuple[Any]:
+        return self.dag_node.get_args()
+
+    @property
+    def kwargs(self) -> Dict[str, Any]:
+        return self.dag_node.get_kwargs()
+
+    @property
+    def num_readers(self) -> int:
+        return len(self.downstream_task_idxs)
+
+    @property
+    def arg_type_hints(self) -> List["ChannelOutputType"]:
+        return [arg_node.type_hint for arg_node in self.arg_nodes]
+
+    def __str__(self) -> str:
+        return f"""
+            Node: {self.dag_node}
+            Arguments: {self.args}
+            Output: {self.output_channels}
+            """
+
+
+class _ExecutableTaskInput:
+    """Represents an input to an ExecutableTask.
+
+    Args:
+        input_variant: either an unresolved input (when type is ChannelInterface)
+            , or a resolved input value (when type is Any)
+        channel_idx: if input_variant is an unresolved input, this is the index
+            into the input channels list.
+    """
+
+    def __init__(
+        self,
+        input_variant: Union[ChannelInterface, Any],
+        channel_idx: Optional[int],
+    ):
+        self.input_variant = input_variant
+        self.channel_idx = channel_idx
+
+    def resolve(self, channel_results: Any) -> Any:
+        """
+        Resolve the input value from the channel results.
+
+        Args:
+            channel_results: The results from reading the input channels.
+        """
+
+        if isinstance(self.input_variant, ChannelInterface):
+            value = channel_results[self.channel_idx]
+        else:
+            value = self.input_variant
+        return value
+
+
+@DeveloperAPI
+class ExecutableTask:
+    """A task that can be executed in a compiled DAG, and it
+    corresponds to an actor method.
+    """
+
+    def __init__(
+        self,
+        task: "CompiledTask",
+        resolved_args: List[Any],
+        resolved_kwargs: Dict[str, Any],
+    ):
+        """
+        Args:
+            task: The CompiledTask that this ExecutableTask corresponds to.
+            resolved_args: The arguments to the method. Arguments that are
+                not Channels will get passed through to the actor method.
+                If the argument is a channel, it will be replaced by the
+                value read from the channel before the method executes.
+            resolved_kwargs: The keyword arguments to the method. Currently, we
+                do not support binding kwargs to other DAG nodes, so the values
+                of the dictionary cannot be Channels.
+        """
+        from ray.dag import CollectiveOutputNode
+
+        self.method_name = task.dag_node.get_method_name()
+        self.bind_index = task.dag_node._get_bind_index()
+        self.output_channels = task.output_channels
+        self.output_idxs = task.output_idxs
+        self.input_type_hints: List[ChannelOutputType] = task.arg_type_hints
+        self.output_type_hint: ChannelOutputType = task.dag_node.type_hint
+
+        # The NCCL collective operation.
+        self.collective_op: Optional["ray.dag.CollectiveOperation"] = None
+        if isinstance(task.dag_node, CollectiveOutputNode):
+            self.collective_op = task.dag_node.collective_op
+
+        self.input_channels: List[ChannelInterface] = []
+        self.task_inputs: List[_ExecutableTaskInput] = []
+        self.resolved_kwargs: Dict[str, Any] = resolved_kwargs
+        # A unique index which can be used to index into `idx_to_task` to get
+        # the corresponding task.
+        self.task_idx = task.idx
+
+        # Reverse map for input_channels: maps an input channel to
+        # its index in input_channels.
+        input_channel_to_idx: dict[ChannelInterface, int] = {}
+
+        for arg in resolved_args:
+            if isinstance(arg, ChannelInterface):
+                channel = arg
+                if channel in input_channel_to_idx:
+                    # The same channel was added before, so reuse the index.
+                    channel_idx = input_channel_to_idx[channel]
+                else:
+                    # Add a new channel to the list of input channels.
+                    self.input_channels.append(channel)
+                    channel_idx = len(self.input_channels) - 1
+                    input_channel_to_idx[channel] = channel_idx
+
+                task_input = _ExecutableTaskInput(arg, channel_idx)
+            else:
+                task_input = _ExecutableTaskInput(arg, None)
+            self.task_inputs.append(task_input)
+
+        # Currently DAGs do not support binding kwargs to other DAG nodes.
+        for val in self.resolved_kwargs.values():
+            assert not isinstance(val, ChannelInterface)
+
+        # Input reader to read input data from upstream DAG nodes.
+        self.input_reader: ReaderInterface = SynchronousReader(self.input_channels)
+        # Output writer to write output data to downstream DAG nodes.
+        self.output_writer: WriterInterface = SynchronousWriter(
+            self.output_channels, self.output_idxs
+        )
+        # The intermediate future for a READ or COMPUTE operation,
+        # and `wait()` must be called to get the actual result of the operation.
+        # The result of a READ operation will be used by a COMPUTE operation,
+        # and the result of a COMPUTE operation will be used by a WRITE operation.
+        self._intermediate_future: Optional[DAGOperationFuture] = None
+
+    def cancel(self):
+        """
+        Close all the input channels and the output channel. The exact behavior
+        depends on the type of channel. Typically, it will release the resources
+        used by the channels.
+        """
+        self.input_reader.close()
+        self.output_writer.close()
+
+    def destroy_cuda_event(self):
+        """
+        If this executable task has created a GPU future that is not yet waited on,
+        that future is in the channel context cache. Remove the future from the cache
+        and destroy its CUDA event.
+        """
+        GPUFuture.remove_gpu_future(self.task_idx)
+
+    def prepare(self, overlap_gpu_communication: bool = False):
+        """
+        Prepare the task for execution. The `exec_operation` function can only
+        be called after `prepare` has been called.
+
+        Args:
+            overlap_gpu_communication: Whether to overlap GPU communication with
+                computation during DAG execution to improve performance
+        """
+        for typ_hint in self.input_type_hints:
+            typ_hint.register_custom_serializer()
+        self.output_type_hint.register_custom_serializer()
+        self.input_reader.start()
+        self.output_writer.start()
+
+        self._send_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
+        self._recv_stream: Union["cp.cuda.Stream", nullcontext] = nullcontext()
+        if not overlap_gpu_communication:
+            return
+
+        # Set up send_stream and recv_stream when overlap_gpu_communication
+        # is configured
+        if self.output_type_hint.requires_nccl():
+            nccl_group_id = _get_nccl_group_id(self.output_type_hint)
+            nccl_group = ChannelContext.get_current().communicators.get(nccl_group_id)
+            assert nccl_group is not None
+            self._send_stream = nccl_group.send_stream
+        if self.input_type_hints:
+            for type_hint in self.input_type_hints:
+                if type_hint.requires_nccl():
+                    nccl_group_id = _get_nccl_group_id(type_hint)
+                    nccl_group = ChannelContext.get_current().communicators.get(
+                        nccl_group_id
+                    )
+                    assert nccl_group is not None
+                    if not isinstance(self._recv_stream, nullcontext):
+                        assert self._recv_stream == nccl_group.recv_stream, (
+                            "Currently all torch tensor input channels of a "
+                            "Compiled Graph task should use the same recv cuda stream."
+                        )
+                    self._recv_stream = nccl_group.recv_stream
+
+    def wrap_and_set_intermediate_future(
+        self, val: Any, wrap_in_gpu_future: bool
+    ) -> None:
+        """
+        Wrap the value in a `DAGOperationFuture` and store to the intermediate future.
+        The value corresponds to result of a READ or COMPUTE operation.
+
+        If wrap_in_gpu_future is True, the value will be wrapped in a GPUFuture,
+        Otherwise, the future will be a ResolvedFuture.
+
+        Args:
+            val: The value to wrap in a future.
+            wrap_in_gpu_future: Whether to wrap the value in a GPUFuture.
+        """
+        assert self._intermediate_future is None
+
+        if wrap_in_gpu_future:
+            future = GPUFuture(val, self.task_idx)
+        else:
+            future = ResolvedFuture(val)
+        self._intermediate_future = future
+
+    def reset_and_wait_intermediate_future(self) -> Any:
+        """
+        Reset the intermediate future and wait for the result.
+
+        The wait does not block the CPU because:
+        - If the future is a ResolvedFuture, the result is immediately returned.
+        - If the future is a GPUFuture, the result is only waited by the current
+            CUDA stream, and the CPU is not blocked.
+
+        Returns:
+            The result of a READ or COMPUTE operation from the intermediate future.
+        """
+        future = self._intermediate_future
+        self._intermediate_future = None
+        return future.wait()
+
+    def _read(self, overlap_gpu_communication: bool) -> bool:
+        """
+        Read input data from upstream DAG nodes and cache the intermediate result.
+
+        Args:
+            overlap_gpu_communication: Whether to overlap GPU communication with
+                computation during DAG execution to improve performance.
+
+        Returns:
+            True if system error occurs and exit the loop; otherwise, False.
+        """
+        assert self._intermediate_future is None
+        exit = False
+        try:
+            input_data = self.input_reader.read()
+            # When overlap_gpu_communication is enabled, wrap the result in
+            # a GPUFuture so that this read operation (communication) can
+            # be overlapped with computation.
+            self.wrap_and_set_intermediate_future(
+                input_data, wrap_in_gpu_future=overlap_gpu_communication
+            )
+        except RayChannelError:
+            # Channel closed. Exit the loop.
+            exit = True
+        return exit
+
+    def _compute(
+        self,
+        overlap_gpu_communication: bool,
+        class_handle,
+    ) -> bool:
+        """
+        Retrieve the intermediate result from the READ operation and perform the
+        computation. Then, cache the new intermediate result. The caller must ensure
+        that the last operation executed is READ so that the function retrieves the
+        correct intermediate result.
+
+        Args:
+            overlap_gpu_communication: Whether to overlap GPU communication with
+                computation during DAG execution to improve performance.
+            class_handle: An instance of the class to which the actor belongs. For
+                example, the type of `class_handle` is <class 'xxxx.Worker'> if the
+                actor belongs to the `class Worker` class.
+        Returns:
+            True if system error occurs and exit the loop; otherwise, False.
+        """
+        input_data = self.reset_and_wait_intermediate_future()
+        try:
+            _process_return_vals(input_data, return_single_output=False)
+        except Exception as exc:
+            # Previous task raised an application-level exception.
+            # Propagate it and skip the actual task. We don't need to wrap the
+            # exception in a RayTaskError here because it has already been wrapped
+            # by the previous task.
+            self.wrap_and_set_intermediate_future(
+                exc, wrap_in_gpu_future=overlap_gpu_communication
+            )
+            return False
+
+        resolved_inputs = []
+        for task_input in self.task_inputs:
+            resolved_inputs.append(task_input.resolve(input_data))
+
+        if self.collective_op is not None:
+            # Run a NCCL collective operation.
+            method = self.collective_op.execute
+        else:
+            # Run an actor method.
+            method = getattr(class_handle, self.method_name)
+        try:
+            output_val = method(*resolved_inputs, **self.resolved_kwargs)
+        except Exception as exc:
+            output_val = _wrap_exception(exc)
+
+        # When overlap_gpu_communication is enabled, wrap the result in a GPUFuture
+        # so that this compute operation can be overlapped with communication.
+        self.wrap_and_set_intermediate_future(
+            output_val, wrap_in_gpu_future=overlap_gpu_communication
+        )
+        return False
+
+    def _write(self) -> bool:
+        """
+        Retrieve the intermediate result from the COMPUTE operation and write to its
+        downstream DAG nodes. The caller must ensure that the last operation executed
+        is COMPUTE so that the function retrieves the correct intermediate result.
+
+        Returns:
+            True if system error occurs and exit the loop; otherwise, False.
+        """
+        output_val = self.reset_and_wait_intermediate_future()
+        exit = False
+        try:
+            self.output_writer.write(output_val)
+        except RayChannelError:
+            # Channel closed. Exit the loop.
+            exit = True
+        return exit
+
+    def exec_operation(
+        self,
+        class_handle,
+        op_type: _DAGNodeOperationType,
+        overlap_gpu_communication: bool = False,
+    ) -> bool:
+        """
+        An ExecutableTask corresponds to a DAGNode. It consists of three
+        operations: READ, COMPUTE, and WRITE, which should be executed in
+        order to ensure that each operation can read the correct intermediate
+        result.
+        Args:
+            class_handle: The handle of the class to which the actor belongs.
+            op_type: The type of the operation. Possible types are READ,
+                COMPUTE, and WRITE.
+            overlap_gpu_communication: Whether to overlap GPU communication with
+                computation during DAG execution to improve performance.
+        Returns:
+            True if the next operation should not be executed; otherwise, False.
+        """
+        if op_type == _DAGNodeOperationType.READ:
+            with _device_context_manager():
+                with self._recv_stream:
+                    return self._read(overlap_gpu_communication)
+        elif op_type == _DAGNodeOperationType.COMPUTE:
+            return self._compute(overlap_gpu_communication, class_handle)
+        elif op_type == _DAGNodeOperationType.WRITE:
+            with _device_context_manager():
+                with self._send_stream:
+                    return self._write()
+
+
+@dataclass
+class _ExecutableTaskRecord:
+    actor_classname: str
+    actor_name: str
+    actor_id: str
+    method_name: str
+    bind_index: int
+    operation: str
+    start_t: float
+    end_t: float
+
+    def to_dict(self):
+        return asdict(self)
+
+
+@DeveloperAPI
+class CompiledDAG:
+    """Experimental class for accelerated execution.
+
+    This class should not be called directly. Instead, create
+    a ray.dag and call experimental_compile().
+
+    See REP https://github.com/ray-project/enhancements/pull/48 for more
+    information.
+    """
+
+    @ray.remote(num_cpus=0)
+    class DAGDriverProxyActor:
+        """
+        To support the driver as a reader, the output writer needs to be able to invoke
+        remote functions on the driver. This is necessary so that the output writer can
+        create a reader ref on the driver node, and later potentially create a larger
+        reader ref on the driver node if the channel backing store needs to be resized.
+        However, remote functions cannot be invoked on the driver.
+
+        A Compiled Graph creates an actor from this class when the DAG is initialized.
+        The actor is on the same node as the driver. This class has an empty
+        implementation, though it serves as a way for the output writer to invoke remote
+        functions on the driver node.
+        """
+
+        pass
+
+    def __init__(
+        self,
+        submit_timeout: Optional[float] = None,
+        buffer_size_bytes: Optional[int] = None,
+        enable_asyncio: bool = False,
+        max_inflight_executions: Optional[int] = None,
+        max_buffered_results: Optional[int] = None,
+        overlap_gpu_communication: Optional[bool] = None,
+        default_communicator: Optional[Union[Communicator, str]] = "create",
+    ):
+        """
+        Args:
+            submit_timeout: The maximum time in seconds to wait for execute() calls.
+                None means using default timeout (DAGContext.submit_timeout),
+                0 means immediate timeout (immediate success or timeout without
+                blocking), -1 means infinite timeout (block indefinitely).
+            buffer_size_bytes: The initial buffer size in bytes for messages
+                that can be passed between tasks in the DAG. The buffers will
+                be automatically resized if larger messages are written to the
+                channel.
+            enable_asyncio: Whether to enable asyncio. If enabled, caller must
+                be running in an event loop and must use `execute_async` to
+                invoke the DAG. Otherwise, the caller should use `execute` to
+                invoke the DAG.
+            max_inflight_executions: The maximum number of in-flight executions that
+                can be submitted via `execute` or `execute_async` before consuming
+                the output using `ray.get()`. If the caller submits more executions,
+                `RayCgraphCapacityExceeded` is raised.
+            max_buffered_results: The maximum number of results that can be
+                buffered at the driver. If more results are buffered,
+                `RayCgraphCapacityExceeded` is raised. Note that
+                when result corresponding to an execution is retrieved
+                (by calling `ray.get()` on a `CompiledDAGRef` or
+                `CompiledDAGRef` or await on a `CompiledDAGFuture), results
+                corresponding to earlier executions that have not been retrieved
+                yet are buffered.
+            overlap_gpu_communication: (experimental) Whether to overlap GPU
+                communication with computation during DAG execution. If True, the
+                communication and computation can be overlapped, which can improve
+                the performance of the DAG execution. If None, the default value
+                will be used.
+            _default_communicator: The default communicator to use to transfer
+                tensors. Three types of values are valid. (1) Communicator:
+                For p2p operations, this is the default communicator
+                to use for nodes annotated with `with_tensor_transport()` and when
+                shared memory is not the desired option (e.g., when transport="nccl",
+                or when transport="auto" for communication between two different GPUs).
+                For collective operations, this is the default communicator to use
+                when a custom communicator is not specified.
+                (2) "create": for each collective operation without a custom communicator
+                specified, a communicator is created and initialized on its involved actors,
+                or an already created communicator is reused if the set of actors is the same.
+                For all p2p operations without a custom communicator specified, it reuses
+                an already created collective communicator if the p2p actors are a subset.
+                Otherwise, a new communicator is created.
+                (3) None: a ValueError will be thrown if a custom communicator is not specified.
+
+        Returns:
+            Channel: A wrapper around ray.ObjectRef.
+        """
+        from ray.dag import DAGContext
+
+        ctx = DAGContext.get_current()
+
+        self._enable_asyncio: bool = enable_asyncio
+        self._fut_queue = asyncio.Queue()
+        self._max_inflight_executions = max_inflight_executions
+        if self._max_inflight_executions is None:
+            self._max_inflight_executions = ctx.max_inflight_executions
+        self._max_buffered_results = max_buffered_results
+        if self._max_buffered_results is None:
+            self._max_buffered_results = ctx.max_buffered_results
+        self._dag_id = uuid.uuid4().hex
+        self._submit_timeout: Optional[float] = submit_timeout
+        if self._submit_timeout is None:
+            self._submit_timeout = ctx.submit_timeout
+        self._get_timeout: Optional[float] = ctx.get_timeout
+        self._buffer_size_bytes: Optional[int] = buffer_size_bytes
+        if self._buffer_size_bytes is None:
+            self._buffer_size_bytes = ctx.buffer_size_bytes
+        self._overlap_gpu_communication: Optional[bool] = overlap_gpu_communication
+        if self._overlap_gpu_communication is None:
+            self._overlap_gpu_communication = ctx.overlap_gpu_communication
+        self._create_default_communicator = False
+        if isinstance(default_communicator, str):
+            if default_communicator == "create":
+                self._create_default_communicator = True
+                default_communicator = None
+            else:
+                raise ValueError(
+                    "The only allowed string for default_communicator is 'create', "
+                    f"got {default_communicator}"
+                )
+        elif default_communicator is not None and not isinstance(
+            default_communicator, Communicator
+        ):
+            raise ValueError(
+                "The default_communicator must be None, a string, or a Communicator, "
+                f"got {type(default_communicator)}"
+            )
+        self._default_communicator: Optional[Communicator] = default_communicator
+
+        # Dict from passed-in communicator to set of type hints that refer to it.
+        self._communicator_to_type_hints: Dict[
+            Communicator,
+            Set["ray.experimental.channel.torch_tensor_type.TorchTensorType"],
+        ] = defaultdict(set)
+        # Dict from set of actors to created communicator ID.
+        # These communicators are created by Compiled Graph, rather than passed in.
+        # Communicators are only created when self._create_default_communicator is True.
+        self._actors_to_created_communicator_id: Dict[
+            Tuple["ray.actor.ActorHandle"], str
+        ] = {}
+
+        # Set of actors involved in P2P communication using an unresolved communicator.
+        self._p2p_actors_with_unresolved_communicators: Set[
+            "ray.actor.ActorHandle"
+        ] = set()
+        # Set of DAG nodes involved in P2P communication using an unresolved communicator.
+        self._p2p_dag_nodes_with_unresolved_communicators: Set[
+            "ray.dag.DAGNode"
+        ] = set()
+        # Set of collective operations using an unresolved communicator.
+        self._collective_ops_with_unresolved_communicators: Set[
+            "ray.dag.collective_node._CollectiveOperation"
+        ] = set()
+
+        self._default_type_hint: ChannelOutputType = SharedMemoryType(
+            buffer_size_bytes=self._buffer_size_bytes,
+            # We conservatively set num_shm_buffers to _max_inflight_executions.
+            # It means that the DAG can be underutilized, but it guarantees there's
+            # no false positive timeouts.
+            num_shm_buffers=self._max_inflight_executions,
+        )
+        if not isinstance(self._buffer_size_bytes, int) or self._buffer_size_bytes <= 0:
+            raise ValueError(
+                "`buffer_size_bytes` must be a positive integer, found "
+                f"{self._buffer_size_bytes}"
+            )
+
+        # Used to ensure that the future returned to the
+        # caller corresponds to the correct DAG output. I.e.
+        # order of futures added to fut_queue should match the
+        # order of inputs written to the DAG.
+        self._dag_submission_lock = asyncio.Lock()
+
+        # idx -> CompiledTask.
+        self.idx_to_task: Dict[int, "CompiledTask"] = {}
+        # DAGNode -> idx.
+        self.dag_node_to_idx: Dict["ray.dag.DAGNode", int] = {}
+        # idx counter.
+        self.counter: int = 0
+
+        # Attributes that are set during preprocessing.
+        # Preprocessing identifies the input node and output node.
+        self.input_task_idx: Optional[int] = None
+        self.output_task_idx: Optional[int] = None
+        # List of task indices that are input attribute nodes.
+        self.input_attr_task_idxs: List[int] = []
+        # Denotes whether execute/execute_async returns a list of refs/futures.
+        self._returns_list: bool = False
+        # Number of expected positional args and kwargs that may be passed to
+        # dag.execute.
+        self._input_num_positional_args: Optional[int] = None
+        self._input_kwargs: Tuple[str, ...] = None
+
+        # Cached attributes that are set during compilation.
+        self.dag_input_channels: Optional[List[ChannelInterface]] = None
+        self.dag_output_channels: Optional[List[ChannelInterface]] = None
+        self._dag_submitter: Optional[WriterInterface] = None
+        self._dag_output_fetcher: Optional[ReaderInterface] = None
+
+        # ObjectRef for each worker's task. The task is an infinite loop that
+        # repeatedly executes the method specified in the DAG.
+        self.worker_task_refs: Dict["ray.actor.ActorHandle", "ray.ObjectRef"] = {}
+        self.actor_to_tasks: Dict[
+            "ray.actor.ActorHandle", List["CompiledTask"]
+        ] = defaultdict(list)
+        # Mapping from actor handle to its GPU IDs.
+        # This is used for type hint resolution for with_tensor_transport("auto").
+        self.actor_to_gpu_ids: Dict["ray.actor.ActorHandle", List[str]] = {}
+        self.actor_to_executable_tasks: Dict[
+            "ray.actor.ActorHandle", List["ExecutableTask"]
+        ] = {}
+        # Mapping from the actor handle to the execution schedule which is a list
+        # of operations to be executed.
+        self.actor_to_execution_schedule: Dict[
+            "ray.actor.ActorHandle", List[_DAGNodeOperation]
+        ] = defaultdict(list)
+        # Mapping from the actor handle to the node ID that the actor is on.
+        # A None actor handle means the actor is the driver.
+        self.actor_to_node_id: Dict[Optional["ray.actor.ActorHandle"], str] = {}
+        # The index of the current execution. It is incremented each time
+        # the DAG is executed.
+        self._execution_index: int = -1
+        # The maximum index of finished executions.
+        # All results with higher indexes have not been generated yet.
+        self._max_finished_execution_index: int = -1
+        # execution_index -> {channel_index -> result}
+        self._result_buffer: Dict[int, Dict[int, Any]] = defaultdict(dict)
+        # channel to possible inner channel
+        self._channel_dict: Dict[ChannelInterface, ChannelInterface] = {}
+
+        def _create_proxy_actor() -> "ray.actor.ActorHandle":
+            # Creates the driver actor on the same node as the driver.
+            #
+            # To support the driver as a reader, the output writer needs to be able to
+            # invoke remote functions on the driver (e.g., to create the reader ref, to
+            # create a reader ref for a larger object when the channel backing store is
+            # resized, etc.). The driver actor serves as a way for the output writer
+            # to invoke remote functions on the driver node.
+            return CompiledDAG.DAGDriverProxyActor.options(
+                scheduling_strategy=NodeAffinitySchedulingStrategy(
+                    ray.get_runtime_context().get_node_id(), soft=False
+                )
+            ).remote()
+
+        self._proxy_actor = _create_proxy_actor()
+        # Set to True when `teardown` API is called.
+        self._is_teardown = False
+        # Execution index to set of channel indices for CompiledDAGRefs
+        # or CompiledDAGFuture whose destructor has been called. A "None"
+        # channel index means there is only one channel, and its destructor
+        # has been called.
+        self._destructed_ref_idxs: Dict[int, Set[Optional[int]]] = dict()
+        # Execution index to set of channel indices for CompiledDAGRefs
+        # or CompiledDAGFuture whose get() has been called. A "None"
+        # channel index means there is only one channel, and its get()
+        # has been called.
+        self._got_ref_idxs: Dict[int, Set[Optional[int]]] = dict()
+
+    @property
+    def is_teardown(self) -> bool:
+        return self._is_teardown
+
+    def get_id(self) -> str:
+        """
+        Get the unique ID of the compiled DAG.
+        """
+        return self._dag_id
+
+    def __str__(self) -> str:
+        return f"CompiledDAG({self._dag_id})"
+
+    def _add_node(self, node: "ray.dag.DAGNode") -> None:
+        idx = self.counter
+        self.idx_to_task[idx] = CompiledTask(idx, node)
+        self.dag_node_to_idx[node] = idx
+        self.counter += 1
+
+    def _preprocess(self) -> None:
+        """Before compiling, preprocess the DAG to build an index from task to
+        upstream and downstream tasks, and to set the input and output node(s)
+        of the DAG.
+
+        This function is idempotent.
+        """
+        from ray.dag import (
+            DAGNode,
+            ClassMethodNode,
+            CollectiveOutputNode,
+            FunctionNode,
+            InputAttributeNode,
+            InputNode,
+            MultiOutputNode,
+        )
+
+        self.input_task_idx, self.output_task_idx = None, None
+
+        input_attributes: Set[str] = set()
+        # Find the input node and input attribute nodes in the DAG.
+        for idx, task in self.idx_to_task.items():
+            if isinstance(task.dag_node, InputNode):
+                assert self.input_task_idx is None, "More than one InputNode found"
+                self.input_task_idx = idx
+                # handle_unused_attributes:
+                # Save input attributes in a set.
+                input_node = task.dag_node
+                input_attributes.update(input_node.input_attribute_nodes.keys())
+            elif isinstance(task.dag_node, InputAttributeNode):
+                self.input_attr_task_idxs.append(idx)
+
+        # Find the (multi-)output node to the DAG.
+        for idx, task in self.idx_to_task.items():
+            if idx == self.input_task_idx or isinstance(
+                task.dag_node, InputAttributeNode
+            ):
+                continue
+            if (
+                len(task.downstream_task_idxs) == 0
+                and task.dag_node.is_cgraph_output_node
+            ):
+                assert self.output_task_idx is None, "More than one output node found"
+                self.output_task_idx = idx
+
+        assert self.output_task_idx is not None
+        output_node = self.idx_to_task[self.output_task_idx].dag_node
+        # Add an MultiOutputNode to the end of the DAG if it's not already there.
+        if not isinstance(output_node, MultiOutputNode):
+            output_node = MultiOutputNode([output_node])
+            self._add_node(output_node)
+            self.output_task_idx = self.dag_node_to_idx[output_node]
+        else:
+            self._returns_list = True
+
+        # TODO: Support no-input DAGs (use an empty object to signal).
+        if self.input_task_idx is None:
+            raise NotImplementedError(
+                "Compiled DAGs currently require exactly one InputNode"
+            )
+
+        # Whether the DAG binds directly to the InputNode(), versus binding to
+        # a positional arg or kwarg of the input. For example, a.foo.bind(inp)
+        # instead of a.foo.bind(inp[0]) or a.foo.bind(inp.key).
+        direct_input: Optional[bool] = None
+        # Collect the set of InputNode keys bound to DAG node args.
+        input_positional_args: Set[int] = set()
+        input_kwargs: Set[str] = set()
+        # Set of tasks with annotation of with_tensor_transport("auto").
+        # These only correspond to ClassMethodNodes, but not InputNodes
+        # or InputAttributeNodes.
+        auto_transport_tasks: Set["CompiledTask"] = set()
+
+        # For each task node, set its upstream and downstream task nodes.
+        # Also collect the set of tasks that produce torch.tensors.
+        for task_idx, task in self.idx_to_task.items():
+            dag_node = task.dag_node
+            if not (
+                isinstance(dag_node, InputNode)
+                or isinstance(dag_node, InputAttributeNode)
+                or isinstance(dag_node, MultiOutputNode)
+                or isinstance(dag_node, ClassMethodNode)
+            ):
+                if isinstance(dag_node, FunctionNode):
+                    # TODO(swang): Support non-actor tasks.
+                    raise NotImplementedError(
+                        "Compiled DAGs currently only support actor method nodes"
+                    )
+                else:
+                    raise ValueError(f"Found unsupported node of type {type(dag_node)}")
+
+            if isinstance(dag_node, ClassMethodNode) and dag_node.is_class_method_call:
+                actor_handle = dag_node._get_actor_handle()
+                if actor_handle is None:
+                    raise ValueError(
+                        "Compiled DAGs can only bind methods to an actor "
+                        "that is already created with Actor.remote()"
+                    )
+
+                if actor_handle not in self.actor_to_gpu_ids:
+                    self.actor_to_gpu_ids[actor_handle] = CompiledDAG._get_gpu_ids(
+                        actor_handle
+                    )
+
+                if isinstance(dag_node.type_hint, AutoTransportType):
+                    auto_transport_tasks.add(task)
+
+                # Collect actors for NCCL P2P methods.
+                if dag_node.type_hint.requires_nccl():
+                    self._track_communicator_usage(dag_node, {actor_handle})
+                # Collect NCCL collective operations.
+                if isinstance(dag_node, CollectiveOutputNode):
+                    self._track_communicator_usage(
+                        dag_node,
+                        set(dag_node._collective_op.actor_handles),
+                        collective_op=True,
+                    )
+                    assert not self._overlap_gpu_communication, (
+                        "Currently, the overlap_gpu_communication option is not "
+                        "supported for NCCL collective operations. Please set "
+                        "overlap_gpu_communication=False."
+                    )
+            elif isinstance(dag_node, InputNode) or isinstance(
+                dag_node, InputAttributeNode
+            ):
+                if dag_node.type_hint.requires_nccl():
+                    raise ValueError(
+                        "DAG inputs cannot be transferred via NCCL because "
+                        "the driver cannot participate in the NCCL group"
+                    )
+                if isinstance(dag_node.type_hint, AutoTransportType):
+                    # Currently driver on GPU is not supported, so we always
+                    # use shared memory to transfer tensors.
+                    dag_node.type_hint = TorchTensorType(
+                        device=dag_node.type_hint.device
+                    )
+
+            if type(dag_node.type_hint) is ChannelOutputType:
+                # No type hint specified by the user. Replace
+                # with the default type hint for this DAG.
+                dag_node.type_hint = self._default_type_hint
+
+            for _, val in task.kwargs.items():
+                if isinstance(val, DAGNode):
+                    raise ValueError(
+                        "Compiled DAG currently does not support binding to "
+                        "other DAG nodes as kwargs"
+                    )
+
+            for _, arg in enumerate(task.args):
+                if not isinstance(arg, DAGNode):
+                    continue
+                upstream_node_idx = self.dag_node_to_idx[arg]
+                upstream_task = self.idx_to_task[upstream_node_idx]
+                downstream_actor_handle = None
+                if (
+                    isinstance(dag_node, ClassMethodNode)
+                    and dag_node.is_class_method_call
+                ):
+                    downstream_actor_handle = dag_node._get_actor_handle()
+
+                # Add upstream node as the argument nodes of this task, whose
+                # type hints may be updated when resolved lazily.
+                task.arg_nodes.append(upstream_task.dag_node)
+
+                if isinstance(upstream_task.dag_node, InputAttributeNode):
+                    # Record all of the keys used to index the InputNode.
+                    # During execution, we will check that the user provides
+                    # the same args and kwargs.
+                    if isinstance(upstream_task.dag_node.key, int):
+                        input_positional_args.add(upstream_task.dag_node.key)
+                    elif isinstance(upstream_task.dag_node.key, str):
+                        input_kwargs.add(upstream_task.dag_node.key)
+                    else:
+                        raise ValueError(
+                            "InputNode() can only be indexed using int "
+                            "for positional args or str for kwargs."
+                        )
+
+                    if direct_input is not None and direct_input:
+                        raise ValueError(
+                            "All tasks must either use InputNode() "
+                            "directly, or they must index to specific args or "
+                            "kwargs."
+                        )
+                    direct_input = False
+
+                    # If the upstream node is an InputAttributeNode, treat the
+                    # DAG's input node as the actual upstream node
+                    upstream_task = self.idx_to_task[self.input_task_idx]
+
+                elif isinstance(upstream_task.dag_node, InputNode):
+                    if direct_input is not None and not direct_input:
+                        raise ValueError(
+                            "All tasks must either use InputNode() directly, "
+                            "or they must index to specific args or kwargs."
+                        )
+                    direct_input = True
+
+                upstream_task.downstream_task_idxs[task_idx] = downstream_actor_handle
+
+                if upstream_task.dag_node.type_hint.requires_nccl():
+                    # Here we are processing the args of the DAGNode, so track
+                    # downstream actors only, upstream actor is already tracked
+                    # when processing the DAGNode itself.
+                    self._track_communicator_usage(
+                        upstream_task.dag_node,
+                        {downstream_actor_handle},
+                    )
+        # Check that all specified input attributes, e.g., InputNode()["x"],
+        # are used in the DAG.
+        _check_unused_dag_input_attributes(output_node, input_attributes)
+
+        self._check_leaf_nodes()
+
+        self._resolve_auto_transport(auto_transport_tasks)
+
+        self._init_communicators()
+
+        if direct_input:
+            self._input_num_positional_args = 1
+        elif not input_positional_args:
+            self._input_num_positional_args = 0
+        else:
+            self._input_num_positional_args = max(input_positional_args) + 1
+        self._input_kwargs = tuple(input_kwargs)
+
+    def _init_communicators(self) -> None:
+        """
+        Initialize communicators for the DAG.
+        """
+
+        # First, initialize communicators that are passed in by the user.
+        for communicator, type_hints in self._communicator_to_type_hints.items():
+            communicator_id = _init_communicator(
+                communicator.get_actor_handles(),
+                communicator,
+                self._overlap_gpu_communication,
+            )
+            for type_hint in type_hints:
+                type_hint.set_communicator_id(communicator_id)
+
+        # Then, create communicators for collective operations.
+        # Reuse an already created communicator for the same set of actors.
+        for collective_op in self._collective_ops_with_unresolved_communicators:
+            if not self._create_default_communicator:
+                raise ValueError(
+                    "Communicator creation is not allowed for collective operations."
+                )
+            # using tuple to preserve the order of actors for collective operations
+            actors = tuple(collective_op.actor_handles)
+            if actors in self._actors_to_created_communicator_id:
+                communicator_id = self._actors_to_created_communicator_id[actors]
+            else:
+                communicator_id = _init_communicator(
+                    list(actors),
+                    None,
+                    self._overlap_gpu_communication,
+                )
+                self._actors_to_created_communicator_id[actors] = communicator_id
+            collective_op.type_hint.set_communicator_id(communicator_id)
+
+        # Finally, create a communicator for P2P operations.
+        # Reuse an already created collective op communicator when p2p actors
+        # are a subset of the actors in the collective op communicator.
+        p2p_communicator_id = None
+        if self._p2p_actors_with_unresolved_communicators:
+            for (
+                actors,
+                communicator_id,
+            ) in self._actors_to_created_communicator_id.items():
+                if self._p2p_actors_with_unresolved_communicators.issubset(actors):
+                    p2p_communicator_id = communicator_id
+                    break
+            if p2p_communicator_id is None:
+                p2p_communicator_id = _init_communicator(
+                    list(self._p2p_actors_with_unresolved_communicators),
+                    None,
+                    self._overlap_gpu_communication,
+                )
+            for dag_node in self._p2p_dag_nodes_with_unresolved_communicators:
+                dag_node.type_hint.set_communicator_id(p2p_communicator_id)
+
+    def _track_communicator_usage(
+        self,
+        dag_node: "ray.dag.DAGNode",
+        actors: Set["ray.actor.ActorHandle"],
+        collective_op: bool = False,
+    ) -> None:
+        """
+        Track the usage of a communicator.
+
+        This method first determines the communicator to use: if a custom
+        communicator is specified, use it; if not and a default communicator
+        is available, use it; otherwise, it records necessary information to
+        create a new communicator later.
+
+        This method also performs validation checks on the passed-in communicator.
+
+        Args:
+            dag_node: The DAG node that uses the communicator, this is the node
+                that has the `with_tensor_transport()` type hint for p2p communication,
+                or a `CollectiveOutputNode` for collective operations.
+            actors: The full or partial set of actors that use the communicator.
+                This method should be called one or multiple times so that all actors
+                of the communicator are tracked.
+            collective_op: Whether the communicator is used for a collective operation.
+        """
+        if None in actors:
+            raise ValueError("Driver cannot participate in the NCCL group.")
+        if collective_op:
+            type_hint = dag_node._collective_op.type_hint
+        else:
+            type_hint = dag_node.type_hint
+        communicator = type_hint.get_custom_communicator()
+
+        if communicator is None:
+            if (
+                self._default_communicator is None
+                and not self._create_default_communicator
+            ):
+                if dag_node._original_type_hint is not None:
+                    assert isinstance(dag_node._original_type_hint, AutoTransportType)
+                    raise ValueError(
+                        f"with_tensor_transport(transport='auto') is used for DAGNode {dag_node}, "
+                        "This requires specifying a default communicator or 'create' for "
+                        "_default_communicator when calling experimental_compile()."
+                    )
+                raise ValueError(
+                    f"DAGNode {dag_node} has no custom communicator specified. "
+                    "Please specify a custom communicator for the DAGNode using "
+                    "`with_tensor_transport()`, or specify a communicator or 'create' for "
+                    "_default_communicator when calling experimental_compile()."
+                )
+            communicator = self._default_communicator
+
+        if communicator is None:
+            if collective_op:
+                self._collective_ops_with_unresolved_communicators.add(
+                    dag_node._collective_op
+                )
+            else:
+                self._p2p_dag_nodes_with_unresolved_communicators.add(dag_node)
+                self._p2p_actors_with_unresolved_communicators.update(actors)
+        else:
+            if collective_op:
+                if set(communicator.get_actor_handles()) != actors:
+                    raise ValueError(
+                        "The passed-in communicator must have the same set "
+                        "of actors as the collective operation. "
+                        f"The passed-in communicator has actors {communicator.get_actor_handles()} "
+                        f"while the collective operation has actors {actors}."
+                    )
+            else:
+                if not actors.issubset(set(communicator.get_actor_handles())):
+                    raise ValueError(
+                        "The passed-in communicator must include all of the actors "
+                        "used in the P2P operation. "
+                        f"The passed-in communicator has actors {communicator.get_actor_handles()} "
+                        f"while the P2P operation has actors {actors}."
+                    )
+            self._communicator_to_type_hints[communicator].add(type_hint)
+
+    def _resolve_auto_transport(
+        self,
+        auto_transport_tasks: Set["CompiledTask"],
+    ) -> None:
+        """
+        Resolve the auto transport type hint for the DAG.
+        """
+        type_hint_resolver = TypeHintResolver(self.actor_to_gpu_ids)
+        # Resolve AutoChannelType type hints and track the actors that use NCCL.
+        # This is needed so that the NCCL group can be initialized for these
+        # actors that use NCCL.
+        for task in auto_transport_tasks:
+            writer = task.dag_node._get_actor_handle()
+            readers = task.downstream_task_idxs.values()
+            writer_and_node = (writer, self._get_node_id(writer))
+            reader_and_node_list = [
+                (reader, self._get_node_id(reader)) for reader in readers
+            ]
+            # Update the type hint to the resolved one. This is needed because
+            # the resolved type hint's `register_custom_serializer` will be called
+            # in preparation for channel I/O.
+            task.dag_node.type_hint = type_hint_resolver.resolve(
+                task.dag_node.type_hint,
+                writer_and_node,
+                reader_and_node_list,
+            )
+            if task.dag_node.type_hint.requires_nccl():
+                self._track_communicator_usage(
+                    task.dag_node,
+                    set(readers).union({writer}),
+                )
+
+    def _check_leaf_nodes(self) -> None:
+        """
+        Check if there are leaf nodes in the DAG and raise an error if there are.
+        """
+        from ray.dag import (
+            DAGNode,
+            ClassMethodNode,
+        )
+
+        leaf_nodes: List[DAGNode] = []
+        for _, task in self.idx_to_task.items():
+            if not isinstance(task.dag_node, ClassMethodNode):
+                continue
+            if (
+                len(task.downstream_task_idxs) == 0
+                and not task.dag_node.is_cgraph_output_node
+            ):
+                leaf_nodes.append(task.dag_node)
+        # Leaf nodes are not allowed because the exception thrown by the leaf
+        # node will not be propagated to the driver.
+        if len(leaf_nodes) != 0:
+            raise ValueError(
+                "Compiled DAG doesn't support leaf nodes, i.e., nodes that don't have "
+                "downstream nodes and are not output nodes. There are "
+                f"{len(leaf_nodes)} leaf nodes in the DAG. Please add the outputs of "
+                f"{[leaf_node.get_method_name() for leaf_node in leaf_nodes]} to the "
+                f"the MultiOutputNode."
+            )
+
+    @staticmethod
+    def _get_gpu_ids(actor_handle: "ray.actor.ActorHandle") -> List[str]:
+        """
+        Get the GPU IDs of an actor handle.
+        """
+        accelerator_ids = ray.get(
+            actor_handle.__ray_call__.remote(
+                lambda self: ray.get_runtime_context().get_accelerator_ids()
+            )
+        )
+        return accelerator_ids.get("GPU", [])
+
+    def _get_node_id(self, actor_handle: Optional["ray.actor.ActorHandle"]) -> str:
+        """
+        Get the node ID of an actor handle and cache it.
+
+        Args:
+            actor_handle: The actor handle, or None if the actor handle is the
+                driver.
+        Returns:
+            The node ID of the actor handle or driver.
+        """
+        if actor_handle in self.actor_to_node_id:
+            return self.actor_to_node_id[actor_handle]
+        node_id = None
+        if actor_handle == self._proxy_actor or actor_handle is None:
+            node_id = ray.get_runtime_context().get_node_id()
+        else:
+            node_id = ray.get(
+                actor_handle.__ray_call__.remote(
+                    lambda self: ray.get_runtime_context().get_node_id()
+                )
+            )
+        self.actor_to_node_id[actor_handle] = node_id
+        return node_id
+
+    def _get_or_compile(
+        self,
+    ) -> None:
+        """Compile an execution path. This allocates channels for adjacent
+        tasks to send/receive values. An infinite task is submitted to each
+        actor in the DAG that repeatedly receives from input channel(s) and
+        sends to output channel(s).
+
+        This function is idempotent and will cache the previously allocated
+        channels. After calling this function, _dag_submitter and
+        _dag_output_fetcher will be set and can be used to invoke and fetch
+        outputs for the DAG.
+        """
+        from ray.dag import (
+            DAGNode,
+            InputNode,
+            InputAttributeNode,
+            MultiOutputNode,
+            ClassMethodNode,
+        )
+
+        if self.input_task_idx is None:
+            self._preprocess()
+        assert self.input_task_idx is not None
+
+        if self._dag_submitter is not None:
+            assert self._dag_output_fetcher is not None
+            return
+
+        frontier = [self.input_task_idx]
+        visited = set()
+        # Create output buffers. This loop does a breadth-first search through the DAG.
+        while frontier:
+            cur_idx = frontier.pop(0)
+            if cur_idx in visited:
+                continue
+            visited.add(cur_idx)
+
+            task = self.idx_to_task[cur_idx]
+            if (
+                isinstance(task.dag_node, ClassMethodNode)
+                and task.dag_node.is_class_method_call
+            ):
+                # Create output buffers for the actor method.
+                assert len(task.output_channels) == 0
+                # `output_to_readers` stores the reader tasks for each output of
+                # the current node. If the current node returns one output, the
+                # readers are the downstream nodes of the current node. If the
+                # current node returns multiple outputs, the readers of each
+                # output are the downstream nodes of the ClassMethodNode that
+                # is a class method output.
+                output_to_readers: Dict[CompiledTask, List[CompiledTask]] = defaultdict(
+                    list
+                )
+                for idx in task.downstream_task_idxs:
+                    downstream_task = self.idx_to_task[idx]
+                    downstream_node = downstream_task.dag_node
+                    if (
+                        isinstance(downstream_node, ClassMethodNode)
+                        and downstream_node.is_class_method_output
+                    ):
+                        output_to_readers[downstream_task] = [
+                            self.idx_to_task[idx]
+                            for idx in downstream_task.downstream_task_idxs
+                        ]
+                    else:
+                        if task not in output_to_readers:
+                            output_to_readers[task] = []
+                        output_to_readers[task].append(downstream_task)
+                fn = task.dag_node._get_remote_method("__ray_call__")
+                for output, readers in output_to_readers.items():
+                    reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]] = []
+                    # Use reader_handles_set to deduplicate readers on the
+                    # same actor, because with CachedChannel each actor will
+                    # only read from the upstream channel once.
+                    reader_handles_set = set()
+                    read_by_multi_output_node = False
+                    for reader in readers:
+                        if isinstance(reader.dag_node, MultiOutputNode):
+                            read_by_multi_output_node = True
+                            # inserting at 0 to make sure driver is first reader as
+                            # expected by CompositeChannel read
+                            reader_and_node_list.insert(
+                                0,
+                                (
+                                    self._proxy_actor,
+                                    self._get_node_id(self._proxy_actor),
+                                ),
+                            )
+                        else:
+                            reader_handle = reader.dag_node._get_actor_handle()
+                            if reader_handle not in reader_handles_set:
+                                reader_handle = reader.dag_node._get_actor_handle()
+                                reader_and_node_list.append(
+                                    (reader_handle, self._get_node_id(reader_handle))
+                                )
+                                reader_handles_set.add(reader_handle)
+
+                    # if driver is an actual actor, gets driver actor id
+                    driver_actor_id = (
+                        ray.get_runtime_context().get_actor_id()
+                        if read_by_multi_output_node
+                        else None
+                    )
+                    # Create an output channel for each output of the current node.
+                    output_channel = ray.get(
+                        fn.remote(
+                            do_allocate_channel,
+                            reader_and_node_list,
+                            task.dag_node.type_hint,
+                            driver_actor_id,
+                        )
+                    )
+                    output_idx = None
+                    downstream_node = output.dag_node
+                    if (
+                        isinstance(downstream_node, ClassMethodNode)
+                        and downstream_node.is_class_method_output
+                    ):
+                        output_idx = downstream_node.output_idx
+                    task.output_channels.append(output_channel)
+                    task.output_idxs.append(output_idx)
+                    task.output_node_idxs.append(self.dag_node_to_idx[downstream_node])
+                actor_handle = task.dag_node._get_actor_handle()
+                assert actor_handle is not None
+                self.actor_to_tasks[actor_handle].append(task)
+            elif (
+                isinstance(task.dag_node, ClassMethodNode)
+                and task.dag_node.is_class_method_output
+            ):
+                task_node = task.dag_node
+                upstream_node = task_node.class_method_call
+                assert upstream_node
+                upstream_task = self.idx_to_task[self.dag_node_to_idx[upstream_node]]
+                for i in range(len(upstream_task.output_channels)):
+                    if upstream_task.output_idxs[i] == task_node.output_idx:
+                        task.output_channels.append(upstream_task.output_channels[i])
+                        task.output_idxs.append(upstream_task.output_idxs[i])
+                assert len(task.output_channels) == 1
+            elif isinstance(task.dag_node, InputNode):
+                # A dictionary that maps an InputNode or InputAttributeNode to its
+                # readers and the node on which the reader is running. Use `set` to
+                # deduplicate readers on the same actor because with CachedChannel
+                # each actor will only read from the shared memory once.
+                input_node_to_reader_and_node_set: Dict[
+                    Union[InputNode, InputAttributeNode],
+                    Set[Tuple["ray.actor.ActorHandle", str]],
+                ] = defaultdict(set)
+
+                for idx in task.downstream_task_idxs:
+                    reader_task = self.idx_to_task[idx]
+                    assert isinstance(reader_task.dag_node, ClassMethodNode)
+                    reader_handle = reader_task.dag_node._get_actor_handle()
+                    reader_node_id = self._get_node_id(reader_handle)
+                    for arg in reader_task.args:
+                        if isinstance(arg, InputAttributeNode) or isinstance(
+                            arg, InputNode
+                        ):
+                            input_node_to_reader_and_node_set[arg].add(
+                                (reader_handle, reader_node_id)
+                            )
+
+                # A single channel is responsible for sending the same data to
+                # corresponding consumers. Therefore, we create a channel for
+                # each InputAttributeNode, or a single channel for the entire
+                # input data if there are no InputAttributeNodes.
+                task.output_channels = []
+                for input_dag_node in input_node_to_reader_and_node_set:
+                    reader_and_node_list = list(
+                        input_node_to_reader_and_node_set[input_dag_node]
+                    )
+
+                    output_channel = do_allocate_channel(
+                        self,
+                        reader_and_node_list,
+                        input_dag_node.type_hint,
+                        None,
+                    )
+                    task.output_channels.append(output_channel)
+                    task.output_idxs.append(
+                        None
+                        if isinstance(input_dag_node, InputNode)
+                        else input_dag_node.key
+                    )
+
+                    # Update the InputAttributeNode's `output_channels`, which is
+                    # used to determine whether to create a CachedChannel.
+                    if isinstance(input_dag_node, InputAttributeNode):
+                        input_attr_idx = self.dag_node_to_idx[input_dag_node]
+                        input_attr_task = self.idx_to_task[input_attr_idx]
+                        input_attr_task.output_channels.append(output_channel)
+                        assert len(input_attr_task.output_channels) == 1
+            else:
+                assert isinstance(task.dag_node, InputAttributeNode) or isinstance(
+                    task.dag_node, MultiOutputNode
+                )
+
+            for idx in task.downstream_task_idxs:
+                frontier.append(idx)
+
+        # Validate input channels for tasks that have not been visited
+        for node_idx, task in self.idx_to_task.items():
+            if (
+                node_idx == self.input_task_idx
+                or node_idx == self.output_task_idx
+                or isinstance(task.dag_node, InputAttributeNode)
+            ):
+                continue
+            if node_idx not in visited:
+                has_at_least_one_channel_input = False
+                for arg in task.args:
+                    if isinstance(arg, DAGNode):
+                        has_at_least_one_channel_input = True
+                if not has_at_least_one_channel_input:
+                    raise ValueError(
+                        "Compiled DAGs require each task to take a ray.dag.InputNode "
+                        "or at least one other DAGNode as an input. "
+                        "Invalid task node:\n"
+                        f"{task.dag_node}\n"
+                        "Please bind the task to proper DAG nodes."
+                    )
+
+        from ray.dag.constants import RAY_CGRAPH_ENABLE_DETECT_DEADLOCK
+
+        if RAY_CGRAPH_ENABLE_DETECT_DEADLOCK and self._detect_deadlock():
+            raise ValueError(
+                "This DAG cannot be compiled because it will deadlock on NCCL "
+                "calls. If you believe this is a false positive, please disable "
+                "the graph verification by setting the environment variable "
+                "RAY_CGRAPH_ENABLE_DETECT_DEADLOCK to 0 and file an issue at "
+                "https://github.com/ray-project/ray/issues/new/."
+            )
+
+        input_task = self.idx_to_task[self.input_task_idx]
+        self.dag_input_channels = input_task.output_channels
+        assert self.dag_input_channels is not None
+
+        # Create executable tasks for each actor
+        for actor_handle, tasks in self.actor_to_tasks.items():
+            # Dict from arg to the set of tasks that consume it.
+            arg_to_consumers: Dict[DAGNode, Set[CompiledTask]] = defaultdict(set)
+
+            # Step 1: populate `arg_to_consumers` and perform some validation.
+            for task in tasks:
+                has_at_least_one_channel_input = False
+                for arg in task.args:
+                    if isinstance(arg, DAGNode):
+                        has_at_least_one_channel_input = True
+                        arg_to_consumers[arg].add(task)
+                        arg_idx = self.dag_node_to_idx[arg]
+                        upstream_task = self.idx_to_task[arg_idx]
+                        assert len(upstream_task.output_channels) == 1
+                        arg_channel = upstream_task.output_channels[0]
+                        assert arg_channel is not None
+                # TODO: Support no-input DAGs (use an empty object to signal).
+                if not has_at_least_one_channel_input:
+                    raise ValueError(
+                        "Compiled DAGs require each task to take a "
+                        "ray.dag.InputNode or at least one other DAGNode as an "
+                        "input"
+                    )
+
+            # Step 2: create cached channels if needed
+
+            # Dict from original channel to the channel to be used in execution.
+            # The value of this dict is either the original channel or a newly
+            # created CachedChannel (if the original channel is read more than once).
+            for arg, consumers in arg_to_consumers.items():
+                arg_idx = self.dag_node_to_idx[arg]
+                upstream_task = self.idx_to_task[arg_idx]
+                assert len(upstream_task.output_channels) == 1
+                arg_channel = upstream_task.output_channels[0]
+                assert arg_channel is not None
+                if len(consumers) > 1:
+                    self._channel_dict[arg_channel] = CachedChannel(
+                        len(consumers),
+                        arg_channel,
+                    )
+                else:
+                    self._channel_dict[arg_channel] = arg_channel
+
+            # Step 3: create executable tasks for the actor
+            executable_tasks = []
+            for task in tasks:
+                resolved_args: List[Any] = []
+                for arg in task.args:
+                    if isinstance(arg, DAGNode):
+                        arg_idx = self.dag_node_to_idx[arg]
+                        upstream_task = self.idx_to_task[arg_idx]
+                        assert len(upstream_task.output_channels) == 1
+                        arg_channel = upstream_task.output_channels[0]
+                        assert arg_channel is not None
+                        arg_channel = self._channel_dict[arg_channel]
+                        resolved_args.append(arg_channel)
+                    else:
+                        # Constant arg
+                        resolved_args.append(arg)
+                executable_task = ExecutableTask(
+                    task,
+                    resolved_args,
+                    task.kwargs,
+                )
+                executable_tasks.append(executable_task)
+            # Sort executable tasks based on their bind index, i.e., submission order
+            # so that they will be executed in that order.
+            executable_tasks.sort(key=lambda task: task.bind_index)
+            self.actor_to_executable_tasks[actor_handle] = executable_tasks
+
+        from ray.dag.constants import RAY_CGRAPH_ENABLE_PROFILING
+
+        if RAY_CGRAPH_ENABLE_PROFILING:
+            exec_task_func = do_profile_tasks
+        else:
+            exec_task_func = do_exec_tasks
+
+        # Build an execution schedule for each actor
+        self.actor_to_execution_schedule = self._build_execution_schedule()
+        for actor_handle, executable_tasks in self.actor_to_executable_tasks.items():
+            self.worker_task_refs[actor_handle] = actor_handle.__ray_call__.options(
+                concurrency_group="_ray_system"
+            ).remote(
+                exec_task_func,
+                executable_tasks,
+                self.actor_to_execution_schedule[actor_handle],
+                self._overlap_gpu_communication,
+            )
+
+        assert self.output_task_idx is not None
+        self.dag_output_channels = []
+        for output in self.idx_to_task[self.output_task_idx].args:
+            assert isinstance(output, DAGNode)
+            output_idx = self.dag_node_to_idx[output]
+            task = self.idx_to_task[output_idx]
+            assert len(task.output_channels) == 1
+            self.dag_output_channels.append(task.output_channels[0])
+
+        # Register custom serializers for input, input attribute, and output nodes.
+        self._register_input_output_custom_serializer()
+
+        assert self.dag_input_channels
+        assert self.dag_output_channels
+        assert [
+            output_channel is not None for output_channel in self.dag_output_channels
+        ]
+        # If no MultiOutputNode was specified during the DAG creation, there is only
+        # one output. Return a single output channel instead of a list of
+        # channels.
+        if not self._returns_list:
+            assert len(self.dag_output_channels) == 1
+
+        # Driver should ray.put on input, ray.get/release on output
+        self._monitor = self._monitor_failures()
+        input_task = self.idx_to_task[self.input_task_idx]
+        if self._enable_asyncio:
+            self._dag_submitter = AwaitableBackgroundWriter(
+                self.dag_input_channels,
+                input_task.output_idxs,
+                is_input=True,
+            )
+            self._dag_output_fetcher = AwaitableBackgroundReader(
+                self.dag_output_channels,
+                self._fut_queue,
+            )
+        else:
+            self._dag_submitter = SynchronousWriter(
+                self.dag_input_channels, input_task.output_idxs, is_input=True
+            )
+            self._dag_output_fetcher = SynchronousReader(self.dag_output_channels)
+
+        self._dag_submitter.start()
+        self._dag_output_fetcher.start()
+
+    def _generate_dag_operation_graph_node(
+        self,
+    ) -> Dict["ray.actor.ActorHandle", List[List[_DAGOperationGraphNode]]]:
+        """
+        Generate READ, COMPUTE, and WRITE operations for each DAG node.
+
+        Returns:
+            A dictionary that maps an actor handle to a list of lists of
+            _DAGOperationGraphNode. For the same actor, the index of the
+            outer list corresponds to the index of the ExecutableTask in
+            the list of `executable_tasks` in `actor_to_executable_tasks`,
+            i.e. `exec_task_idx`. In the inner list, the order of operations
+            is READ, COMPUTE, and WRITE.
+
+            Example:
+            {
+                actor1: [
+                    [READ COMPUTE WRITE] # exec_task_idx 0
+                    [READ COMPUTE WRITE] # exec_task_idx 1
+                ]
+            }
+        """
+        from ray.dag.collective_node import CollectiveOutputNode, _CollectiveOperation
+
+        assert self.idx_to_task
+        assert self.actor_to_executable_tasks
+
+        actor_to_operation_nodes: Dict[
+            "ray.actor.ActorHandle", List[List[_DAGOperationGraphNode]]
+        ] = defaultdict(list)
+        collective_op_to_nodes: Dict[
+            _CollectiveOperation, Set[_DAGOperationGraphNode]
+        ] = defaultdict(set)
+        collective_op_to_idxs: Dict[
+            _CollectiveOperation, Tuple[int, _DAGNodeOperationType]
+        ] = defaultdict(set)
+
+        for actor_handle, executable_tasks in self.actor_to_executable_tasks.items():
+            for exec_task_idx, exec_task in enumerate(executable_tasks):
+                # Divide a DAG node into three _DAGOperationGraphNodes: READ, COMPUTE,
+                # and WRITE. Each _DAGOperationGraphNode has a _DAGNodeOperation.
+                task_idx = exec_task.task_idx
+                dag_node = self.idx_to_task[task_idx].dag_node
+                method_name = exec_task.method_name
+                actor_handle = dag_node._get_actor_handle()
+                requires_nccl = dag_node.type_hint.requires_nccl()
+                upstream_requires_nccl = False
+                for upstream_node in dag_node._upstream_nodes:
+                    if upstream_node.type_hint.requires_nccl():
+                        upstream_requires_nccl = True
+                        break
+
+                read_node = _DAGOperationGraphNode(
+                    _DAGNodeOperation(
+                        exec_task_idx, _DAGNodeOperationType.READ, method_name
+                    ),
+                    task_idx,
+                    actor_handle,
+                    upstream_requires_nccl,
+                )
+                compute_node = _DAGOperationGraphNode(
+                    _DAGNodeOperation(
+                        exec_task_idx, _DAGNodeOperationType.COMPUTE, method_name
+                    ),
+                    task_idx,
+                    actor_handle,
+                    isinstance(dag_node, CollectiveOutputNode),
+                )
+                write_node = _DAGOperationGraphNode(
+                    _DAGNodeOperation(
+                        exec_task_idx, _DAGNodeOperationType.WRITE, method_name
+                    ),
+                    task_idx,
+                    actor_handle,
+                    requires_nccl,
+                )
+
+                actor_to_operation_nodes[actor_handle].append(
+                    [read_node, compute_node, write_node]
+                )
+                if isinstance(dag_node, CollectiveOutputNode):
+                    collective_op_to_nodes[dag_node.collective_op].add(compute_node)
+                    collective_op_to_idxs[dag_node.collective_op].add(
+                        (task_idx, _DAGNodeOperationType.COMPUTE)
+                    )
+
+        # Set collective nodes for all the NCCL collective operation nodes.
+        for collective_op, nodes in collective_op_to_nodes.items():
+            idxs = collective_op_to_idxs[collective_op]
+            for node in nodes:
+                node.collective_idxs = idxs
+
+        return actor_to_operation_nodes
+
+    def _build_execution_schedule(
+        self,
+    ) -> Dict["ray.actor.ActorHandle", List[_DAGNodeOperation]]:
+        """
+        Generate an execution schedule for each actor. The schedule is a list of
+        _DAGNodeOperation.
+
+        Step 1: Generate a DAG node operation graph. Refer to the functions
+        `_generate_dag_operation_graph_node` and `_build_dag_node_operation_graph`
+        for more details.
+
+        Step 2: Topological sort
+
+        It is possible to have multiple _DAGOperationGraphNodes with zero in-degree.
+        Refer to the function `_select_next_nodes` for the logic of selecting nodes.
+
+        Then, put the selected nodes into the corresponding actors' schedules.
+
+        The schedule should be intuitive to users, meaning that the execution should
+        perform operations in ascending order of `bind_index` as much as possible.
+
+        [Example]:
+
+        See `test_execution_schedule` for more examples.
+
+        Returns:
+            actor_to_execution_schedule: A dictionary that maps an actor handle to
+                the execution schedule which is a list of operations to be executed.
+        """
+        # Step 1: Build a graph of _DAGOperationGraphNode
+        actor_to_operation_nodes = self._generate_dag_operation_graph_node()
+        graph = _build_dag_node_operation_graph(
+            self.idx_to_task, actor_to_operation_nodes
+        )
+        # Step 2: Generate an execution schedule for each actor using topological sort
+        actor_to_execution_schedule = _generate_actor_to_execution_schedule(graph)
+
+        # Step 3: Overlap GPU communication for the execution schedule if configured
+        actor_to_overlapped_schedule = None
+        if self._overlap_gpu_communication:
+            actor_to_overlapped_schedule = _generate_overlapped_execution_schedule(
+                actor_to_execution_schedule
+            )
+
+        if RAY_CGRAPH_VISUALIZE_SCHEDULE:
+            _visualize_execution_schedule(
+                actor_to_execution_schedule, actor_to_overlapped_schedule, graph
+            )
+
+        if actor_to_overlapped_schedule is not None:
+            return _extract_execution_schedule(actor_to_overlapped_schedule)
+        else:
+            return _extract_execution_schedule(actor_to_execution_schedule)
+
+    def _detect_deadlock(self) -> bool:
+        """
+        TODO (kevin85421): Avoid false negatives.
+
+        Currently, a compiled graph may deadlock if there are NCCL channels, and the
+        readers have control dependencies on the same actor. For example:
+
+        actor1.a ---> actor2.f1
+                 |
+                 ---> actor2.f2
+
+        The control dependency between `actor2.f1` and `actor2.f2` is that `f1` should
+        run before `f2`. If `actor1.a` writes to `actor2.f2` before `actor2.f1`, a
+        deadlock will occur.
+
+        Currently, the execution schedule is not granular enough to detect this
+        deadlock.
+
+        Returns:
+            True if a deadlock is detected; otherwise, False.
+        """
+        logger.debug("Deadlock detection has not been implemented yet.")
+        return False
+
+    def _monitor_failures(self):
+        get_outer = weakref.ref(self)
+
+        class Monitor(threading.Thread):
+            def __init__(self):
+                super().__init__(daemon=True)
+                self.name = "CompiledGraphMonitorThread"
+                # Lock to make sure that we only perform teardown for this DAG
+                # once.
+                self._in_teardown_lock = threading.Lock()
+                self._teardown_done = False
+
+            def _outer_ref_alive(self) -> bool:
+                if get_outer() is None:
+                    logger.error(
+                        "CompiledDAG has been destructed before teardown. "
+                        "This should not occur please report an issue at "
+                        "https://github.com/ray-project/ray/issues/new/.",
+                        stack_info=True,
+                    )
+                    return False
+                return True
+
+            def wait_teardown(self, kill_actors: bool = False):
+                outer = get_outer()
+                if not self._outer_ref_alive():
+                    return
+
+                from ray.dag import DAGContext
+
+                ctx = DAGContext.get_current()
+                teardown_timeout = ctx.teardown_timeout
+                for actor, ref in outer.worker_task_refs.items():
+                    timeout = False
+                    try:
+                        ray.get(ref, timeout=teardown_timeout)
+                    except ray.exceptions.GetTimeoutError:
+                        msg = (
+                            f"Compiled DAG actor {actor} is still running "
+                            f"{teardown_timeout}s after teardown()."
+                        )
+                        if kill_actors:
+                            msg += (
+                                " Force-killing actor. "
+                                "Increase RAY_CGRAPH_teardown_timeout if you want "
+                                "teardown to wait longer."
+                            )
+                            ray.kill(actor)
+                        else:
+                            msg += (
+                                " Teardown may hang. "
+                                "Call teardown with kill_actors=True if force kill "
+                                "is desired."
+                            )
+
+                        logger.warning(msg)
+                        timeout = True
+                    except Exception:
+                        # We just want to check that the task has finished so
+                        # we don't care if the actor task ended in an
+                        # exception.
+                        pass
+
+                    if not timeout:
+                        continue
+
+                    try:
+                        ray.get(ref)
+                    except Exception:
+                        pass
+
+                if kill_actors:
+                    # In the previous loop, we allow the actor tasks to exit first.
+                    # Now, we force kill the actors if not yet.
+                    for actor in outer.worker_task_refs:
+                        logger.info(f"Killing actor: {actor}")
+                        ray.kill(actor)
+
+            def teardown(self, kill_actors: bool = False):
+                with self._in_teardown_lock:
+                    if self._teardown_done:
+                        return
+
+                    outer = get_outer()
+                    if not self._outer_ref_alive():
+                        return
+
+                    logger.info("Tearing down compiled DAG")
+                    outer._dag_submitter.close()
+                    outer._dag_output_fetcher.close()
+
+                    for actor in outer.actor_to_executable_tasks.keys():
+                        logger.info(f"Cancelling compiled worker on actor: {actor}")
+                    # Cancel all actor loops in parallel.
+                    cancel_refs = [
+                        actor.__ray_call__.remote(do_cancel_executable_tasks, tasks)
+                        for actor, tasks in outer.actor_to_executable_tasks.items()
+                    ]
+                    for cancel_ref in cancel_refs:
+                        try:
+                            ray.get(cancel_ref, timeout=30)
+                        except RayChannelError:
+                            # Channel error happens when a channel is closed
+                            # or timed out. In this case, do not log.
+                            pass
+                        except Exception:
+                            logger.exception("Error cancelling worker task")
+                            pass
+
+                    for (
+                        communicator_id
+                    ) in outer._actors_to_created_communicator_id.values():
+                        _destroy_communicator(communicator_id)
+
+                    logger.info("Waiting for worker tasks to exit")
+                    self.wait_teardown(kill_actors=kill_actors)
+
+                    logger.info("Teardown complete")
+                    self._teardown_done = True
+
+            def run(self):
+                try:
+                    outer = get_outer()
+                    if not self._outer_ref_alive():
+                        return
+                    ray.get(list(outer.worker_task_refs.values()))
+                except KeyboardInterrupt:
+                    logger.info(
+                        "Received KeyboardInterrupt, tearing down with kill_actors=True"
+                    )
+                    self.teardown(kill_actors=True)
+                except Exception as e:
+                    logger.debug(f"Handling exception from worker tasks: {e}")
+                    self.teardown()
+
+        monitor = Monitor()
+        monitor.start()
+        return monitor
+
+    def _raise_if_too_many_inflight_executions(self):
+        num_inflight_executions = (
+            self._execution_index - self._max_finished_execution_index
+        )
+        if num_inflight_executions >= self._max_inflight_executions:
+            raise ray.exceptions.RayCgraphCapacityExceeded(
+                "The compiled graph can't have more than "
+                f"{self._max_inflight_executions} in-flight executions, and you "
+                f"currently have {num_inflight_executions} in-flight executions. "
+                "Retrieve an output using ray.get before submitting more requests or "
+                "increase `_max_inflight_executions`. "
+                "`dag.experimental_compile(_max_inflight_executions=...)`"
+            )
+
+    def _has_execution_results(
+        self,
+        execution_index: int,
+    ) -> bool:
+        """Check whether there are results corresponding to the given execution
+        index stored in self._result_buffer. This helps avoid fetching and
+        caching results again.
+
+        Args:
+            execution_index: The execution index corresponding to the result.
+
+        Returns:
+            Whether the result for the given index has been fetched and cached.
+        """
+        return execution_index in self._result_buffer
+
+    def _cache_execution_results(
+        self,
+        execution_index: int,
+        result: Any,
+    ):
+        """Cache execution results in self._result_buffer. Results are converted
+        to dictionary format to allow efficient element removal and calculation of
+        the buffer size. This can only be called once per execution index.
+
+        Args:
+            execution_index: The execution index corresponding to the result.
+            result: The results from all channels to be cached.
+        """
+        if not self._has_execution_results(execution_index):
+            for chan_idx, res in enumerate(result):
+                # avoid caching for any CompiledDAGRef that has already been destructed.
+                if not (
+                    execution_index in self._destructed_ref_idxs
+                    and chan_idx in self._destructed_ref_idxs[execution_index]
+                ):
+                    self._result_buffer[execution_index][chan_idx] = res
+
+    def _get_execution_results(
+        self, execution_index: int, channel_index: Optional[int]
+    ) -> List[Any]:
+        """Retrieve execution results from self._result_buffer and return the result.
+        Results are converted back to original list format ordered by output channel
+        index.
+
+        Args:
+            execution_index: The execution index to retrieve results from.
+            channel_index: The index of the output channel corresponding to the result.
+                Channel indexing is consistent with the order of
+                self.dag_output_channels. None means that the result wraps outputs from
+                all output channels.
+
+        Returns:
+            The execution result corresponding to the given execution index and channel
+            index.
+        """
+        # Although CompiledDAGRef and CompiledDAGFuture guarantee that the same
+        # execution index and channel index combination will not be requested multiple
+        # times and therefore self._result_buffer will always have execution_index as
+        # a key, we still do a sanity check to avoid misuses.
+        assert execution_index in self._result_buffer
+
+        if channel_index is None:
+            # Convert results stored in self._result_buffer back to original
+            # list representation
+            result = [
+                kv[1]
+                for kv in sorted(
+                    self._result_buffer.pop(execution_index).items(),
+                    key=lambda kv: kv[0],
+                )
+            ]
+        else:
+            result = [self._result_buffer[execution_index].pop(channel_index)]
+
+        if execution_index not in self._got_ref_idxs:
+            self._got_ref_idxs[execution_index] = set()
+        self._got_ref_idxs[execution_index].add(channel_index)
+        self._clean_up_buffers(execution_index)
+        return result
+
+    def _delete_execution_results(self, execution_index: int, channel_index: int):
+        """
+        Delete the execution results for the given execution index and channel index.
+        This method should be called when a CompiledDAGRef or CompiledDAGFuture is
+        destructed.
+
+        Note that this method maintains metadata for the deleted execution results,
+        and only actually deletes the buffers lazily when the buffer is not needed
+        anymore.
+
+        Args:
+            execution_index: The execution index to destruct results from.
+            channel_index: The index of the output channel corresponding to the result.
+        """
+        if execution_index not in self._destructed_ref_idxs:
+            self._destructed_ref_idxs[execution_index] = set()
+        self._destructed_ref_idxs[execution_index].add(channel_index)
+        self._clean_up_buffers(execution_index)
+
+    def _try_release_result_buffer(self, execution_index: int):
+        """
+        Try to release the result buffer for the given execution index.
+        """
+
+        should_release = False
+        got_channel_idxs = self._got_ref_idxs.get(execution_index, set())
+        if None in got_channel_idxs:
+            assert len(got_channel_idxs) == 1, (
+                "when None exists in got_channel_idxs, it means all channels, and "
+                "it should be the only value in the set",
+            )
+            should_release = True
+        else:
+            destructed_channel_idxs = self._destructed_ref_idxs.get(
+                execution_index, set()
+            )
+            processed_channel_idxs = got_channel_idxs.union(destructed_channel_idxs)
+            # No more processing is needed for this execution index.
+            should_release = processed_channel_idxs == set(
+                range(len(self.dag_output_channels))
+            )
+
+        if not should_release:
+            return False
+
+        self._result_buffer.pop(execution_index, None)
+        self._destructed_ref_idxs.pop(execution_index, None)
+        self._got_ref_idxs.pop(execution_index, None)
+        return True
+
+    def _try_release_native_buffer(
+        self, idx_to_release: int, timeout: Optional[float] = None
+    ) -> bool:
+        """
+        Try to release the native buffer for the given execution index.
+
+        Args:
+            idx_to_release: The execution index to release buffers from.
+            timeout: The maximum time in seconds to wait for the release.
+
+        Returns:
+            Whether the buffers have been released.
+        """
+        if idx_to_release != self._max_finished_execution_index + 1:
+            # Native buffer can only be released for the next execution index.
+            return False
+
+        destructed_channel_idxs = self._destructed_ref_idxs.get(idx_to_release, set())
+        should_release = False
+        if None in destructed_channel_idxs:
+            assert len(destructed_channel_idxs) == 1, (
+                "when None exists in destructed_channel_idxs, it means all channels, "
+                "and it should be the only value in the set",
+            )
+            should_release = True
+        elif len(destructed_channel_idxs) == len(self.dag_output_channels):
+            should_release = True
+
+        if not should_release:
+            return False
+
+        # refs corresponding to idx_to_release are all destructed,
+        # and they are never fetched or cached.
+        assert idx_to_release not in self._result_buffer
+        assert idx_to_release not in self._got_ref_idxs
+
+        try:
+            self._dag_output_fetcher.release_channel_buffers(timeout)
+        except RayChannelTimeoutError as e:
+            raise RayChannelTimeoutError(
+                "Releasing native buffers corresponding to a stale CompiledDAGRef "
+                "is taking a long time. If this is expected, increase "
+                f"RAY_CGRAPH_get_timeout which is currently {self._get_timeout} "
+                "seconds. Otherwise, this may indicate that the execution "
+                "is hanging."
+            ) from e
+        self._destructed_ref_idxs.pop(idx_to_release)
+
+        return True
+
+    def _try_release_buffer(
+        self, idx_to_release: int, timeout: Optional[float] = None
+    ) -> bool:
+        """
+        Try to release the buffer for the given execution index.
+        First try to release the native buffer, then try to release the result buffer.
+
+        Args:
+            idx_to_release: The execution index to release buffers from.
+            timeout: The maximum time in seconds to wait for the release.
+
+        Returns:
+            Whether the native buffer or result buffer has been released.
+        """
+        if self._try_release_native_buffer(idx_to_release, timeout):
+            # Releasing native buffer means the corresponding execution result
+            # is consumed (and discarded).
+            self._max_finished_execution_index += 1
+            return True
+        return self._try_release_result_buffer(idx_to_release)
+
+    def _try_release_buffers(self):
+        """
+        Repeatedly release buffer if possible.
+
+        This method starts from _max_finished_execution_index + 1 and tries to release
+        as many buffers as possible. If a native buffer is released,
+        _max_finished_execution_index will be incremented.
+        """
+        timeout = self._get_timeout
+        while True:
+            start_time = time.monotonic()
+            if not self._try_release_buffer(
+                self._max_finished_execution_index + 1, timeout
+            ):
+                break
+
+            if timeout != -1:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+
+    def _clean_up_buffers(self, idx_to_release: int):
+        """
+        Clean up native and result buffers.
+
+        This method:
+        1. Tries to release the buffer for the given execution index.
+           This index is the specific one that requires a clean up,
+           e.g., right after get() is called or a CompiledDAGRef/CompiledDAGFuture
+           is destructed.
+        2. Tries to release all buffers starting from _max_finished_execution_index + 1.
+           This step is to clean up buffers that are no longer needed.
+
+        Args:
+            idx_to_release: The execution index that requires a clean up,
+                e.g., right after get() is called or a CompiledDAGRef/CompiledDAGFuture
+                is destructed.
+        """
+        self._try_release_buffer(idx_to_release)
+        self._try_release_buffers()
+
+    def _execute_until(
+        self,
+        execution_index: int,
+        channel_index: Optional[int] = None,
+        timeout: Optional[float] = None,
+    ):
+        """Repeatedly execute this DAG until the given execution index and
+        buffer results for all CompiledDagRef's.
+        If the DAG has already been executed up to the given index, it will do nothing.
+
+        Note: If this comes across execution indices for which the corresponding
+        CompiledDAGRef's have been destructed, it will release the buffer and not
+        cache the result.
+
+        Args:
+            execution_index: The execution index to execute until.
+            channel_index: The index of the output channel to get the result from.
+                Channel indexing is consistent with the order of
+                self.dag_output_channels. None means wrapping results from all output
+                channels into a single list.
+            timeout: The maximum time in seconds to wait for the execution.
+                None means using default timeout (DAGContext.get_timeout),
+                0 means immediate timeout (immediate success or timeout without
+                blocking), -1 means infinite timeout (block indefinitely).
+
+        TODO(rui): catch the case that user holds onto the CompiledDAGRefs
+        """
+        if timeout is None:
+            timeout = self._get_timeout
+        while self._max_finished_execution_index < execution_index:
+            if len(self._result_buffer) >= self._max_buffered_results:
+                raise RayCgraphCapacityExceeded(
+                    "The compiled graph can't have more than "
+                    f"{self._max_buffered_results} buffered results, and you "
+                    f"currently have {len(self._result_buffer)} buffered results. "
+                    "Call `ray.get()` on CompiledDAGRef's (or await on "
+                    "CompiledDAGFuture's) to retrieve results, or increase "
+                    f"`_max_buffered_results` if buffering is desired, note that "
+                    "this will increase driver memory usage."
+                )
+            start_time = time.monotonic()
+
+            # Fetch results from each output channel up to execution_index and cache
+            # them separately to enable individual retrieval
+            # If a CompiledDagRef for a specific execution index has been destructed,
+            # release the channel buffers for that execution index instead of caching
+            try:
+                if not self._try_release_native_buffer(
+                    self._max_finished_execution_index + 1, timeout
+                ):
+                    result = self._dag_output_fetcher.read(timeout)
+                    self._cache_execution_results(
+                        self._max_finished_execution_index + 1,
+                        result,
+                    )
+                # We have either released the native buffer or fetched and
+                # cached the result buffer, therefore we always increment
+                # _max_finished_execution_index.
+                self._max_finished_execution_index += 1
+            except RayChannelTimeoutError as e:
+                raise RayChannelTimeoutError(
+                    "If the execution is expected to take a long time, increase "
+                    f"RAY_CGRAPH_get_timeout which is currently {self._get_timeout} "
+                    "seconds. Otherwise, this may indicate that the execution is "
+                    "hanging."
+                ) from e
+
+            if timeout != -1:
+                timeout -= time.monotonic() - start_time
+                timeout = max(timeout, 0)
+
+    def execute(
+        self,
+        *args,
+        **kwargs,
+    ) -> Union[CompiledDAGRef, List[CompiledDAGRef]]:
+        """Execute this DAG using the compiled execution path.
+
+        Args:
+            args: Args to the InputNode.
+            kwargs: Kwargs to the InputNode
+
+        Returns:
+            A list of Channels that can be used to read the DAG result.
+
+        Raises:
+            RayChannelTimeoutError: If the execution does not complete within
+                self._submit_timeout seconds.
+
+        NOTE: Not thread-safe due to _execution_index etc.
+        """
+        if self._enable_asyncio:
+            raise ValueError("Use execute_async if enable_asyncio=True")
+
+        self._get_or_compile()
+
+        self._check_inputs(args, kwargs)
+        if len(args) == 1 and len(kwargs) == 0:
+            # When serializing a tuple, the Ray serializer invokes pickle5, which adds
+            # several microseconds of overhead. One common case for Compiled Graphs is
+            # passing a single argument (oftentimes of of type `bytes`, which requires
+            # no serialization). To avoid imposing this overhead on this common case, we
+            # create a fast path for this case that avoids pickle5.
+            inp = args[0]
+        else:
+            inp = CompiledDAGArgs(args=args, kwargs=kwargs)
+
+        # We want to release any buffers we can at this point based on the
+        # max_finished_execution_index so that the number of inflight executions
+        # is up to date.
+        self._try_release_buffers()
+        self._raise_if_too_many_inflight_executions()
+        try:
+            self._dag_submitter.write(inp, self._submit_timeout)
+        except RayChannelTimeoutError as e:
+            raise RayChannelTimeoutError(
+                "If the execution is expected to take a long time, increase "
+                f"RAY_CGRAPH_submit_timeout which is currently {self._submit_timeout} "
+                "seconds. Otherwise, this may indicate that execution is hanging."
+            ) from e
+
+        self._execution_index += 1
+
+        if self._returns_list:
+            ref = [
+                CompiledDAGRef(self, self._execution_index, channel_index)
+                for channel_index in range(len(self.dag_output_channels))
+            ]
+        else:
+            ref = CompiledDAGRef(self, self._execution_index)
+
+        return ref
+
+    def _check_inputs(self, args: Tuple[Any, ...], kwargs: Dict[str, Any]) -> None:
+        """
+        Helper method to check that the DAG args provided by the user during
+        execution are valid according to the defined DAG.
+        """
+        if len(args) != self._input_num_positional_args:
+            raise ValueError(
+                "dag.execute() or dag.execute_async() must be "
+                f"called with {self._input_num_positional_args} positional args, got "
+                f"{len(args)}"
+            )
+
+        for kwarg in self._input_kwargs:
+            if kwarg not in kwargs:
+                raise ValueError(
+                    "dag.execute() or dag.execute_async() "
+                    f"must be called with kwarg `{kwarg}`"
+                )
+
+    async def execute_async(
+        self,
+        *args,
+        **kwargs,
+    ) -> Union[CompiledDAGFuture, List[CompiledDAGFuture]]:
+        """Execute this DAG using the compiled execution path.
+
+        NOTE: Not thread-safe.
+
+        Args:
+            args: Args to the InputNode.
+            kwargs: Kwargs to the InputNode.
+
+        Returns:
+            A list of Channels that can be used to read the DAG result.
+        """
+        if not self._enable_asyncio:
+            raise ValueError("Use execute if enable_asyncio=False")
+
+        self._get_or_compile()
+        self._check_inputs(args, kwargs)
+        async with self._dag_submission_lock:
+            if len(args) == 1 and len(kwargs) == 0:
+                # When serializing a tuple, the Ray serializer invokes pickle5, which
+                # adds several microseconds of overhead. One common case for accelerated
+                # DAGs is passing a single argument (oftentimes of of type `bytes`,
+                # which requires no serialization). To avoid imposing this overhead on
+                # this common case, we create a fast path for this case that avoids
+                # pickle5.
+                inp = args[0]
+            else:
+                inp = CompiledDAGArgs(args=args, kwargs=kwargs)
+
+            self._raise_if_too_many_inflight_executions()
+            await self._dag_submitter.write(inp)
+            # Allocate a future that the caller can use to get the result.
+            fut = asyncio.Future()
+            await self._fut_queue.put(fut)
+
+        self._execution_index += 1
+
+        if self._returns_list:
+            fut = [
+                CompiledDAGFuture(self, self._execution_index, fut, channel_index)
+                for channel_index in range(len(self.dag_output_channels))
+            ]
+        else:
+            fut = CompiledDAGFuture(self, self._execution_index, fut)
+
+        return fut
+
+    def _visualize_ascii(self) -> str:
+        """
+        Visualize the compiled graph in
+        ASCII format with directional markers.
+
+        This function generates an ASCII visualization of a Compiled Graph,
+        where each task node is labeled,
+        and edges use `<` and `>` markers to show data flow direction.
+
+        This method is called by:
+            - `compiled_dag.visualize(format="ascii")`
+
+
+
+        High-Level Algorithm:
+        - Topological Sorting: Sort nodes topologically to organize
+            them into layers based on dependencies.
+        - Grid Initialization: Set up a 2D grid canvas with dimensions based
+            on the number of layers and the maximum number of nodes per layer.
+        - Node Placement: Position each node on the grid according to its
+            layer and relative position within that layer.
+            Spacing is added for readability, and directional markers (`<` and `>`)
+            are added to edges to show input/output flow clearly.
+
+        This method should be called
+          **after** compiling the graph with `experimental_compile()`.
+
+        Returns:
+            ASCII representation of the CG with Nodes Information,
+            Edges Information and Graph Built.
+
+        Limitations:
+        - Note: This is only used for quick visualization for small graphs.
+            For complex graph (i.e. more than 20 tasks), please use graphviz.
+        - Scale: Works best for smaller CGs (typically fewer than 20 tasks).
+            Larger CGs may result in dense, less readable ASCII
+            outputs due to limited space for node and edge rendering.
+        - Shape: Ideal for relatively shallow CGs with clear dependency paths.
+            For deep, highly branched or densely connected CGs,
+            readability may suffer.
+        - Edge Overlap: In cases with high fan-out (i.e., nodes with many children)
+            or fan-in (nodes with many parents), edge lines may intersect or overlap
+            in the ASCII visualization, potentially obscuring some connections.
+        - Multi-output Tasks: Multi-output tasks can be visualized, but positioning
+            may cause line breaks or overlap when a task has multiple outputs that
+            feed into nodes at varying depths.
+
+        Example:
+            Basic Visualization:
+            ```python
+            # Print the CG structure in ASCII format
+            print(compiled_dag.visualize(format="ascii"))
+            ```
+
+            Example of Ordered Visualization (task is build in order
+                to reduce line intersection):
+            ```python
+            with InputNode() as i:
+                o1, o2, o3 = a.return_three.bind(i)
+                o4 = b.echo.bind(o1)
+                o5 = b.echo.bind(o2)
+                o6, o7 = b.return_two.bind(o3)
+                dag = MultiOutputNode([o4, o5, o6, o7])
+
+            compiled_dag = dag.experimental_compile()
+            compiled_dag.visualize(format="ascii",view=True)
+
+
+            # Output:
+            # 0:InputNode
+            # |
+            # 1:Actor_54777d:return_three
+            # |---------------------------->|---------------------------->|                                                  # noqa
+            # 2:Output[0]                   3:Output[1]                   4:Output[2]                                        # noqa
+            # |                             |                             |                                                  # noqa
+            # 5:Actor_c927c9:echo           6:Actor_c927c9:echo           7:Actor_c927c9:return_two                          # noqa
+            # |                             |                             |---------------------------->|                    # noqa
+            # |                             |                             9:Output[0]                   10:Output[1]         # noqa
+            # |<----------------------------|-----------------------------|-----------------------------|                    # noqa
+            # 8:MultiOutputNode
+            ```
+
+            Example of Anti-pattern Visualization (There are intersections):
+            # We can swtich the nodes ordering to reduce intersections, i.e. swap o2 and o3
+            ```python
+            with InputNode() as i:
+                o1, o2, o3 = a.return_three.bind(i)
+                o4 = b.echo.bind(o1)
+                o5 = b.echo.bind(o3)
+                o6, o7 = b.return_two.bind(o2)
+                dag = MultiOutputNode([o4, o5, o6, o7])
+            compiled_dag = dag.experimental_compile()
+            compiled_dag.visualize(format="ascii",view=True)
+
+            # Output (Nodes 5, 7, 9, 10 should connect to Node 8):
+            # 0:InputNode
+            # |
+            # 1:Actor_84835a:return_three
+            # |---------------------------->|---------------------------->|                            # noqa
+            # 2:Output[0]                   3:Output[1]                   4:Output[2]                  # noqa
+            # |                             |                             |                            # noqa
+            # 5:Actor_02a6a1:echo           6:Actor_02a6a1:return_two     7:Actor_02a6a1:echo          # noqa
+            # |                             |---------------------------->|                            # noqa
+            # |                             9:Output[0]                   10:Output[1]                 # noqa
+            # |<----------------------------------------------------------|                            # noqa
+            # 8:MultiOutputNode
+            ```
+        """
+
+        from ray.dag import (
+            InputAttributeNode,
+            InputNode,
+            MultiOutputNode,
+            ClassMethodNode,
+            DAGNode,
+        )
+
+        # Check that the DAG has been compiled
+        if not hasattr(self, "idx_to_task") or not self.idx_to_task:
+            raise ValueError(
+                "The DAG must be compiled before calling 'visualize()'. "
+                "Please call 'experimental_compile()' first."
+            )
+
+        # Check that each CompiledTask has a valid dag_node
+        for idx, task in self.idx_to_task.items():
+            if not hasattr(task, "dag_node") or not isinstance(task.dag_node, DAGNode):
+                raise ValueError(
+                    f"Task at index {idx} does not have a valid 'dag_node'. "
+                    "Ensure that 'experimental_compile()' completed successfully."
+                )
+
+        from collections import defaultdict, deque
+
+        # Create adjacency list representation of the DAG
+        # Adjacency list for DAG; maps a node index to its downstream nodes.
+        adj_list: Dict[int, List[int]] = defaultdict(list)
+        # Indegree count for topological sorting; maps a node index to its indegree.
+        indegree: Dict[int, int] = defaultdict(int)
+
+        # Tracks whether a node is a multi-output node.
+        is_multi_output: Dict[int, bool] = defaultdict(bool)
+        # Maps child node indices to their parent node indices.
+        child2parent: Dict[int, int] = defaultdict(int)
+        ascii_visualization = ""
+        # Node information; maps a node index to its descriptive label.
+        node_info: Dict[int, str] = {}
+        # Edge information; tuples of (upstream_index, downstream_index, edge_label).
+        edge_info: List[Tuple[int, int, str]] = []
+
+        for idx, task in self.idx_to_task.items():
+            dag_node = task.dag_node
+            label = f"Task {idx}  "
+
+            # Determine the type and label of the node
+            if isinstance(dag_node, InputNode):
+                label += "InputNode"
+            elif isinstance(dag_node, InputAttributeNode):
+                label += f"InputAttributeNode[{dag_node.key}]"
+            elif isinstance(dag_node, MultiOutputNode):
+                label += "MultiOutputNode"
+            elif isinstance(dag_node, ClassMethodNode):
+                if dag_node.is_class_method_call:
+                    method_name = dag_node.get_method_name()
+                    actor_handle = dag_node._get_actor_handle()
+                    actor_id = (
+                        actor_handle._actor_id.hex()[:6] if actor_handle else "unknown"
+                    )
+                    label += f"Actor: {actor_id}... Method: {method_name}"
+                elif dag_node.is_class_method_output:
+                    label += f"ClassMethodOutputNode[{dag_node.output_idx}]"
+                else:
+                    label += "ClassMethodNode"
+            else:
+                label += type(dag_node).__name__
+
+            node_info[idx] = label
+
+            for arg_index, arg in enumerate(dag_node.get_args()):
+                if isinstance(arg, DAGNode):
+                    upstream_task_idx = self.dag_node_to_idx[arg]
+
+                    # Get the type hint for this argument
+                    if arg_index < len(task.arg_type_hints):
+                        if task.arg_type_hints[arg_index].requires_nccl():
+                            type_hint = "Nccl"
+                        else:
+                            type_hint = type(task.arg_type_hints[arg_index]).__name__
+                    else:
+                        type_hint = "UnknownType"
+
+                    adj_list[upstream_task_idx].append(idx)
+                    indegree[idx] += 1
+                    edge_info.append((upstream_task_idx, idx, type_hint))
+
+        width_adjust = 0
+        for upstream_task_idx, child_idx_list in adj_list.items():
+            # Mark as multi-output if the node has more than one output path
+            if len(child_idx_list) > 1:
+                for child in child_idx_list:
+                    is_multi_output[child] = True
+                    child2parent[child] = upstream_task_idx
+                width_adjust = max(width_adjust, len(child_idx_list))
+
+        # Topological sort to determine layers
+        layers = defaultdict(list)
+        zero_indegree = deque([idx for idx in self.idx_to_task if indegree[idx] == 0])
+        layer_index = 0
+
+        while zero_indegree:
+            next_layer = deque()
+            while zero_indegree:
+                task_idx = zero_indegree.popleft()
+                layers[layer_index].append(task_idx)
+                for downstream in adj_list[task_idx]:
+                    indegree[downstream] -= 1
+                    if indegree[downstream] == 0:
+                        next_layer.append(downstream)
+            zero_indegree = next_layer
+            layer_index += 1
+
+        # Print detailed node information
+        ascii_visualization += "Nodes Information:\n"
+        for idx, info in node_info.items():
+            ascii_visualization += f'{idx} [label="{info}"] \n'
+
+        # Print edges
+        ascii_visualization += "\nEdges Information:\n"
+        for upstream_task, downstream_task, type_hint in edge_info:
+            if type_hint == "Nccl":
+                edgs_channel = "+++"
+            else:
+                edgs_channel = "---"
+            ascii_visualization += (
+                f"{upstream_task} {edgs_channel}>" f" {downstream_task}\n"
+            )
+
+        # Add the legend to the output
+        ascii_visualization += "\nLegend:\n"
+        ascii_visualization += "+++> : Represents Nccl-type data channels\n"
+        ascii_visualization += "---> : Represents Shared Memory data channels\n"
+
+        # Find the maximum width (number of nodes in any layer)
+        max_width = max(len(layer) for layer in layers.values()) + width_adjust
+        height = len(layers)
+
+        # Build grid for ASCII visualization
+        grid = [[" " for _ in range(max_width * 20)] for _ in range(height * 2 - 1)]
+
+        # Place nodes in the grid with more details
+        task_to_pos = {}
+        for layer_num, layer_tasks in layers.items():
+            layer_y = layer_num * 2  # Every second row is for nodes
+            for col_num, task_idx in enumerate(layer_tasks):
+                task = self.idx_to_task[task_idx]
+                task_info = f"{task_idx}:"
+
+                # Determine if it's an actor method or a regular task
+                if isinstance(task.dag_node, ClassMethodNode):
+                    if task.dag_node.is_class_method_call:
+                        method_name = task.dag_node.get_method_name()
+                        actor_handle = task.dag_node._get_actor_handle()
+                        actor_id = (
+                            actor_handle._actor_id.hex()[:6]
+                            if actor_handle
+                            else "unknown"
+                        )
+                        task_info += f"Actor_{actor_id}:{method_name}"
+                    elif task.dag_node.is_class_method_output:
+                        task_info += f"Output[{task.dag_node.output_idx}]"
+                    else:
+                        task_info += "UnknownMethod"
+                else:
+                    task_info += type(task.dag_node).__name__
+
+                adjust_col_num = 0
+                if task_idx in is_multi_output:
+                    adjust_col_num = layers[layer_num - 1].index(child2parent[task_idx])
+                col_x = (col_num + adjust_col_num) * 30  # Every 30th column for spacing
+                # Place the task information into the grid
+                for i, char in enumerate(task_info):
+                    if col_x + i < len(grid[0]):  # Ensure we don't overflow the grid
+                        grid[layer_y][col_x + i] = char
+
+                task_to_pos[task_idx] = (layer_y, col_x)
+
+        # Connect the nodes with lines
+        for upstream_task, downstream_tasks in adj_list.items():
+            upstream_y, upstream_x = task_to_pos[upstream_task]
+            for downstream_task in downstream_tasks:
+                downstream_y, downstream_x = task_to_pos[downstream_task]
+
+                # Draw vertical line
+                for y in range(upstream_y + 1, downstream_y):
+                    if grid[y][upstream_x] == " ":
+                        grid[y][upstream_x] = "|"
+
+                    # Draw horizontal line with directional arrows
+                if upstream_x != downstream_x:
+                    for x in range(
+                        min(upstream_x, downstream_x) + 1,
+                        max(upstream_x, downstream_x),
+                    ):
+                        grid[downstream_y - 1][x] = (
+                            "-"
+                            if grid[downstream_y - 1][x] == " "
+                            else grid[downstream_y - 1][x]
+                        )
+
+                    # Add arrows to indicate flow direction
+                    if downstream_x > upstream_x:
+                        grid[downstream_y - 1][downstream_x - 1] = ">"
+                    else:
+                        grid[downstream_y - 1][downstream_x + 1] = "<"
+
+                # Draw connection to the next task
+                grid[downstream_y - 1][downstream_x] = "|"
+
+        # Ensure proper multi-output task connection
+        for idx, task in self.idx_to_task.items():
+            if isinstance(task.dag_node, MultiOutputNode):
+                output_tasks = task.dag_node.get_args()
+                for i, output_task in enumerate(output_tasks):
+                    if isinstance(output_task, DAGNode):
+                        output_task_idx = self.dag_node_to_idx[output_task]
+                        if output_task_idx in task_to_pos:
+                            output_y, output_x = task_to_pos[output_task_idx]
+                            grid[output_y - 1][output_x] = "|"
+
+        # Convert grid to string for printing
+        ascii_visualization += "\nGraph Built:\n"
+        ascii_visualization += "\n".join("".join(row) for row in grid)
+
+        return ascii_visualization
+
+    def get_channel_details(
+        self, channel: ChannelInterface, downstream_actor_id: str
+    ) -> str:
+        """
+        Get details about outer and inner channel types and channel ids
+        based on the channel and the downstream actor ID.
+        Used for graph visualization.
+        Args:
+            channel: The channel to get details for.
+            downstream_actor_id: The downstream actor ID.
+        Returns:
+            A string with details about the channel based on its connection
+            to the actor provided.
+        """
+        channel_details = type(channel).__name__
+        # get outer channel
+        if channel in self._channel_dict and self._channel_dict[channel] != channel:
+            channel = self._channel_dict[channel]
+            channel_details += f"\n{type(channel).__name__}"
+            if type(channel) is CachedChannel:
+                channel_details += f", {channel._channel_id[:6]}..."
+        # get inner channel
+        if (
+            type(channel) is CompositeChannel
+            and downstream_actor_id in channel._channel_dict
+        ):
+            inner_channel = channel._channel_dict[downstream_actor_id]
+            channel_details += f"\n{type(inner_channel).__name__}"
+            if type(inner_channel) is IntraProcessChannel:
+                channel_details += f", {inner_channel._channel_id[:6]}..."
+        return channel_details
+
+    def visualize(
+        self,
+        filename="compiled_graph",
+        format="png",
+        view=False,
+        channel_details=False,
+    ) -> str:
+        """
+        Visualize the compiled graph by showing tasks and their dependencies.
+        This method should be called **after** the graph has been compiled using
+        `experimental_compile()`.
+
+        Args:
+            filename: For non-ASCII formats, the output file name (without extension).
+                For ASCII format, the visualization will be printed to the console,
+                and this argument is ignored.
+            format: The format of the output file (e.g., 'png', 'pdf', 'ascii').
+            view: For non-ASCII formats: Whether to open the file with the default
+                viewer. For ASCII format: Whether to print the visualization and return
+                None or return the ascii visualization string directly.
+            channel_details: If True, adds channel details to edges.
+
+        Returns:
+            The string representation of the compiled graph. For Graphviz-based formats
+            (e.g., 'png', 'pdf', 'jpeg'), returns the Graphviz DOT string representation
+            of the compiled graph. For ASCII format, returns the ASCII string
+            representation of the compiled graph.
+
+        Raises:
+            ValueError: If the graph is empty or not properly compiled.
+            ImportError: If the `graphviz` package is not installed.
+
+        """
+        if format == "ascii":
+            if channel_details:
+                raise ValueError(
+                    "Parameters 'channel_details' are"
+                    " not compatible with 'ascii' format."
+                )
+            ascii_visualiztion_str = self._visualize_ascii()
+            if view:
+                print(ascii_visualiztion_str)
+            return ascii_visualiztion_str
+        try:
+            import graphviz
+        except ImportError:
+            raise ImportError(
+                "Please install graphviz to visualize the compiled graph. "
+                "You can install it by running `pip install graphviz`."
+            )
+        from ray.dag import (
+            InputAttributeNode,
+            InputNode,
+            MultiOutputNode,
+            ClassMethodNode,
+            DAGNode,
+        )
+
+        # Check that the DAG has been compiled
+        if not hasattr(self, "idx_to_task") or not self.idx_to_task:
+            raise ValueError(
+                "The DAG must be compiled before calling 'visualize()'. "
+                "Please call 'experimental_compile()' first."
+            )
+
+        # Check that each CompiledTask has a valid dag_node
+        for idx, task in self.idx_to_task.items():
+            if not hasattr(task, "dag_node") or not isinstance(task.dag_node, DAGNode):
+                raise ValueError(
+                    f"Task at index {idx} does not have a valid 'dag_node'. "
+                    "Ensure that 'experimental_compile()' completed successfully."
+                )
+
+        # Dot file for debugging
+        dot = graphviz.Digraph(name="compiled_graph", format=format)
+        # Give every actor a unique color, colors between 24k -> 40k tested as readable
+        # other colors may be too dark, especially when wrapping back around to 0
+        actor_id_to_color = defaultdict(
+            lambda: f"#{((len(actor_id_to_color) * 2000 + 24000) % 0xFFFFFF):06X}"
+        )
+        # Add nodes with task information
+        for idx, task in self.idx_to_task.items():
+            dag_node = task.dag_node
+            # Initialize the label and attributes
+            label = f"Task {idx}\n"
+            shape = "oval"  # Default shape
+            style = "filled"
+            fillcolor = ""
+
+            # Handle different types of dag_node
+            if isinstance(dag_node, InputNode):
+                label += "InputNode"
+                shape = "rectangle"
+                fillcolor = "lightblue"
+            elif isinstance(dag_node, InputAttributeNode):
+                label += f"InputAttributeNode[{dag_node.key}]"
+                shape = "rectangle"
+                fillcolor = "lightblue"
+            elif isinstance(dag_node, MultiOutputNode):
+                label += "MultiOutputNode"
+                shape = "rectangle"
+                fillcolor = "yellow"
+            elif isinstance(dag_node, ClassMethodNode):
+                if dag_node.is_class_method_call:
+                    # Class Method Call Node
+                    method_name = dag_node.get_method_name()
+                    actor = dag_node._get_actor_handle()
+                    if actor:
+                        class_name = (
+                            actor._ray_actor_creation_function_descriptor.class_name
+                        )
+                        actor_id = actor._actor_id.hex()
+                        label += f"Actor: {class_name}\n"
+                        label += f"ID: {actor_id[:6]}...\n"
+                        label += f"Method: {method_name}"
+                        fillcolor = actor_id_to_color[actor_id]
+                    else:
+                        label += f"Method: {method_name}"
+                        fillcolor = "lightgreen"
+                    shape = "oval"
+                elif dag_node.is_class_method_output:
+                    # Class Method Output Node
+                    label += f"ClassMethodOutputNode[{dag_node.output_idx}]"
+                    shape = "rectangle"
+                    fillcolor = "orange"
+                else:
+                    # Unexpected ClassMethodNode
+                    label += "ClassMethodNode"
+                    shape = "diamond"
+                    fillcolor = "red"
+            else:
+                # Unexpected node type
+                label += type(dag_node).__name__
+                shape = "diamond"
+                fillcolor = "red"
+
+            # Add the node to the graph with attributes
+            dot.node(str(idx), label, shape=shape, style=style, fillcolor=fillcolor)
+            channel_type_str = (
+                (
+                    type(dag_node.type_hint).__name__
+                    if dag_node.type_hint
+                    else "UnknownType"
+                )
+                + "\n"
+                if channel_details
+                else None
+            )
+
+            # This logic is built on the assumption that there will only be multiple
+            # output channels if the task has multiple returns
+            # case: task with one output
+            if len(task.output_channels) == 1:
+                for downstream_node in task.dag_node._downstream_nodes:
+                    downstream_idx = self.dag_node_to_idx[downstream_node]
+                    edge_label = None
+                    if channel_details:
+                        edge_label = channel_type_str
+                        edge_label += self.get_channel_details(
+                            task.output_channels[0],
+                            (
+                                downstream_node._get_actor_handle()._actor_id.hex()
+                                if type(downstream_node) is ClassMethodNode
+                                else self._proxy_actor._actor_id.hex()
+                            ),
+                        )
+                    dot.edge(str(idx), str(downstream_idx), label=edge_label)
+            # case: multi return, output channels connect to class method output nodes
+            elif len(task.output_channels) > 1:
+                assert len(task.output_idxs) == len(task.output_channels)
+                for output_channel, downstream_idx in zip(
+                    task.output_channels, task.output_node_idxs
+                ):
+                    edge_label = None
+                    if channel_details:
+                        edge_label = channel_type_str
+                        edge_label += self.get_channel_details(
+                            output_channel,
+                            task.dag_node._get_actor_handle()._actor_id.hex(),
+                        )
+                    dot.edge(str(idx), str(downstream_idx), label=edge_label)
+            if type(task.dag_node) is InputAttributeNode:
+                # Add an edge from the InputAttributeNode to the InputNode
+                dot.edge(str(self.input_task_idx), str(idx))
+        dot.render(filename, view=view)
+        return dot.source
+
+    def _register_input_output_custom_serializer(self):
+        """
+        Register custom serializers for input, input attribute, and output nodes.
+        """
+        assert self.input_task_idx is not None
+        assert self.output_task_idx is not None
+
+        # Register custom serializers for input node.
+        input_task = self.idx_to_task[self.input_task_idx]
+        input_task.dag_node.type_hint.register_custom_serializer()
+
+        # Register custom serializers for input attribute nodes.
+        for input_attr_task_idx in self.input_attr_task_idxs:
+            input_attr_task = self.idx_to_task[input_attr_task_idx]
+            input_attr_task.dag_node.type_hint.register_custom_serializer()
+
+        # Register custom serializers for output nodes.
+        for output in self.idx_to_task[self.output_task_idx].args:
+            output.type_hint.register_custom_serializer()
+
+    def teardown(self, kill_actors: bool = False):
+        """
+        Teardown and cancel all actor tasks for this DAG. After this
+        function returns, the actors should be available to execute new tasks
+        or compile a new DAG.
+
+        Note: This method is automatically called when the CompiledDAG is destructed
+        or the script exits. However, this should be explicitly called before compiling
+        another graph on the same actors. Python may not garbage collect the
+        CompiledDAG object immediately when you may expect.
+        """
+        if self._is_teardown:
+            return
+
+        monitor = getattr(self, "_monitor", None)
+        if monitor is not None:
+            from ray.dag import DAGContext
+
+            ctx = DAGContext.get_current()
+            monitor.teardown(kill_actors=kill_actors)
+            monitor.join(timeout=ctx.teardown_timeout)
+            # We do not log a warning here if the thread is still alive because
+            # wait_teardown already logs upon teardown_timeout.
+
+        self._is_teardown = True
+
+    def __del__(self):
+        self.teardown()
+
+
+@DeveloperAPI
+def build_compiled_dag_from_ray_dag(
+    dag: "ray.dag.DAGNode",
+    submit_timeout: Optional[float] = None,
+    buffer_size_bytes: Optional[int] = None,
+    enable_asyncio: bool = False,
+    max_inflight_executions: Optional[int] = None,
+    max_buffered_results: Optional[int] = None,
+    overlap_gpu_communication: Optional[bool] = None,
+    default_communicator: Optional[Union[Communicator, str]] = "create",
+) -> "CompiledDAG":
+    compiled_dag = CompiledDAG(
+        submit_timeout,
+        buffer_size_bytes,
+        enable_asyncio,
+        max_inflight_executions,
+        max_buffered_results,
+        overlap_gpu_communication,
+        default_communicator,
+    )
+
+    def _build_compiled_dag(node):
+        compiled_dag._add_node(node)
+        return node
+
+    root = dag._find_root()
+    root.traverse_and_apply(_build_compiled_dag)
+    compiled_dag._get_or_compile()
+    global _compiled_dags
+    _compiled_dags[compiled_dag.get_id()] = compiled_dag
+    return compiled_dag
diff --git a/vllm_musa/hack/diff/collective_fusion.diff b/vllm_musa/hack/diff/collective_fusion.diff
new file mode 100644
index 000000000..8371f8d96
--- /dev/null
+++ b/vllm_musa/hack/diff/collective_fusion.diff
@@ -0,0 +1,22 @@
+diff --git a/vllm/compilation/collective_fusion.py b/vllm/compilation/collective_fusion.py
+index f754fc238..78ffcfcd0 100644
+--- a/vllm/compilation/collective_fusion.py
++++ b/vllm/compilation/collective_fusion.py
+@@ -6,7 +6,7 @@ import torch
+ import torch._inductor.pattern_matcher as pm
+ import torch.fx as fx
+ from torch._inductor.pattern_matcher import PatternMatcherPass
+-from torch.distributed._symmetric_memory import enable_symm_mem_for_group
++# from torch.distributed._symmetric_memory import enable_symm_mem_for_group
+
+ from vllm.config import VllmConfig
+ from vllm.distributed import get_tp_group
+@@ -104,7 +104,7 @@ class AsyncTPPass(VllmInductorPass):
+         super().__init__(config)
+
+         # Enable symmetric memory for the TP process group
+-        enable_symm_mem_for_group(get_tp_group().device_group.group_name)
++        # enable_symm_mem_for_group(get_tp_group().device_group.group_name)
+         self.patterns: PatternMatcherPass = PatternMatcherPass(
+             pass_name="async_tp_pass")
+         GEMMReduceScatterPattern(self.model_dtype,
diff --git a/vllm_musa/hack/diff/config.diff b/vllm_musa/hack/diff/config.diff
new file mode 100644
index 000000000..6ec41665b
--- /dev/null
+++ b/vllm_musa/hack/diff/config.diff
@@ -0,0 +1,53 @@
+diff --git a/vllm/config.py b/vllm/config.py
+index 724f69a38..44fbb621f 100644
+--- a/vllm/config.py
++++ b/vllm/config.py
+@@ -1131,7 +1131,7 @@ class ModelConfig:
+         if not hasattr(self.hf_text_config, "model_type"):
+             return False
+         elif self.hf_text_config.model_type in \
+-            ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp'):
++            ('deepseek_v2', 'deepseek_v3', 'deepseek_mtp', 'kimi_k2'):
+             return self.hf_text_config.kv_lora_rank is not None
+         elif self.hf_text_config.model_type == 'eagle':
+             # if the model is an EAGLE module, check for the
+@@ -4081,6 +4081,9 @@ class CompilationConfig:
+     splitting certain operations such as attention into subgraphs. Thus this
+     flag cannot be used together with splitting_ops. This may provide
+     performance benefits for smaller models."""
++    simple_cuda_graph: bool = False
++    """Whether to use simple CUDA graph, which uses full CUDA graphs without
++    torch.compile. This can speed up the startup time of vLLM."""
+
+     pass_config: PassConfig = field(default_factory=PassConfig)
+     """Custom inductor passes, see PassConfig for more details"""
+@@ -4205,6 +4208,13 @@ class CompilationConfig:
+         if isinstance(self.pass_config, dict):
+             self.pass_config = PassConfig(**self.pass_config)
+
++        if self.simple_cuda_graph:
++            # When using simple CUDA graph, we skip torch.compile.
++            self.level = CompilationLevel.NO_COMPILATION
++            # Simple CUDA graph does not support piecewise CUDA graphs.
++            self.full_cuda_graph = True
++            self.use_cudagraph = True
++
+     def init_backend(self, vllm_config: "VllmConfig") -> Union[str, Callable]:
+         if self.level == CompilationLevel.NO_COMPILATION:
+             raise ValueError("No compilation level is set.")
+@@ -4559,9 +4569,14 @@ class VllmConfig:
+             not self.model_config.enforce_eager:
+             # By default, V1 uses piecewise CUDA graphs. If full_cuda_graph
+             # is set to True, full CUDA graphs will be used.
++            self.use_cudagraph = True
+             self.compilation_config.cudagraph_num_of_warmups = 1
+-            self.compilation_config.level = CompilationLevel.PIECEWISE
+-            self.compilation_config.set_splitting_ops_for_v1()
++            if self.compilation_config.simple_cuda_graph:
++                self.compilation_config.full_cuda_graph = True
++                self.compilation_config.level = CompilationLevel.NO_COMPILATION
++            else:
++                self.compilation_config.level = CompilationLevel.PIECEWISE
++                self.compilation_config.set_splitting_ops_for_v1()
+
+         self._set_cudagraph_sizes()
diff --git a/vllm_musa/hack/diff/deepseek_v2.diff b/vllm_musa/hack/diff/deepseek_v2.diff
new file mode 100644
index 000000000..5f11ef7c7
--- /dev/null
+++ b/vllm_musa/hack/diff/deepseek_v2.diff
@@ -0,0 +1,52 @@
+diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py
+index 2fa1294b7..af9b78ff9 100644
+--- a/vllm/model_executor/models/deepseek_v2.py
++++ b/vllm/model_executor/models/deepseek_v2.py
+@@ -739,14 +739,20 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
+         self.num_expert_groups = config.n_group
+
+         self.moe_layers: list[FusedMoE] = []
++        example_moe = None
+         for layer in self.model.layers:
++            if isinstance(layer, PPMissingLayer):
++                continue
++
+             assert isinstance(layer, DeepseekV2DecoderLayer)
+             if isinstance(layer.mlp, DeepseekV2MoE):
++                # Pick last one layer since the first ones may be dense layers.
++                example_moe = layer.mlp
+                 self.moe_layers.append(layer.mlp.experts)
+
+-        # Pick last one layer since the first ones may be dense layers.
+-        example_moe = typing.cast(
+-            DeepseekV2MoE, self.model.layers[config.num_hidden_layers - 1].mlp)
++        if example_moe is None:
++            raise RuntimeError("No DeepseekV2MoE layer found in model.layers.")
++
+         self.num_logical_experts = example_moe.n_logical_experts
+         self.num_physical_experts = example_moe.n_physical_experts
+         self.num_local_physical_experts = example_moe.n_local_physical_experts
+@@ -882,15 +888,14 @@ class DeepseekV2ForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):
+                     # available replicas.
+                     weight_loader = typing.cast(Callable[..., bool],
+                                                 param.weight_loader)
+-                    success = weight_loader(param,
+-                                            loaded_weight,
+-                                            name_mapped,
+-                                            shard_id=shard_id,
+-                                            expert_id=expert_id,
+-                                            return_success=True)
+-                    if success:
+-                        name = name_mapped
+-                        break
++                    weight_loader(
++                        param,
++                        loaded_weight,
++                        name_mapped,
++                        shard_id=shard_id,
++                        expert_id=expert_id)
++                    name = name_mapped
++                    break
+                 else:
+                     if is_expert_weight:
+                         # We've checked that this is an expert weight
diff --git a/vllm_musa/hack/diff/distributed_utils.diff b/vllm_musa/hack/diff/distributed_utils.diff
new file mode 100644
index 000000000..d31286683
--- /dev/null
+++ b/vllm_musa/hack/diff/distributed_utils.diff
@@ -0,0 +1,19 @@
+diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py
+index 67f71643d..a49e17a09 100644
+--- a/vllm/distributed/utils.py
++++ b/vllm/distributed/utils.py
+@@ -20,8 +20,7 @@ from typing import Any, Optional
+ import torch
+ from torch.distributed import ProcessGroup, TCPStore
+ from torch.distributed.distributed_c10d import (Backend, PrefixStore,
+-                                                _get_default_timeout,
+-                                                _unregister_process_group)
++                                                _get_default_timeout)
+ from torch.distributed.rendezvous import rendezvous
+
+ import vllm.envs as envs
+@@ -533,4 +532,3 @@ def stateless_destroy_torch_distributed_process_group(
+         from torch.distributed.distributed_c10d import _shutdown_backend
+         _shutdown_backend(pg)
+
+-    _unregister_process_group(pg.group_name)
diff --git a/vllm_musa/hack/diff/env_override.diff b/vllm_musa/hack/diff/env_override.diff
new file mode 100644
index 000000000..b6731ed26
--- /dev/null
+++ b/vllm_musa/hack/diff/env_override.diff
@@ -0,0 +1,10 @@
+diff --git a/vllm/env_override.py b/vllm/env_override.py
+index ef425d433..653956df4 100644
+--- a/vllm/env_override.py
++++ b/vllm/env_override.py
+@@ -38,4 +38,4 @@ os.environ['PYTORCH_NVML_BASED_CUDA_CHECK'] = '1'
+ # see https://github.com/vllm-project/vllm/issues/10480
+ os.environ['TORCHINDUCTOR_COMPILE_THREADS'] = '1'
+ # see https://github.com/vllm-project/vllm/issues/10619
+-torch._inductor.config.compile_threads = 1
++import torch_musa
diff --git a/vllm_musa/hack/diff/fusion.diff b/vllm_musa/hack/diff/fusion.diff
new file mode 100644
index 000000000..8ef6cedc8
--- /dev/null
+++ b/vllm_musa/hack/diff/fusion.diff
@@ -0,0 +1,55 @@
+diff --git a/vllm/compilation/fusion.py b/vllm/compilation/fusion.py
+index 951a2861e..69df0dfc5 100644
+--- a/vllm/compilation/fusion.py
++++ b/vllm/compilation/fusion.py
+@@ -29,8 +29,8 @@ def empty_fp32(*args, **kwargs):
+     return torch.empty(*args, **kwargs, dtype=torch.float32, device="cuda")
+
+
+-RMS_OP = torch.ops._C.rms_norm.default
+-RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
++RMS_OP = None
++RMS_ADD_OP = None
+
+
+ # Use proxy as NamedTuple direct subclasses cannot have static members
+@@ -87,12 +87,12 @@ kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, False, GroupShape.PER_TENSOR, True)
+ kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, False, GroupShape.PER_TOKEN, True)
+
+ QUANT_OPS: dict[QuantKey, OpOverload] = {
+-    kFp8StaticTensorSym:
+-    torch.ops._C.static_scaled_fp8_quant.default,  # noqa: E501
+-    kFp8DynamicTensorSym:
+-    torch.ops._C.dynamic_scaled_fp8_quant.default,  # noqa: E501
+-    kFp8DynamicTokenSym:
+-    torch.ops._C.dynamic_per_token_scaled_fp8_quant.default,  # noqa: E501
++    # kFp8StaticTensorSym:
++    # torch.ops._C.static_scaled_fp8_quant,  # noqa: E501
++    # kFp8DynamicTensorSym:
++    # torch.ops._C.dynamic_scaled_fp8_quant,  # noqa: E501
++    # kFp8DynamicTokenSym:
++    # torch.ops._C.dynamic_per_token_scaled_fp8_quant,  # noqa: E501
+ }
+
+
+@@ -111,14 +111,14 @@ class FusedRMSQuantKey(NamedTuple):
+
+
+ FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {
+-    FusedRMSQuantKey(kFp8StaticTensorSym, False):
+-    torch.ops._C.rms_norm_static_fp8_quant.default,  # noqa: E501
+-    FusedRMSQuantKey(kFp8StaticTensorSym, True):
+-    torch.ops._C.fused_add_rms_norm_static_fp8_quant.default,  # noqa: E501
+-    FusedRMSQuantKey(kFp8DynamicTokenSym, False):
+-    torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa: E501
+-    FusedRMSQuantKey(kFp8DynamicTokenSym, True):
+-    torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa: E501
++#     FusedRMSQuantKey(kFp8StaticTensorSym, False):
++#     torch.ops._C.rms_norm_static_fp8_quant,  # noqa: E501
++#     FusedRMSQuantKey(kFp8StaticTensorSym, True):
++#     torch.ops._C.fused_add_rms_norm_static_fp8_quant,  # noqa: E501
++#     FusedRMSQuantKey(kFp8DynamicTokenSym, False):
++#     torch.ops._C.rms_norm_dynamic_per_token_quant,  # noqa: E501
++#     FusedRMSQuantKey(kFp8DynamicTokenSym, True):
++#     torch.ops._C.rms_norm_dynamic_per_token_quant,  # noqa: E501
+ }
diff --git a/vllm_musa/hack/diff/gptq.diff b/vllm_musa/hack/diff/gptq.diff
new file mode 100644
index 000000000..56b0b3632
--- /dev/null
+++ b/vllm_musa/hack/diff/gptq.diff
@@ -0,0 +1,13 @@
+diff --git a/vllm/model_executor/layers/quantization/gptq.py b/vllm/model_executor/layers/quantization/gptq.py
+index d3ab1be3b..5f5ee0f21 100644
+--- a/vllm/model_executor/layers/quantization/gptq.py
++++ b/vllm/model_executor/layers/quantization/gptq.py
+@@ -86,7 +86,7 @@ class GPTQConfig(QuantizationConfig):
+
+     @classmethod
+     def get_supported_act_dtypes(cls) -> list[torch.dtype]:
+-        return [torch.half]
++        return [torch.half, torch.bfloat16]
+
+     @classmethod
+     # Need to figure it out
diff --git a/vllm_musa/hack/diff/requirements_common.diff b/vllm_musa/hack/diff/requirements_common.diff
new file mode 100644
index 000000000..4c9031907
--- /dev/null
+++ b/vllm_musa/hack/diff/requirements_common.diff
@@ -0,0 +1,19 @@
+diff --git a/requirements/common.txt b/requirements/common.txt
+index 8bc0be777..25d35ea17 100644
+--- a/requirements/common.txt
++++ b/requirements/common.txt
+@@ -2,12 +2,12 @@ regex # Replace re for higher-performance regex matching
+ cachetools
+ psutil
+ sentencepiece  # Required for LLaMA tokenizer.
+-numpy
++numpy==1.26.1
+ requests >= 2.26.0
+ tqdm
+ blake3
+ py-cpuinfo
+-transformers >= 4.51.1
++transformers == 4.52.4
+ huggingface-hub[hf_xet] >= 0.33.0  # Required for Xet downloads.
+ tokenizers >= 0.21.1  # Required for fast incremental detokenization.
+ protobuf # Required by LlamaTokenizer.
diff --git a/vllm_musa/hack/diff/vllm_utils.diff b/vllm_musa/hack/diff/vllm_utils.diff
new file mode 100644
index 000000000..a128dc66c
--- /dev/null
+++ b/vllm_musa/hack/diff/vllm_utils.diff
@@ -0,0 +1,12 @@
+diff --git a/vllm/utils/__init__.py b/vllm/utils/__init__.py
+index 9550b056f..25a88b591 100644
+--- a/vllm/utils/__init__.py
++++ b/vllm/utils/__init__.py
+@@ -2256,6 +2256,7 @@ def direct_register_custom_op(
+         dispatch_key: str = "CUDA",
+         tags: tuple[torch.Tag, ...] = (),
+ ):
++    return
+     """
+     `torch.library.custom_op` can have significant overhead because it
+     needs to consider complicated dispatching logic. This function
diff --git a/vllm_musa/hack/hack.sh b/vllm_musa/hack/hack.sh
new file mode 100644
index 000000000..5f25fa04d
--- /dev/null
+++ b/vllm_musa/hack/hack.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+echo "Hack some musa sdk!"
\ No newline at end of file
diff --git a/vllm_musa/hack/impl.py b/vllm_musa/hack/impl.py
new file mode 100644
index 000000000..9dc5552c9
--- /dev/null
+++ b/vllm_musa/hack/impl.py
@@ -0,0 +1,1088 @@
+import dataclasses
+import functools
+import inspect
+import sys
+import typing
+import weakref
+
+from torchgen.model import FunctionSchema, OperatorName, SchemaKind, BaseType, ListType, BaseTy
+
+import torch
+import torch._C as _C
+import torch.library as library
+from torch._library.abstract_impl import AbstractImplCtx
+from torch.library import get_ctx
+
+from .autograd import autograd_kernel_indirection, construct_autograd_kernel
+
+"""
+For a detailed guide on custom ops, please see
+https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
+
+This file includes pieces of the implementation of our custom operator API.
+"""
+
+__all__ = ["custom_op", "CustomOp", "get_ctx", "AbstractImplCtx"]
+
+
+SUPPORTED_DEVICE_TYPE_TO_KEY = {
+    "cpu": "CPU",
+    "cuda": "CUDA",
+}
+
+# We will not let users register CustomOps with anything that could look like
+# PyTorch internals to avoid confusion.
+RESERVED_NS = {
+    "prim",
+    "prims",
+    "aten",
+    "at",
+    "torch",
+    "pytorch",
+}
+
+
+def custom_op(
+    qualname: str, manual_schema: typing.Optional[str] = None
+) -> typing.Callable:
+    r"""Creates a new CustomOp object.
+
+    WARNING: if you're a user, please do not use this directly
+    (instead use the torch._custom_ops APIs).
+    Also please see the following for a detailed guide on custom ops.
+    https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
+
+    In PyTorch, defining an op (short for "operator") is a two step-process:
+    - we need to define (create) the op
+    - we need to implement behavior for how the operator interacts with
+      various PyTorch subsystems, like CPU/CUDA Tensors, Autograd, etc.
+
+    This entrypoint defines the CustomOp object (the first step);
+    you must then perform the second step by calling various methods on
+    the CustomOp object.
+
+    This API is used as a decorator (see examples).
+
+    Arguments:
+        qualname (str): Should be a string that looks like
+            "namespace::operator_name". Operators in PyTorch need a namespace to
+            avoid name collisions; a given operator may only be created once.
+            If you are writing a Python library, we recommend the namespace to
+            be the name of your top-level module. The operator_name must be
+            the same as the name of the function you pass to custom_op
+            (see examples).
+        manual_schema (Optional[str]): Each PyTorch operator needs a schema that
+            tells PyTorch the types of the inputs/outputs. If None (default),
+            we will infer the schema from the type annotations on the function
+            (see examples). Otherwise, if you don't want to use type annotations,
+            you may provide us the schema string.
+
+    Example::
+        >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+        >>> import numpy as np
+        >>> from torch import Tensor
+        >>>
+        >>> # Step 1: define the CustomOp.
+        >>> # We need to provide the decorator a "prototype function"
+        >>> # (a function with Python ellipses as the body).
+        >>> @custom_op("my_library::numpy_sin")
+        >>> def numpy_sin(x: Tensor) -> Tensor:
+        >>>     ...
+        >>>
+        >>> # numpy_sin is now an instance of class CustomOp
+        >>> print(type(numpy_sin))
+        >>>
+        >>> # Step 2: Register an implementation for various PyTorch subsystems
+        >>>
+        >>> # Register an implementation for CPU tensors
+        >>> @numpy_sin.impl('cpu')
+        >>> def numpy_sin_impl_cpu(x):
+        >>>     return torch.from_numpy(np.sin(x.numpy()))
+        >>>
+        >>> # Register an implementation for CUDA tensors
+        >>> @numpy_sin.impl('cuda')
+        >>> def numpy_sin_impl_cuda(x):
+        >>>     return torch.from_numpy(np.sin(x.cpu().numpy())).to(x.device)
+        >>>
+        >>> x = torch.randn(3)
+        >>> numpy_sin(x)  # calls numpy_sin_impl_cpu
+        >>>
+        >>> x_cuda = x.cuda()
+        >>> numpy_sin(x)  # calls numpy_sin_impl_cuda
+
+    """
+
+    def inner(func):
+        if not inspect.isfunction(func):
+            raise ValueError(
+                f"custom_op(...)(func): Expected `func` to be a Python "
+                f"function, got: {type(func)}"
+            )
+
+        ns, name = parse_qualname(qualname)
+        validate_namespace(ns)
+        if func.__name__ != name:
+            raise ValueError(
+                f"custom_op(qualname='{qualname}', ...)(func): expected `func` "
+                f"to have name '{name}' but got '{func.__name__}'. "
+                f"Please either change the name of `func` or the qualname that "
+                f"is passed to `custom_op`"
+            )
+
+        schema = infer_schema(func, mutates_args=()) if manual_schema is None else manual_schema
+        schema_str = f"{name}{schema}"
+        function_schema = FunctionSchema.parse(schema_str)
+        validate_schema(function_schema)
+        if manual_schema is not None:
+            validate_function_matches_schema(function_schema, func)
+
+        lib = library.Library(ns, "FRAGMENT")
+        lib.define(schema_str)
+        ophandle = find_ophandle_or_throw(ns, function_schema.name)
+        result = CustomOp(lib, ns, function_schema, name, ophandle, _private_access=True)
+
+        result.__name__ = func.__name__
+        result.__module__ = func.__module__
+        result.__doc__ = func.__doc__
+
+        library.impl(lib, result._opname, "Autograd")(
+            autograd_kernel_indirection(weakref.proxy(result))
+        )
+
+        torch._C._dispatch_set_report_error_callback(
+            ophandle, functools.partial(report_error_callback, weakref.proxy(result))
+        )
+
+        return result
+
+    return inner
+
+
+# Global dictionary holding references to all CustomOp objects
+# Yes, it keeps all CustomOps alive (see NOTE [CustomOp lifetime])
+# Used to query the CustomOp associated with a specific C++ dispatcher operator.
+# An example usage is FakeTensor: FakeTensor checks if a specific operator
+# has an implementation registered via the CustomOp API.
+# Indexed by qualname (e.g. aten::foo)
+global_registry: typing.Dict[str, "CustomOp"] = {}
+
+
+class CustomOp:
+    r"""Class for custom operators in PyTorch.
+
+    Use the CustomOp API to create user-defined custom operators that behave
+    just like regular PyTorch operators (e.g. torch.sin, torch.mm) when it
+    comes to various PyTorch subsystems (like torch.compile).
+
+    To construct a `CustomOp`, use `custom_op`.
+    """
+
+    def __init__(self, lib, cpp_ns, schema, operator_name, ophandle, *, _private_access=False):
+        super().__init__()
+        if not _private_access:
+            raise RuntimeError(
+                "The CustomOp constructor is private and we do not guarantee "
+                "BC for it. Please use custom_op(...) to create a CustomOp object"
+            )
+        name = f"{cpp_ns}::{operator_name}"
+        self._schema = schema
+        self._cpp_ns = cpp_ns
+        self._lib: library.Library = lib
+        self._ophandle: _C._DispatchOperatorHandle = ophandle
+        # Has the name of the op, e.g. "foo". We cache here for convenience.
+        self._opname: str = operator_name
+        # this is _opname but with namespace. e.g. "custom::foo"
+        self._qualname: str = name
+        self.__name__ = None  # mypy requires this
+        # NB: Some of these impls are registered as kernels to DispatchKeys.
+        # Modifying the _impls dict directly won't do anything in that case.
+        self._impls: typing.Dict[str, typing.Optional[FuncAndLocation]] = {}
+        # See NOTE [CustomOp autograd kernel indirection]
+        self._registered_autograd_kernel_indirection = False
+
+        global_registry[self._qualname] = self
+
+    def _register_autograd_kernel_indirection(self):
+        assert not self._registered_autograd_kernel_indirection
+        self._lib.impl(self._opname, autograd_kernel_indirection(weakref.proxy(self)), "Autograd")
+        self._registered_autograd_kernel_indirection = True
+
+    # Records the impl and the source location in self._impls
+    # Note that this doesn't cause torch.library to use the impl, that
+    # needs to be done in a separate self._lib.impl call.
+    def _register_impl(self, kind, func, stacklevel=2):
+        if self._has_impl(kind):
+            func_and_location = self._impls[kind]
+            assert func_and_location is not None  # Pacify mypy
+            location = func_and_location.location
+            raise RuntimeError(
+                f"Attempting to register a {kind} impl for operator {self._qualname} "
+                f"that already has a {kind} impl registered from Python at "
+                f"{location}. This is not supported."
+            )
+        frame = inspect.getframeinfo(sys._getframe(stacklevel))
+        location = f"{frame.filename}:{frame.lineno}"
+        self._impls[kind] = FuncAndLocation(func, location)
+
+    def _get_impl(self, kind):
+        return self._impls[kind]
+
+    def _has_impl(self, kind):
+        return kind in self._impls
+
+    def _destroy(self):
+        # NOTE: [CustomOp lifetime]
+        # A CustomOp, once created, lives forever. The mechanism is that the
+        # global registry holds a reference to it. However, to make testing
+        # easier, we want to be able to destroy CustomOp objects.
+        # CustomOp._destroy does the job, though it leaves the CustomOp
+        # in a garbage state.
+        del self._lib
+
+        opnamespace = getattr(torch.ops, self._cpp_ns)
+        if hasattr(opnamespace, self._opname):
+            delattr(opnamespace, self._opname)
+
+        del global_registry[self._qualname]
+
+    def __repr__(self):
+        return f'<CustomOp(op="{self._qualname}")>'
+
+    def __call__(self, *args, **kwargs):
+        # Bypass torch.ops.* and directly do OperatorHandle::callBoxed.
+        # Using torch.ops.* is a bit of a pain (it can be slow and it has lifetime
+        # issues from caching operators that make testing CustomOp difficult).
+        result = _C._dispatch_call_boxed(self._ophandle, *args, **kwargs)
+        return result
+
+    def impl(
+        self, device_types: typing.Union[str, typing.Iterable[str]], _stacklevel=2,
+    ) -> typing.Callable:
+        r"""Register an implementation for a device type for this CustomOp object.
+
+        WARNING: if you're a user, please do not use this directly
+        (instead use the torch._custom_ops APIs).
+        Also please see the following for a detailed guide on custom ops.
+        https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
+
+        If the CustomOp is passed multiple Tensor inputs with different device
+        types, it will dispatch to the registered implementation for the highest
+        priority device type among those present.
+        The supported device types, in order of priority, are {'cuda', 'cpu'}.
+
+        This API is used as a decorator (see examples).
+
+        Arguments:
+            device_types (str or Iterable[str]): the device type(s) to register the function for.
+
+        Examples::
+            >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA)
+            >>> import numpy as np
+            >>> from torch import Tensor
+            >>>
+            >>> @custom_op("my_library::numpy_cos")
+            >>> def numpy_cos(x: Tensor) -> Tensor:
+            >>>     ...
+            >>>
+            >>> # Register an implementation for CPU Tensors
+            >>> @numpy_cos.impl('cpu')
+            >>> def numpy_cos_impl_cpu(x):
+            >>>     return torch.from_numpy(np.cos(x.numpy()))
+            >>>
+            >>> # Register an implementation for CUDA Tensors
+            >>> @numpy_cos.impl('cuda')
+            >>> def numpy_cos_impl_cuda(x):
+            >>>     return torch.from_numpy(np.cos(x.cpu().numpy())).to(x.device)
+            >>>
+            >>> x = torch.randn(3)
+            >>> numpy_cos(x)  # calls numpy_cos_impl_cpu
+            >>>
+            >>> x_cuda = x.cuda()
+            >>> numpy_cos(x)  # calls numpy_cos_impl_cuda
+
+        """
+        if isinstance(device_types, str):
+            device_types = [device_types]
+        for device_type in device_types:
+            validate_device_type(device_type)
+
+        def inner(f):
+            for device_type in set(device_types):
+                self._check_doesnt_have_library_impl(device_type)
+                self._register_impl(device_type, f, stacklevel=_stacklevel)
+                dispatch_key = SUPPORTED_DEVICE_TYPE_TO_KEY[device_type]
+                library.impl(self._lib, self._opname, dispatch_key)(f)
+            return f
+
+        return inner
+
+    def _check_doesnt_have_library_impl(self, device_type):
+        if self._has_impl(device_type):
+            return
+        key = SUPPORTED_DEVICE_TYPE_TO_KEY[device_type]
+        if _C._dispatch_has_computed_kernel_for_dispatch_key(self._qualname, key):
+            raise RuntimeError(
+                f"impl(..., device_types={device_type}): the operator {self._qualname} "
+                f"already has an implementation for this device type via a "
+                f"pre-existing torch.library or TORCH_LIBRARY registration.")
+
+    def impl_factory(self) -> typing.Callable:
+        r"""Register an implementation for a factory function."""
+
+        def inner(f):
+            self._register_impl("factory", f)
+            library.impl(self._lib, self._opname, "BackendSelect")(f)
+            return f
+
+        return inner
+
+    def impl_abstract(self, _stacklevel=2) -> typing.Callable:
+        r"""Register an abstract implementation for this operator.
+
+        WARNING: please do not use this directly (and instead use the torch._custom_ops
+        APIs). Also please see the following for a detailed guide on custom ops.
+        https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
+
+        An "abstract implementation" specifies the behavior of this operator on
+        Tensors that carry no data. Given some input Tensors with certain properties
+        (sizes/strides/storage_offset/device), it specifies what the properties of
+        the output Tensors are.
+
+        The abstract implementation has the same signature as the operator.
+        It is run for both FakeTensors and meta tensors. To write an abstract
+        implementation, assume that all Tensor inputs to the operator are
+        regular CPU/CUDA/Meta tensors, but they do not have storage, and
+        you are trying to return regular CPU/CUDA/Meta tensor(s) as output.
+        The abstract implementation must consist of only PyTorch operations
+        (and may not directly access the storage or data of any input or
+        intermediate Tensors).
+
+        This API is used as a decorator (see examples).
+
+        Examples::
+            >>> import numpy as np
+            >>> from torch import Tensor
+            >>>
+            >>> # Example 1: an operator without data-dependent output shape
+            >>> @custom_op('my_library::custom_linear')
+            >>> def custom_linear(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
+            >>>     ...
+            >>>
+            >>> @custom_linear.impl_abstract()
+            >>> def custom_linear_abstract(x, weight):
+            >>>     assert x.dim() == 2
+            >>>     assert weight.dim() == 2
+            >>>     assert bias.dim() == 1
+            >>>     assert x.shape[1] == weight.shape[1]
+            >>>     assert weight.shape[0] == bias.shape[0]
+            >>>     assert x.device == weight.device
+            >>>
+            >>>     return (x @ weight.t()) + bias
+            >>>
+            >>> # Example 2: an operator with data-dependent output shape
+            >>> @custom_op('my_library::custom_nonzero')
+            >>> def custom_nonzero(x: Tensor) -> Tensor:
+            >>>     ...
+            >>>
+            >>> @custom_nonzero.impl_abstract()
+            >>> def custom_nonzero_abstract(x):
+            >>>     # Number of nonzero-elements is data-dependent.
+            >>>     # Since we cannot peek at the data in an abstract impl,
+            >>>     # we use the ctx object to construct a new symint that
+            >>>     # represents the data-dependent size.
+            >>>     ctx = torch._custom_op.get_ctx()
+            >>>     nnz = ctx.create_unbacked_symint()
+            >>>     shape = [x.dim(), nnz]
+            >>>     result = x.new_empty(shape, dtype=torch.long)
+            >>>     return result
+            >>>
+            >>> @custom_nonzero.impl(['cpu', 'cuda'])
+            >>> def custom_nonzero_impl(x):
+            >>>     x_np = to_numpy(x)
+            >>>     res = np.stack(np.nonzero(x_np), axis=1)
+            >>>     # unbacked symbolic ints in PyTorch must be >= 2, so we
+            >>>     # constrain the range to at least 2
+            >>>     if res.shape[0] <= 1:
+            >>>         raise RuntimeError("not supported")
+            >>>     return torch.tensor(res, device=x.device)
+
+        """
+
+        def inner(f):
+            self._check_doesnt_have_library_meta_impl()
+            self._register_impl("abstract", f, stacklevel=_stacklevel)
+            location = self._get_impl("abstract").location
+
+            qualname = self._qualname
+
+            # Handle DispatchKey.Meta registration
+            @functools.wraps(f)
+            def f_with_ctx(*args, **kwargs):
+                def error_on_ctx():
+                    raise RuntimeError(
+                        f"Attempted to call get_ctx() for the meta implementation "
+                        f"for {qualname}."
+                        f"You have presumably called get_ctx() because the operator "
+                        f"has a data-dependent output shape; if so, there is no "
+                        f"such meta implementation and this error is the correct "
+                        f"behavior. Otherwise, please remove the call to get_ctx() "
+                        f"in the implementation registered with impl_abstract "
+                        f"at {location}"
+                    )
+
+                with torch._library.abstract_impl.set_ctx_getter(error_on_ctx):
+                    return f(*args, **kwargs)
+
+            self._lib.impl(self._opname, f_with_ctx, "Meta")
+            return f
+
+        return inner
+
+    def _check_can_register_backward(self):
+        def error(detail):
+            raise RuntimeError(
+                f"Cannot use torch._custom_ops APIs to register backward "
+                f"formula for {detail}. Got operator "
+                f"{self._qualname} with schema: {schema}"
+            )
+
+        schema = self._schema
+        if schema.kind() != SchemaKind.functional:
+            error("non-functional operator")
+
+        rets = schema.returns
+        if not schema.returns:
+            error("operator with no returns")
+
+        assert len(rets) > 0
+        is_non_mutating_view = any(
+            r.annotation is not None and not r.annotation.is_write for r in rets
+        )
+        if is_non_mutating_view:
+            error("operator that returns views")
+
+        # We make assumptions about the schema's return types.
+        allowed_return_types = {
+            BaseType(BaseTy.int): "int",
+            BaseType(BaseTy.SymInt): "SymInt",
+            BaseType(BaseTy.bool): "bool",
+            BaseType(BaseTy.float): "float",
+            BaseType(BaseTy.Tensor): "Tensor",
+            ListType(BaseType(BaseTy.Tensor), None): "List[Tensor]",
+        }
+        for ret in schema.returns:
+            if ret.type in allowed_return_types:
+                continue
+            error(f"operator with return not in {list(allowed_return_types.values())} (got {ret.type})")
+
+    def _check_doesnt_have_library_autograd_impl(self):
+        if self._registered_autograd_kernel_indirection:
+            return
+
+        if _C._dispatch_has_kernel_for_dispatch_key(self._qualname, "CompositeImplicitAutograd"):
+            raise RuntimeError(
+                f"impl_backward/impl_save_for_backward: the operator {self._qualname} "
+                f"already has an implementation for this device type via a "
+                f"pre-existing registration to DispatchKey::CompositeImplicitAutograd."
+                f"CompositeImplicitAutograd operators do not need an autograd formula; "
+                f"instead, the operator will decompose into its constituents and those "
+                f"can have autograd formulas defined on them.")
+
+        # We can improve this by adding "all Autograd<BACKEND> keys", but
+        # realistically people will just be using this API for CPU/CUDA for now.
+        for key in ["Autograd", "AutogradCPU", "AutogradCUDA"]:
+            if _C._dispatch_has_kernel_for_dispatch_key(self._qualname, key):
+                raise RuntimeError(
+                    f"impl_backward/impl_save_for_backward: "
+                    f"the operator {self._qualname} already has an Autograd kernel "
+                    f"registered to DispatchKey::{key} vi a pre-existing "
+                    f"torch.library or TORCH_LIBRARY registration. Please either "
+                    f"remove those registrations or don't use the torch._custom_ops APIs")
+
+    def _check_doesnt_have_library_meta_impl(self):
+        if self._has_impl("abstract"):
+            return
+
+        # If the user's operator is CompositeExplicitAutograd,
+        # allow them to impl_abstract. This is being pragmatic
+        # (existing custom ops may have CompositeExplicitAutograd
+        # registration that don't work with Meta kernels, so this
+        # gives them an escape hatch).
+        if (
+            _C._dispatch_has_kernel_for_dispatch_key(self._qualname, "CompositeExplicitAutograd")
+            and not _C._dispatch_has_kernel_for_dispatch_key(self._qualname, "Meta")
+        ):
+            return
+
+        # Otherwise, if the user's already has a Meta kernel or their
+        # op is CompositeImplicitAutograd or some other alias dispatch key,
+        # raise.
+
+        # Special case for CompositeImplicitAutograd
+        if _C._dispatch_has_kernel_for_dispatch_key(self._qualname, "CompositeImplicitAutograd"):
+            raise RuntimeError(
+                f"impl_abstract(...): the operator {self._qualname} "
+                f"already has an implementation for this device type via a "
+                f"pre-existing registration to DispatchKey::CompositeImplicitAutograd."
+                f"CompositeImplicitAutograd operators do not need an abstract impl; "
+                f"instead, the operator will decompose into its constituents and those "
+                f"can have abstract impls defined on them.")
+
+        if _C._dispatch_has_kernel_for_dispatch_key(self._qualname, "Meta"):
+            raise RuntimeError(
+                f"impl_abstract(...): the operator {self._qualname} "
+                f"already has an DispatchKey::Meta implementation via a "
+                f"pre-existing torch.library or TORCH_LIBRARY registration. "
+                f"Please either remove that registration or don't call impl_abstract.")
+
+    # NOTE ["backward", "save_for_backward", and "autograd"]
+    # As a part of the explicit autograd API, a user must provide us
+    # a "save_for_backward" function and a "backward" function.
+    # When both of these have been provided, then we automatically
+    # construct the "autograd" kernel.
+    def _register_autograd_kernel(self):
+        assert self._has_impl("backward")
+        assert self._has_impl("save_for_backward")
+        kernel = construct_autograd_kernel(
+            self._schema,
+            self._output_differentiability,
+            self,
+            get_op(self._qualname),
+            self._get_impl("save_for_backward").func,
+            self._get_impl("backward").func)
+        self._register_impl("autograd", kernel)
+
+    def impl_save_for_backward(self, _stacklevel=2):
+        r"""Register a function that tells us what to save for backward.
+
+        Please see impl_backward for more details.
+        """
+        def inner(f):
+            self._check_can_register_backward()
+            self._check_doesnt_have_library_autograd_impl()
+            if not self._registered_autograd_kernel_indirection:
+                self._register_autograd_kernel_indirection()
+            self._register_impl("save_for_backward", f, stacklevel=_stacklevel)
+            if self._has_impl("backward"):
+                self._register_autograd_kernel()
+        return inner
+
+    def impl_backward(self, output_differentiability=None, _stacklevel=2):
+        r"""Registers a backward formula.
+
+        WARNING: if you're a user, please do not use this directly
+        (instead use the torch._custom_ops APIs).
+        Also please see the following for a detailed guide on custom ops.
+        https://docs.google.com/document/d/1aGWtgxV3HppuxQAdddyPrs74_aEntpkYt9MalnCKnhk
+
+        In order for the CustomOp to work with autograd, you need to register
+        a backward formula. There are two pieces to this:
+        1. You must give us a function to specify what to save for backward.
+           Call this the "save for backward" function.
+        2. You must give us a function that computes gradients. Call this the
+           "backward" function.
+
+        Use `impl_save_for_backward` to define a "save for backward" function
+        that specifies what gets saved for backward. The function should accept
+        two arguments ``(inputs, output)`` and return the quantities to be saved
+        for backward.
+
+        During runtime, when you call the CustomOp, PyTorch will invoke the
+        "save for backward" function with the inputs and output of the CustomOp.
+
+        Use `impl_backward` to define the "backward" function. The backward
+        function must accept ``(ctx, saved, *grads)``:
+        - ``ctx`` is a context object where we may provide information
+        - ``saved`` is exactly what gets returned from the "save for backward"
+          function
+        - ``grads`` is one or more gradients. The number of gradients matches
+          the number of outputs of the CustomOp.
+
+        The backward function must return a dict that maps the name of
+        an input to the CustomOp to its corresponding gradient. All inputs that
+        were declared to be Tensors in the CustomOp definition must be accounted
+        for in the dict. The gradient may be a Tensor or None.
+
+        """
+        if output_differentiability is not None:
+            def yell():
+                raise RuntimeError(
+                    f"impl_backward(output_differentiability): expected "
+                    f"output_differentiability to be a list of bools with "
+                    f"length equal to the number of outputs of this CustomOp "
+                    f"got: {output_differentiability}")
+
+            if not isinstance(output_differentiability, list):
+                yell()
+            for diff in output_differentiability:
+                if not isinstance(diff, bool):
+                    yell()
+            if len(self._schema.returns) != len(output_differentiability):
+                yell()
+
+        def inner(f):
+            self._check_can_register_backward()
+            self._check_doesnt_have_library_autograd_impl()
+            if not self._registered_autograd_kernel_indirection:
+                self._register_autograd_kernel_indirection()
+            self._register_impl("backward", f, stacklevel=_stacklevel)
+            self._output_differentiability = output_differentiability
+            if self._has_impl("save_for_backward"):
+                self._register_autograd_kernel()
+        return inner
+
+
+@dataclasses.dataclass
+class FuncAndLocation:
+    func: typing.Callable
+    location: str
+
+
+def find_ophandle_or_throw(cpp_ns: str, operator_name: OperatorName):
+    overload_name = (
+        "" if operator_name.overload_name is None else operator_name.overload_name
+    )
+    return _C._dispatch_find_schema_or_throw(
+        f"{cpp_ns}::{str(operator_name.name)}", overload_name
+    )
+
+
+def validate_namespace(ns: str) -> None:
+    if "." in ns:
+        raise ValueError(
+            f'custom_op(..., ns="{ns}"): expected ns to not contain any . (and be a '
+            f"valid variable name)"
+        )
+    if ns in RESERVED_NS:
+        raise ValueError(
+            f"custom_op(..., ns='{ns}'): '{ns}' is a reserved namespace, "
+            f"please choose something else. "
+        )
+
+def validate_schema(schema: FunctionSchema) -> None:
+    if not torch._library.utils.is_functional_schema(schema):
+        raise ValueError(
+            f"custom_op only supports functional operators "
+            f"(ops that do not mutate any inputs, do not return "
+            f"views of the inputs, and has at least one return). "
+            f"Got the following non-functional schema: {schema}"
+        )
+
+    # For simplicity: don't allow self arguments
+    if schema.arguments.self_arg is not None:
+        raise ValueError(
+            f"custom_op does not support arguments named 'self'. Please "
+            f"rename your argument. Got: {schema}"
+        )
+
+
+def parse_qualname(qualname: str) -> typing.Tuple[str, str]:
+    names = qualname.split("::", 1)
+    if len(names) != 2:
+        raise ValueError(f"Expected there to be a namespace in {qualname}, i.e. The "
+                         f"operator name should look something like ns::foo")
+    if '.' in names[1]:
+        raise ValueError(f"The torch.custom_ops APIs do not handle overloads, "
+                         f"i.e. operator names with '.' in them. "
+                         f"Please name your operator something like ns::foo. "
+                         f"Got: {qualname}")
+    return names[0], names[1]
+
+
+def validate_device_type(device_type: str) -> None:
+    if device_type not in SUPPORTED_DEVICE_TYPE_TO_KEY:
+        raise ValueError(
+            f"CustomOp.impl(device_types=[{device_type}, ...]): we only support device_type "
+            f"in {SUPPORTED_DEVICE_TYPE_TO_KEY.keys()}."
+        )
+
+
+def supported_param(param: inspect.Parameter) -> bool:
+    return param.kind in (
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        inspect.Parameter.KEYWORD_ONLY,
+    )
+
+
+def validate_function_matches_schema(
+    schema: FunctionSchema, func: typing.Callable
+) -> None:
+    sig = inspect.signature(func)
+
+    if not all(supported_param(p) for _, p in sig.parameters.items()):
+        raise ValueError(
+            f"custom_op(..., manual_schema)(func): positional-only args, "
+            f"varargs, and kwargs are not supported. Please rewrite `func` "
+            f"to not have them. Got `func` with signature: {sig}"
+        )
+
+    if (
+        any(
+            p.annotation is not inspect.Parameter.empty
+            for _, p in sig.parameters.items()
+        )
+        or sig.return_annotation is not inspect.Signature.empty
+    ):
+        raise ValueError(
+            f"custom_op(..., manual_schema)(func): When passing in a manual "
+            f"schema, we expect `func` to have no type annotations to avoid "
+            f"ambiguity. Got `func` with signature: {sig}"
+        )
+
+    positional = [
+        (name, param)
+        for name, param in sig.parameters.items()
+        if param.kind == inspect.Parameter.POSITIONAL_OR_KEYWORD
+    ]
+    kwargonly = [
+        (name, param)
+        for name, param in sig.parameters.items()
+        if param.kind == inspect.Parameter.KEYWORD_ONLY
+    ]
+
+    def error():
+        raise ValueError(
+            f"custom_op(..., manual_schema)(func): When passing in a manual "
+            f"schema, we expect `func`'s signature to match `manual_schema` "
+            f"(aside from type annotations). "
+            f"func's signature: {sig}, manual_schema: {schema}"
+        )
+
+    def error_default_args():
+        raise ValueError(
+            f"custom_op(..., manual_schema)(func): "
+            f"neither func nor manual_schema should have default "
+            f"arguments. Got "
+            f"func's signature: {sig}, manual_schema: {schema}"
+        )
+
+    def compare(sig_args, schema_args):
+        if len(sig_args) != len(schema_args):
+            error()
+        for (name, param), arg in zip(sig_args, schema_args):
+            if name != arg.name:
+                error()
+            if param.default is not inspect.Parameter.empty or arg.default is not None:
+                error_default_args()
+
+    compare(positional, schema.arguments.flat_positional)
+    compare(kwargonly, schema.arguments.flat_kwarg_only)
+
+
+def infer_schema(
+    prototype_function: typing.Callable,
+    /,
+    *,
+    mutates_args,
+    op_name: typing.Optional[str] = None,
+) -> str:
+    r"""Parses the schema of a given function with type hints. The schema is inferred from the
+    function's type hints, and can be used to define a new operator.
+
+    We make the following assumptions:
+
+    * None of the outputs alias any of the inputs or each other.
+    * | String type annotations "device, dtype, Tensor, types" without library specification are
+      | assumed to be torch.*. Similarly, string type annotations "Optional, List, Sequence, Union"
+      | without library specification are assumed to be typing.*.
+    * | Only the args listed in ``mutates_args`` are being mutated. If ``mutates_args`` is "unknown",
+      | it assumes that all inputs to the operator are being mutates.
+
+    Callers (e.g. the custom ops API) are responsible for checking these assumptions.
+
+    Args:
+        prototype_function: The function from which to infer a schema for from its type annotations.
+        op_name (Optional[str]): The name of the operator in the schema. If ``name`` is None, then the
+            name is not included in the inferred schema. Note that the input schema to
+            ``torch.library.Library.define`` requires a operator name.
+        mutates_args ("unknown" | Iterable[str]): The arguments that are mutated in the function.
+
+    Returns:
+        The inferred schema.
+
+    Example:
+        >>> def foo_impl(x: torch.Tensor) -> torch.Tensor:
+        >>>     return x.sin()
+        >>>
+        >>> infer_schema(foo_impl, op_name="foo", mutates_args={})
+        foo(Tensor x) -> Tensor
+        >>>
+        >>> infer_schema(foo_impl, mutates_args={})
+        (Tensor x) -> Tensor
+    """
+    UNKNOWN_MUTATES = "unknown"
+    sig = inspect.signature(prototype_function)
+
+    def error_fn(what):
+        raise ValueError(
+            f"infer_schema(func): {what} " f"Got func with signature {sig})"
+        )
+
+    def convert_type_string(annotation_type: str):
+        try:
+            return eval(annotation_type)
+        except Exception as e:
+            error_fn(
+                f"Unsupported type annotation {annotation_type}. It is not a type."
+            )
+
+    params = []
+    seen_args = set()
+    saw_kwarg_only_arg = False
+    for idx, (name, param) in enumerate(sig.parameters.items()):
+        if not supported_param(param):
+            error_fn("We do not support positional-only args, varargs, or varkwargs.")
+
+        if param.kind == inspect.Parameter.KEYWORD_ONLY:
+            # The first time we see a kwarg-only arg, add "*" to the schema.
+            if not saw_kwarg_only_arg:
+                params.append("*")
+                saw_kwarg_only_arg = True
+
+        if param.annotation is inspect.Parameter.empty:
+            error_fn(f"Parameter {name} must have a type annotation.")
+
+        # The annotation might be converted to a string by annotation,
+        # we convert it to the actual type.
+        annotation_type = param.annotation
+        if type(annotation_type) == str:
+            annotation_type = convert_type_string(annotation_type)
+
+        if annotation_type not in SUPPORTED_PARAM_TYPES.keys():
+            if annotation_type.__origin__ is tuple:
+                list_type = tuple_to_list(annotation_type)
+                example_type_str = "\n\n"
+                # Only suggest the list type if this type is supported.
+                if list_type in SUPPORTED_PARAM_TYPES.keys():
+                    example_type_str = f"For example, {list_type}.\n\n"
+                error_fn(
+                    f"Parameter {name} has unsupported type {param.annotation}. "
+                    f"We do not support Tuple inputs in schema. As a workaround, please try to use List instead. "
+                    f"{example_type_str}"
+                    f"The valid types are: {SUPPORTED_PARAM_TYPES.keys()}."
+                )
+            else:
+                error_fn(
+                    f"Parameter {name} has unsupported type {param.annotation}. "
+                    f"The valid types are: {SUPPORTED_PARAM_TYPES.keys()}."
+                )
+
+        schema_type = SUPPORTED_PARAM_TYPES[annotation_type]
+        if type(mutates_args) == str:
+            if mutates_args != UNKNOWN_MUTATES:
+                raise ValueError(
+                    "mutates_args must either be a sequence of the names of "
+                    "the arguments that are mutated or the string 'unknown'. "
+                )
+            if schema_type.startswith("Tensor"):
+                schema_type = f"Tensor(a{idx}!){schema_type[len('Tensor'):]}"
+        elif name in mutates_args:
+            if not schema_type.startswith("Tensor"):
+                error_fn(
+                    f"Parameter {name} is in mutable_args but only Tensors or collections of Tensors can be mutated"
+                )
+            schema_type = f"Tensor(a{idx}!){schema_type[len('Tensor'):]}"
+        seen_args.add(name)
+        if param.default is inspect.Parameter.empty:
+            params.append(f"{schema_type} {name}")
+        else:
+            default_repr = None
+            if param.default is None or isinstance(param.default, (int, float, bool)):
+                default_repr = str(param.default)
+            elif isinstance(param.default, (str, torch.device)):
+                default_repr = f'"{param.default}"'
+            elif isinstance(param.default, torch.dtype):
+                dtype_repr = str(param.default)
+                torch_dot = "torch."
+                assert dtype_repr.startswith(torch_dot)
+                default_repr = dtype_repr[len(torch_dot) :]
+            else:
+                error_fn(
+                    f"Parameter {name} has an unsupported default value type {type(param.default)}. "
+                    f"Please file an issue on GitHub so we can prioritize this."
+                )
+            params.append(f"{schema_type} {name}={default_repr}")
+    if mutates_args != UNKNOWN_MUTATES:
+        mutates_args_not_seen = set(mutates_args) - seen_args
+        if len(mutates_args_not_seen) > 0:
+            error_fn(
+                f"{mutates_args_not_seen} in mutates_args were not found in "
+                f"the custom op's signature. "
+                f"mutates_args should contain the names of all args that the "
+                f"custom op mutates, or just the string 'unknown' if you don't know."
+            )
+    return_annotation = sig.return_annotation
+    if type(return_annotation) == str:
+        return_annotation = convert_type_string(return_annotation)
+    ret = parse_return(return_annotation, error_fn)
+    if op_name is not None:
+        return f"{op_name}({', '.join(params)}) -> {ret}"
+    return f"({', '.join(params)}) -> {ret}"
+
+def derived_types(
+    base_type, cpp_type, list_base, optional_base_list, optional_list_base
+):
+    result = [
+        (base_type, cpp_type),
+        (typing.Optional[base_type], f"{cpp_type}?"),
+    ]
+    if list_base:
+        result.append((typing.Sequence[base_type], f"{cpp_type}[]"))  # type: ignore[valid-type]
+    if optional_base_list:
+        result.append((typing.Sequence[typing.Optional[base_type]], f"{cpp_type}?[]"))  # type: ignore[valid-type]
+    if optional_list_base:
+        result.append((typing.Optional[typing.Sequence[base_type]], f"{cpp_type}[]?"))  # type: ignore[valid-type]
+    return result
+
+
+def get_supported_param_types():
+    data = [
+        # (python type, schema type, type[] variant, type?[] variant, type[]? variant
+        (torch.Tensor, "Tensor", True, True, False),
+        (int, "SymInt", True, False, True),
+        (float, "float", True, False, True),
+        (bool, "bool", True, False, True),
+        (str, "str", False, False, False),
+        (torch.types.Number, "Scalar", True, False, False),
+        (torch.dtype, "ScalarType", False, False, False),
+        (torch.device, "Device", False, False, False),
+    ]
+    result = []
+    for line in data:
+        result.extend(derived_types(*line))
+    return dict(result)
+
+
+SUPPORTED_RETURN_TYPES = {
+    torch.Tensor: "Tensor",
+    typing.List[torch.Tensor]: "Tensor[]",
+    int: "SymInt",
+    float: "float",
+    bool: "bool",
+    torch.types.Number: "Scalar",
+}
+
+
+def parse_return(annotation, error_fn):
+    origin = typing.get_origin(annotation)
+    if origin is not tuple:
+        if annotation not in SUPPORTED_RETURN_TYPES.keys():
+            error_fn(
+                f"Return has unsupported type {annotation}. "
+                f"The valid types are: {SUPPORTED_RETURN_TYPES}."
+            )
+        return SUPPORTED_RETURN_TYPES[annotation]
+
+    args = typing.get_args(annotation)
+    for arg in args:
+        if arg not in SUPPORTED_RETURN_TYPES:
+            error_fn(
+                f"Return has unsupported type {annotation}. "
+                f"The valid types are: {SUPPORTED_RETURN_TYPES}."
+            )
+
+    return "(" + ", ".join([SUPPORTED_RETURN_TYPES[arg] for arg in args]) + ")"
+
+
+SUPPORTED_PARAM_TYPES = get_supported_param_types()
+
+
+def report_error_callback(custom_op: typing.Any, key: str) -> None:
+    if key == "Undefined":
+        raise NotImplementedError(
+            f"{custom_op}: There were no Tensor inputs to this operator "
+            f"(e.g. you passed an empty list of Tensors). If your operator is a "
+            f"factory function (that is, it takes no Tensors and constructs "
+            f"a new one), then please use CustomOp.impl_factory to register "
+            f"an implementation for it"
+        )
+    if key == "Meta":
+        raise NotImplementedError(
+            f"{custom_op}: when running with device='Meta' tensors: there is no "
+            f"abstract impl registered for this CustomOp. Please register one via "
+            f"CustomOp.impl_abstract to get this CustomOp to work with Meta tensors"
+        )
+    if key in ("CPU", "CUDA"):
+        device = key.lower()
+        raise NotImplementedError(
+            f"{custom_op}: when running with device='{device}' tensors: there is no "
+            f"{device} impl registered for this CustomOp. Please register one via "
+            f"CustomOp.impl(device_type='{device}')"
+        )
+    raise NotImplementedError(
+        f"{custom_op}: No implementation for dispatch key {key}. It is likely "
+        f"that we have not added this functionality yet, please either open an "
+        f"issue or if you're feeling adventurous, use the low-level "
+        f"torch.library API"
+    )
+
+
+def custom_op_from_existing(op):
+    ns = op.namespace
+    lib = torch.library.Library(ns, "FRAGMENT")
+    name = op.name().split("::")[-1]
+    schema_str = str(op._schema)
+    # CustomOp expects the schema string without the namespace
+    schema_str = schema_str.split("::")[-1]
+    schema = FunctionSchema.parse(schema_str)
+    return CustomOp(lib, ns, schema, name, op, _private_access=True)
+
+
+def get_op(qualname):
+    def error_not_found():
+        raise ValueError(
+            f"Could not find the operator {qualname}. Please make sure you have "
+            f"already registered the operator and (if registered from C++) "
+            f"loaded it via torch.ops.load_library.")
+
+    ns, name = parse_qualname(qualname)
+    if not hasattr(torch.ops, ns):
+        error_not_found()
+    opnamespace = getattr(torch.ops, ns)
+    if not hasattr(opnamespace, name):
+        error_not_found()
+    packet = getattr(opnamespace, name)
+    if not hasattr(packet, 'default'):
+        error_not_found()
+    return packet.default
+
+
+def _find_custom_op(qualname, also_check_torch_library=False):
+    if qualname in global_registry:
+        return global_registry[qualname]
+    if not also_check_torch_library:
+        raise RuntimeError(
+            f"Could not find custom op \"{qualname}\". Did you register it via "
+            f"the torch._custom_ops API?")
+    overload = get_op(qualname)
+    result = custom_op_from_existing(overload)
+    return result
+
+
+def get_abstract_impl(qualname):
+    if qualname not in torch._custom_op.impl.global_registry:
+        return None
+    custom_op = torch._custom_op.impl.global_registry[qualname]
+    if custom_op is None:
+        return None
+    if not custom_op._has_impl("abstract"):
+        return None
+    return custom_op._get_impl("abstract").func
+
+
+def _custom_op_with_schema(qualname, schema):
+    ns, name = qualname.split("::")
+    schema_str = f"{name}{schema}"
+    function_schema = FunctionSchema.parse(schema_str)
+    validate_schema(function_schema)
+
+    lib = library.Library(ns, "FRAGMENT")
+    lib.define(schema_str)
+    ophandle = find_ophandle_or_throw(ns, function_schema.name)
+    result = CustomOp(lib, ns, function_schema, name, ophandle, _private_access=True)
+    result._register_autograd_kernel_indirection()
+
+    torch._C._dispatch_set_report_error_callback(
+        ophandle, functools.partial(report_error_callback, weakref.proxy(result))
+    )
+    return get_op(qualname)
diff --git a/vllm_musa/install.sh b/vllm_musa/install.sh
new file mode 100644
index 000000000..eaa344256
--- /dev/null
+++ b/vllm_musa/install.sh
@@ -0,0 +1,50 @@
+#!/bin/bash
+
+DEVELOP_MODE=0
+for arg in "$@"; do
+  if [ "$arg" == "-e" ]; then
+    DEVELOP_MODE=1
+    break
+  fi
+done
+
+SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+PROJECT_ROOT=$(dirname "$SCRIPT_DIR")
+
+VLLM_DIR="$PROJECT_ROOT/vllm"
+
+install_vllm_musa_command() {
+  if [ $DEVELOP_MODE -eq 1 ]; then
+    echo "📦 使用开发模式安装: pip install -e . -v"
+    pip install -e . -v
+  else
+    echo "📦 使用普通模式安装: pip install . -v"
+    pip install . -v
+  fi
+}
+
+
+install_vllm_musa_command
+# hack for ray==2.46.0
+RAY_PATH=$(python3 -c "import ray; print(ray.__file__)" | sed 's|/[^/]*$||')
+cp "./hack/compiled_dag_node.py" "$RAY_PATH/dag/compiled_dag_node.py"
+
+if command -v conda &> /dev/null; then
+    echo "检测到conda，使用conda安装gcc 12.1.0..."
+    conda install -c conda-forge gcc=12.1.0 -y
+else
+    echo "未检测到conda，尝试使用系统包管理器安装gcc/g++ 12..."
+    apt update
+    apt install -y gcc-12 g++-12
+fi
+
+# hack for triton_musa
+TRITON_PATH=$(python3 -c "import triton; print(triton.__file__)" | sed 's|/[^/]*$||')
+NV_TRITON_BACKEND="$TRITON_PATH/backends/nvidia/"
+if [ -d "$NV_TRITON_BACKEND" ]; then
+    rm -r "$NV_TRITON_BACKEND"
+else
+    echo "✅ 当前环境中没有 '$NV_TRITON_BACKEND' 目录，无需删除。"
+fi
+
+echo "Install vLLM_MUSA Done!"
diff --git a/vllm_musa/mypy.ini b/vllm_musa/mypy.ini
new file mode 100644
index 000000000..fe0fd6611
--- /dev/null
+++ b/vllm_musa/mypy.ini
@@ -0,0 +1,10 @@
+[mypy]
+; warn_return_any = True
+warn_unused_configs = True
+
+; Suppress all missing import errors from torch_npu for mypy.
+[mypy-torch_npu.*]
+ignore_missing_imports = True
+
+[mypy-transformers.*]
+ignore_missing_imports = True
diff --git a/vllm_musa/packages.txt b/vllm_musa/packages.txt
new file mode 100644
index 000000000..84955cda6
--- /dev/null
+++ b/vllm_musa/packages.txt
@@ -0,0 +1,3 @@
+git
+vim
+wget
diff --git a/vllm_musa/pyproject.toml.bak b/vllm_musa/pyproject.toml.bak
new file mode 100644
index 000000000..ecfadeb6c
--- /dev/null
+++ b/vllm_musa/pyproject.toml.bak
@@ -0,0 +1,10 @@
+[build-system]
+# Should be mirrored in requirements.txt
+requires = [
+    "setuptools>=64",
+    "setuptools-scm>=8",
+    "decorator",
+    "pyyaml",
+    "scipy",
+]
+build-backend = "setuptools.build_meta"
diff --git a/vllm_musa/pytest.ini b/vllm_musa/pytest.ini
new file mode 100644
index 000000000..d00c3c5b6
--- /dev/null
+++ b/vllm_musa/pytest.ini
@@ -0,0 +1,58 @@
+[pytest]
+minversion = 6.0
+norecursedirs =
+    vllm-empty/tests/prefix_caching
+    vllm-empty/tests/weight_loading
+    vllm-empty/tests/samplers
+    vllm-empty/tests/kernels
+    vllm-empty/tests/quantization
+    vllm-empty/tests/tool_use
+    vllm-empty/tests/runai_model_streamer_test
+    vllm-empty/tests/kv_transfer
+    vllm-empty/tests/plugins
+    vllm-empty/tests/plugins_tests
+    vllm-empty/tests/prompt_adapter
+    vllm-empty/tests/compile
+    vllm-empty/tests/lora
+    vllm-empty/tests/models
+    vllm-empty/tests/mistral_tool_use
+    vllm-empty/tests/standalone_tests
+    vllm-empty/tests/async_engine
+    vllm-empty/tests/mq_llm_engine
+    vllm-empty/tests/tokenization
+    vllm-empty/tests/core
+    vllm-empty/tests/tracing
+    vllm-empty/tests/engine
+    vllm-empty/tests/tensorizer_loader
+    vllm-empty/tests/entrypoints
+    vllm-empty/tests/model_executor
+    vllm-empty/tests/encoder_decoder
+    vllm-empty/tests/v1
+    vllm-empty/tests/spec_decode
+    vllm-empty/tests/multi_step
+    vllm-empty/tests/vllm_test_utils
+    vllm-empty/tests/tpu
+    vllm-empty/tests/distributed
+    vllm-empty/tests/basic_correctness
+    vllm-empty/tests/worker
+    vllm-empty/tests/metrics
+    vllm-empty/tests/neuron
+
+addopts = --ignore=vllm-empty/tests/test_utils.py
+          --ignore=vllm-empty/tests/test_config.py
+          --ignore=vllm-empty/tests/test_scalartype.py
+          --ignore=vllm-empty/tests/test_embedded_commit.py
+          --ignore=vllm-empty/tests/test_inputs.py
+          --ignore=vllm-empty/tests/test_sharded_state_loader.py
+          --ignore=vllm-empty/tests/test_logger.py
+          --ignore=vllm-empty/tests/test_logits_processor.py
+          --ignore=vllm-empty/tests/test_regression.py
+          --ignore=vllm-empty/tests/models/decoder_only/language/test_aqlm.py
+          --ignore=vllm-empty/tests/models/decoder_only/language/test_gptq_marlin.py
+          --ignore=vllm-empty/tests/models/decoder_only/language/test_gptq_marlin_24.py
+          --ignore=vllm-empty/tests/multimodal/test_processing.py
+          --ignore=vllm-empty/tests/multimodal/test_processor_kwargs.py
+          --ignore=vllm-empty/tests/multimodal/test_utils.py
+
+testpaths =
+    vllm-empty/tests
diff --git a/vllm_musa/requirements-dev.txt b/vllm_musa/requirements-dev.txt
new file mode 100644
index 000000000..9c8109487
--- /dev/null
+++ b/vllm_musa/requirements-dev.txt
@@ -0,0 +1,4 @@
+-r requirements-lint.txt
+modelscope
+pytest >= 6.0
+pytest-asyncio
diff --git a/vllm_musa/requirements-lint.txt b/vllm_musa/requirements-lint.txt
new file mode 100644
index 000000000..711bb50a0
--- /dev/null
+++ b/vllm_musa/requirements-lint.txt
@@ -0,0 +1,15 @@
+# formatting
+yapf==0.32.0
+toml==0.10.2
+tomli==2.0.2
+ruff==0.6.5
+codespell==2.3.0
+isort==5.13.2
+clang-format==18.1.5
+sphinx-lint==1.0.0
+
+# type checking
+mypy==1.11.1
+types-PyYAML
+types-requests
+types-setuptools
diff --git a/vllm_musa/requirements.txt b/vllm_musa/requirements.txt
new file mode 100644
index 000000000..f60d064b0
--- /dev/null
+++ b/vllm_musa/requirements.txt
@@ -0,0 +1,10 @@
+decorator
+pyyaml
+scipy
+setuptools
+setuptools-scm
+pandas
+ray==2.46.0
+pyarrow==20.0.0
+blobfile
+numba
\ No newline at end of file
diff --git a/vllm_musa/run.sh b/vllm_musa/run.sh
new file mode 100644
index 000000000..a3fcba092
--- /dev/null
+++ b/vllm_musa/run.sh
@@ -0,0 +1,132 @@
+#!/bin/bash
+
+if [[ $1 == "--help" ]]; then
+    echo "Usage: run.sh [TP_SIZE] [PP_SIZE] [MODEL_PATH] [HOSTFILE] [VLLM_PP_LAYER_PARTITION]"
+    echo ""
+    echo "Parameters:"
+    echo "  TP_SIZE                      Number of Tensor Parallelism"
+    echo "  PP_SIZE                      Number of Pipeline Parallelism"
+    echo "  MODEL_PATH                   Path to the model"
+    echo "  HOSTFILE                     Host file for distributed training"
+    echo "  VLLM_PP_LAYER_PARTITION      Optional partition scheme (comma-separated values); omit to skip"
+    echo ""
+    echo "Example:"
+    echo "  ./run.sh 2 4 /path/to/model /path/to/hostfile 13,12,12,12,12"
+    exit 0
+fi
+
+set -u
+TP_SIZE=$1
+PP_SIZE=$2
+MODEL_PATH=$3
+HOSTFILE=$4
+VLLM_PP_LAYER_PARTITION="${5:-}"
+set +u
+
+MODEL_NAME=deepseek
+MAX_MODEL_LEN=1024
+GPU_MEMORY_UTILIZATION=0.70
+WORLD_SIZE=$(($PP_SIZE * $TP_SIZE))
+SSH_PORT=62275
+RAY_PORT=63794
+
+
+# For S5000 RoCE
+#   MUSA_BLOCK_SCHEDULE_MODE=1
+#   MCCL_IB_GID_INDEX=3
+env_array=(
+    RAY_DEDUP_LOGS=0
+    MUSA_LAUNCH_BLOCKING=0
+    MUSA_USERQ=1
+    VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE=shm
+    RAY_CGRAPH_get_timeout=3000
+    VLLM_USE_V1=1
+    VLLM_TORCH_PROFILER_DIR=/home/dist/qzg/
+    MUSA_PRINT_ENV=1
+    MUSA_HOME="/usr/local/musa"
+    MTHREADS_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+    TRITON_CACHE_DIR="/tmp/triton"
+    LIBRARY_PATH="/opt/intel/oneapi/mkl/lib/intel64:${LIBRARY_PATH}"
+    LD_LIBRARY_PATH="/usr/lib/x86_64-linux-gnu:/usr/local/musa/lib"
+    VLLM_NCCL_SO_PATH="/usr/local/musa/lib/libmccl.so.2"
+    MCCL_PROTOS=2
+)
+    # MCCL_IB_HCA='=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_6:1,mlx5_7:1,mlx5_8:1,mlx5_9:1'
+    # MCCL_ALGOS=1
+    # MCCL_BUFFSIZE=20480000
+    # MCCL_IB_GID_INDEX=3
+    # MUSA_BLOCK_SCHEDULE_MODE=1
+
+if [[ -n "$VLLM_PP_LAYER_PARTITION" ]]; then
+    env_array+=(VLLM_PP_LAYER_PARTITION="$VLLM_PP_LAYER_PARTITION")
+fi
+
+for item in "${env_array[@]}"; do
+    echo "export $item"
+    eval "export $item"
+done
+
+pkill -f /opt/conda/envs/py310/bin/python3
+ray stop
+rm -rf ${TRITON_CACHE_DIR}/*
+
+CURRENT_TIME=$(date "+%Y-%m-%d_%H:%M:%S")
+echo $CURRENT_TIME
+mkdir -p ./output/$CURRENT_TIME
+
+set -u
+  WORK_HOME="$PWD"
+  EXPNAME="${MODEL_NAME}_pp${PP_SIZE}_tp${TP_SIZE}_gpus${WORLD_SIZE}"
+  LOG_FILE=$WORK_HOME/output/$CURRENT_TIME/$EXPNAME.log
+set +u
+
+hostlist=$(grep -v '^#\|^$' $HOSTFILE | awk '{print $1}' | xargs)
+
+first_host=true
+first_host_ip=127.0.0.1
+for host in ${hostlist[@]}; do
+  echo ray start $host
+  ((COUNT++))
+  if $first_host; then
+    first_host=false
+    first_host_ip=$host
+    ssh -p $SSH_PORT $host "${env_array[@]} ray start --head --port=${RAY_PORT} --dashboard-host='0.0.0.0' --num-gpus 8"
+    sleep 3s
+  else
+    ssh -p $SSH_PORT $host "${env_array[@]} ray start --address ${first_host_ip}:${RAY_PORT} --num-gpus 8"
+  fi
+done
+
+ray status
+
+# if $ENABLE_GRAPH; then
+# vllm serve $MODEL_PATH  \
+#     --trust-remote-code \
+#     --max_model_len $MAX_MODEL_LEN \
+#     --num-gpu-blocks-override $MAX_MODEL_LEN \
+#     --gpu-memory-utilization $GPU_MEMORY_UTILIZATION  \
+#     --served-model-name $MODEL_NAME \
+#     --distributed-executor-backend ray \
+#     --port 8000 \
+#     --block_size 64 \
+#     --max-num-seqs $MAX_SEQUENCES \
+#     --compilation-config '{"cudagraph_capture_sizes": [1], "simple_cuda_graph": true}' \
+#     -tp $TP_SIZE \
+#     -pp $PP_SIZE 2>&1 | tee -a $LOG_FILE
+# else
+# vllm serve $MODEL_PATH  \
+#     --trust-remote-code \
+#     --max_model_len $MAX_MODEL_LEN \
+#     --num-gpu-blocks-override $MAX_MODEL_LEN \
+#     --gpu-memory-utilization $GPU_MEMORY_UTILIZATION  \
+#     --served-model-name $MODEL_NAME \
+#     --distributed-executor-backend ray \
+#     --port 8000 \
+#     --block_size 64 \
+#     --max-num-seqs $MAX_SEQUENCES \
+#     --enforce-eager \
+#     --compilation-config '{"cudagraph_capture_sizes": [1], "simple_cuda_graph": false}' \
+#     --speculative-config '{"method": "deepseek_mtp", "num_speculative_tokens": 1}' \
+#     -tp $TP_SIZE \
+#     -pp $PP_SIZE 2>&1 | tee -a $LOG_FILE
+# fi
\ No newline at end of file
diff --git a/vllm_musa/setup.py b/vllm_musa/setup.py
new file mode 100644
index 000000000..d9db28c59
--- /dev/null
+++ b/vllm_musa/setup.py
@@ -0,0 +1,187 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+# Adapted from https://github.com/vllm-project/vllm/blob/main/setup.py
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+from typing import List
+
+from setuptools import find_packages, setup
+from setuptools_scm import get_version
+
+ROOT_DIR = os.path.dirname(__file__)
+try:
+    VERSION = get_version(write_to="vllm_musa/_version.py")
+except LookupError:
+    # The checkout action in github action CI does not checkout the tag. It
+    # only checks out the commit. In this case, we set a dummy version.
+    VERSION = "0.0.0"
+
+
+def get_path(*filepath) -> str:
+    return os.path.join(ROOT_DIR, *filepath)
+
+
+def read_readme() -> str:
+    """Read the README file if present."""
+    p = get_path("README.md")
+    if os.path.isfile(p):
+        with open(get_path("README.md"), encoding="utf-8") as f:
+            return f.read()
+    else:
+        return ""
+
+
+def get_requirements() -> List[str]:
+    """Get Python package dependencies from requirements.txt."""
+
+    def _read_requirements(filename: str) -> List[str]:
+        with open(get_path(filename)) as f:
+            requirements = f.read().strip().split("\n")
+        resolved_requirements = []
+        for line in requirements:
+            if line.startswith("-r "):
+                resolved_requirements += _read_requirements(line.split()[1])
+            elif line.startswith("--"):
+                continue
+            else:
+                resolved_requirements.append(line)
+        return resolved_requirements
+
+    try:
+        requirements = _read_requirements("requirements.txt")
+    except ValueError:
+        print("Failed to read requirements.txt in vllm_musa.")
+    return requirements
+
+ext_modules = []
+cmdclass = {}
+
+import torch
+import torch_musa
+import torch_musa.utils.musa_extension as torch_musa_ext
+from torch_musa.utils.musa_extension import MUSAExtension, BuildExtension
+
+class NinjaBuildExtension(BuildExtension):
+    def __init__(self, *args, **kwargs) -> None:
+        # do not override env MAX_JOBS if already exists
+        if not os.environ.get("MAX_JOBS"):
+            max_num_jobs_cores = max(1, os.cpu_count())
+            os.environ["MAX_JOBS"] = str(max_num_jobs_cores)
+
+        super().__init__(*args, **kwargs)
+
+
+cmdclass["build_ext"] = NinjaBuildExtension
+cxx_flags = [
+    "-O3",
+    "-Wno-switch-bool",
+]
+mcc_flags = [
+    "-Od3",
+    "-O2",
+    "-DNDEBUG",
+    "-fno-strict-aliasing",
+    "-fmusa-flush-denormals-to-zero",
+    "-std=c++17",
+    "-mllvm",
+    "-mtgpu-enable-max-ilp-scheduling-strategy=0",
+    "-mllvm",
+    "-mtgpu-enchanced-minreg-schedule=1",
+    "-mllvm",
+    "-mtgpu-enable-cse=0"
+]
+kernel_sources = [
+    "csrc/ops.mu",
+    "csrc/layernorm_kernels.mu",
+    "csrc/rmsnorm_cast/rmsnorm_cast.mu",
+    "csrc/pos_encoding_kernels.mu",
+    "csrc/cache_kernels.mu",
+    "csrc/attention/paged_attention_v1.mu",
+    "csrc/attention/paged_attention_v2.mu",
+    "csrc/attention/merge_attn_states.mu",
+    "csrc/moe/moe_align_sum_kernels.mu",
+    "csrc/moe/topk_softmax_kernels.mu",
+    "csrc/moe/mutlass/moe_mutlass.mu",
+    "csrc/moe/moe_gemv_swiglu.mu",
+    "csrc/moe/deep_gemm/deep_gemm.mu",
+    "csrc/activation_kernels.mu",
+    "csrc/element_wise/ternary.mu",
+    "csrc/linear/gemv_swiglu.mu",
+    "csrc/linear/gemv_swiglu_w4a16.mu",
+    "csrc/linear/gemv_w4a16.mu",
+    "csrc/linear/mp22_w4a16_groupwise_gemm/mixed_input_gemm.mu",
+    "csrc/custom_all_reduce.mu",
+    "csrc/flash_attn/flash_attn_fwd.mu",
+    "csrc/flash_attn/v1_flash_attn_fwd.mu",
+    "csrc/flash_attn/varlen_fa_padding_kernel.mu",
+    "csrc/quantization/matmul_mudnn.cpp",
+    "csrc/quantization/fp8/common.mu",
+    "csrc/decode_mla/decode_mla.mu",
+    "csrc/quantization/fp8/per_token_group_quant_8bit.mu",
+    "csrc/pos_encoding_contiguous_kernels.mu",
+    "csrc/layernorm_kernels_musa.mu",
+    "csrc/moe/mp22_fused_grouped_topk_softmax.mu",
+    "csrc/moe/mp31_fused_grouped_topk_softmax.mu",
+]
+ext_modules = [
+    MUSAExtension(
+        name="vllm_musa._kernels",
+        sources=kernel_sources,
+        include_dirs=[get_path("csrc/flash_attn", "include"), get_path("csrs/moe/deep_gemm")],
+        extra_compile_args={
+            "cxx": cxx_flags,
+            "mcc": mcc_flags,
+        },
+        library_dirs=[get_path("csrc/moe/deep_gemm")],
+        libraries=["musa_asm_lib_gen"]
+    )
+]
+
+setup(
+    name='vllm_musa',
+    # Follow:
+    # https://packaging.python.org/en/latest/specifications/version-specifiers
+    version=VERSION,
+    author="vLLM-MUSA team",
+    license="Apache 2.0",
+    description=("vLLM MUSA backend plugin"),
+    long_description=read_readme(),
+    long_description_content_type="text/markdown",
+    url="https://sh-code.mthreads.com/mcc-ai-serving/vllm_musa",
+    project_urls={
+        "Homepage": "https://sh-code.mthreads.com/mcc-ai-serving/vllm_musa",
+    },
+    classifiers=[
+        "Programming Language :: Python :: 3.9",
+        "Programming Language :: Python :: 3.10",
+        "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
+        "License :: OSI Approved :: Apache Software License",
+        "Intended Audience :: Developers",
+        "Intended Audience :: Information Technology",
+        "Intended Audience :: Science/Research",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+        "Topic :: Scientific/Engineering :: Information Analysis",
+    ],
+    packages=find_packages(exclude=("docs", "examples", "tests*")),
+    ext_modules=ext_modules,
+    cmdclass=cmdclass,
+    python_requires=">=3.9",
+    install_requires=get_requirements(),
+    extras_require={},
+    entry_points={'vllm.platform_plugins': ["musa = vllm_musa:register"]})
diff --git a/vllm_musa/vllm_musa/__init__.py b/vllm_musa/vllm_musa/__init__.py
new file mode 100644
index 000000000..d50dfc2d5
--- /dev/null
+++ b/vllm_musa/vllm_musa/__init__.py
@@ -0,0 +1,87 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+def patch_torch():
+    import os
+    import torch
+    import torch_musa
+
+    torch.cuda.is_available = torch.musa.is_available
+    torch.cuda.current_device = torch.musa.current_device
+    torch.cuda.device_count = torch.musa.device_count
+    torch.cuda.set_device = torch.musa.set_device
+    torch.cuda.DoubleTensor = torch.musa.DoubleTensor
+    torch.cuda.FloatTensor = torch.musa.FloatTensor
+    torch.cuda.LongTensor = torch.musa.LongTensor
+    torch.cuda.HalfTensor = torch.musa.HalfTensor
+    torch.cuda.BFloat16Tensor = torch.musa.BFloat16Tensor
+    torch.cuda.IntTensor = torch.musa.IntTensor
+    torch.cuda.synchronize = torch.musa.synchronize
+    torch.cuda.get_rng_state = torch.musa.get_rng_state
+    torch.cuda.set_rng_state = torch.musa.set_rng_state
+    torch.cuda.synchronize = torch.musa.synchronize
+    torch.cuda.empty_cache = torch.musa.empty_cache
+    torch.Tensor.cuda = torch.Tensor.musa
+    torch.cuda.manual_seed = torch.musa.manual_seed
+    torch.cuda.Event = torch.musa.Event
+    torch.cuda.Stream = torch.musa.Stream
+    torch.cuda.get_device_properties = torch.musa.get_device_properties
+    torch.cuda.memory_allocated = torch.musa.memory_allocated
+    torch.cuda.max_memory_allocated = torch.musa.memory_allocated
+    torch.cuda.memory_reserved = torch.musa.memory_reserved
+    torch.cuda.max_memory_reserved = torch.musa.max_memory_reserved
+    torch.cuda.mem_get_info = torch.musa.mem_get_info
+    torch.cuda.reset_peak_memory_stats = torch.musa.reset_peak_memory_stats
+    torch.cuda.CUDAGraph = torch.musa.MUSAGraph
+    torch.cuda.graph_pool_handle = torch.musa.graph_pool_handle
+    torch.cuda.graph = torch.musa.graph
+    torch.cuda.current_stream = torch.musa.current_stream
+    torch.cuda.set_stream = torch.musa.set_stream
+    torch.Tag.needs_fixed_stride_order = None
+
+    torch.ops.vllm._apply_gguf_embedding = lambda : ""
+    torch.ops.vllm._fused_mul_mat_gguf = lambda : ""
+    torch.ops.vllm._fused_moe_gguf = lambda : ""
+    torch.ops.vllm.apply_bnb_4bit = lambda : ""
+
+    os.environ["NVTE_TORCH_COMPILE"] = "0"
+    os.environ["TORCHDYNAMO_DISABLE"] = "1"
+
+    def noop(*args, **kwargs):
+        return lambda func: func
+    torch.compile = noop
+    torch.jit.script = noop
+
+
+def register():
+    """Register the MTGPU platform.
+
+        To ensure that the module is correctly replaced, add it at the beginning
+        If replace module, add it to here, for example:
+        import vllm_musa.patch_module
+        and `vllm_musa/pathc_module.py` file defined like below:
+        ```python
+        import sys
+        import types
+        sys.modules[
+            'vllm.model_executor.layers.fused_moe.fused_moe'] = types.ModuleType('fused_moe_module')
+        ```
+    """
+
+    patch_torch()
+
+    return "vllm_musa.platforms.musa.MusaPlatform"
diff --git a/vllm_musa/vllm_musa/_musa_custom_ops.py b/vllm_musa/vllm_musa/_musa_custom_ops.py
new file mode 100644
index 000000000..4d16a4c40
--- /dev/null
+++ b/vllm_musa/vllm_musa/_musa_custom_ops.py
@@ -0,0 +1,801 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import List, Optional, Tuple
+
+import torch
+from vllm_musa.platforms.musa import on_ph1, on_qy2
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+try:
+    from . import _kernels
+except ImportError as e:
+    logger.warning("Failed to import _kernels from . with %r", e)
+
+
+# layer norm ops
+def rms_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor,
+             epsilon: float) -> None:
+    _kernels.rms_norm(out, input, weight, epsilon)
+
+
+def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
+                       weight: torch.Tensor, epsilon: float) -> None:
+    _kernels.fused_add_rms_norm(input, residual, weight, epsilon)
+
+def musa_fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
+                       weight: torch.Tensor, epsilon: float) -> None:
+    _kernels.musa_fused_add_rms_norm(input, residual, weight, epsilon)
+
+def rmsnorm_cast(
+        input: torch.Tensor, weight: torch.Tensor,
+        epsilon: float, quant_size: int
+):
+    M, N = input.shape
+    device = input.device
+    output_fp8 = torch.empty(M, N, device=device, dtype=torch.float8_e4m3fn)
+    inv_scale = torch.empty((M, (N + quant_size-1) // quant_size), device=device, dtype=torch.float32)
+    _kernels.rmsnorm_cast(output_fp8, inv_scale, input, weight, epsilon, quant_size)
+    return output_fp8, inv_scale
+
+def rmsnorm_add_cast(
+        input: torch.Tensor, resdual: torch.Tensor, weight: torch.Tensor,
+        epsilon: float, quant_size: int
+):
+    M, N = input.shape
+    device = input.device
+    output_fp8 = torch.empty(M, N, device=device, dtype=torch.float8_e4m3fn)
+    inv_scale = torch.empty((M, (N + quant_size-1) // quant_size), device=device, dtype=torch.float32)
+    _kernels.rmsnorm_add_cast(output_fp8, inv_scale, input, resdual, weight, epsilon, quant_size)
+    return output_fp8, inv_scale
+
+# pos encoding ops
+def rotary_embedding(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+) -> None:
+    _kernels.rotary_embedding(positions, query, key, head_size,
+                                  cos_sin_cache, is_neox)
+
+def batched_rotary_embedding(positions: torch.Tensor, query: torch.Tensor,
+                             key: torch.Tensor, head_size: int,
+                             cos_sin_cache: torch.Tensor, is_neox: bool,
+                             rot_dim: int,
+                             cos_sin_cache_offsets: torch.Tensor) -> None:
+    _kernels.batched_rotary_embedding(positions, query, key, head_size,
+                                          cos_sin_cache, is_neox, rot_dim,
+                                          cos_sin_cache_offsets)
+
+
+def rotary_embedding_contiguous(
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    head_size: int,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+) -> None:
+    _kernels.rotary_embedding_contiguous(positions, query, key, head_size,
+                                  cos_sin_cache, is_neox)
+
+
+def batched_rotary_embedding_contiguous(positions: torch.Tensor, query: torch.Tensor,
+                             key: torch.Tensor, head_size: int,
+                             cos_sin_cache: torch.Tensor, is_neox: bool,
+                             rot_dim: int,
+                             cos_sin_cache_offsets: torch.Tensor) -> None:
+    _kernels.batched_rotary_embedding_contiguous(positions, query, key, head_size,
+                                          cos_sin_cache, is_neox, rot_dim,
+                                          cos_sin_cache_offsets)
+
+# cache
+def reshape_and_cache(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+) -> None:
+    _kernels.reshape_and_cache(key, value, key_cache,
+                                value_cache, slot_mapping,
+                                kv_cache_dtype, k_scale, v_scale)
+
+def reshape_and_cache_flash(
+    key: torch.Tensor,
+    value: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+) -> None:
+    _kernels.reshape_and_cache_flash(key, value, key_cache,
+                                                   value_cache, slot_mapping,
+                                                   kv_cache_dtype, k_scale,
+                                                   v_scale)
+
+def copy_blocks(key_caches: List[torch.Tensor],
+                value_caches: List[torch.Tensor],
+                block_mapping: torch.Tensor) -> None:
+    _kernels.copy_blocks(key_caches, value_caches, block_mapping)
+
+
+def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
+                block_mapping: torch.Tensor) -> None:
+    _kernels.swap_blocks(src, dst, block_mapping)
+
+
+def copy_blocks_mla(kv_caches: List[torch.Tensor],
+                    block_mapping: torch.Tensor) -> None:
+    _kernels.copy_blocks_mla(kv_caches, block_mapping)
+
+
+def concat_and_cache_mla(
+    kv_c: torch.Tensor,
+    k_pe: torch.Tensor,
+    kv_cache: torch.Tensor,
+    slot_mapping: torch.Tensor,
+    kv_cache_dtype: str,
+    scale: torch.Tensor,
+) -> None:
+    _kernels.concat_and_cache_mla(kv_c, k_pe, kv_cache,
+                                                slot_mapping, kv_cache_dtype,
+                                                scale)
+
+def gather_cache(src_cache: torch.Tensor,
+                 dst: torch.Tensor,
+                 block_table: torch.Tensor,
+                 cu_seq_lens: torch.Tensor,
+                 batch_size: int,
+                 seq_starts: Optional[torch.Tensor] = None) -> None:
+    _kernels.gather_cache(src_cache, dst, block_table,
+                          cu_seq_lens, batch_size, seq_starts)
+
+
+# paged attention v1
+def paged_attention_v1(
+    out: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    _kernels.paged_attention_v1(
+        out, query, key_cache, value_cache, num_kv_heads, scale, block_tables,
+        seq_lens, block_size, max_seq_len, alibi_slopes, kv_cache_dtype,
+        k_scale, v_scale, tp_rank, blocksparse_local_blocks,
+        blocksparse_vert_stride, blocksparse_block_size,
+        blocksparse_head_sliding_step)
+
+
+def paged_attention_v2(
+    out: torch.Tensor,
+    exp_sum: torch.Tensor,
+    max_logits: torch.Tensor,
+    tmp_out: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    num_kv_heads: int,
+    scale: float,
+    block_tables: torch.Tensor,
+    seq_lens: torch.Tensor,
+    block_size: int,
+    max_seq_len: int,
+    alibi_slopes: Optional[torch.Tensor],
+    kv_cache_dtype: str,
+    k_scale: torch.Tensor,
+    v_scale: torch.Tensor,
+    tp_rank: int = 0,
+    blocksparse_local_blocks: int = 0,
+    blocksparse_vert_stride: int = 0,
+    blocksparse_block_size: int = 64,
+    blocksparse_head_sliding_step: int = 0,
+) -> None:
+    _kernels.paged_attention_v2(
+        out, exp_sum, max_logits, tmp_out, query, key_cache, value_cache,
+        num_kv_heads, scale, block_tables, seq_lens, block_size, max_seq_len,
+        alibi_slopes, kv_cache_dtype, k_scale, v_scale, tp_rank,
+        blocksparse_local_blocks, blocksparse_vert_stride,
+        blocksparse_block_size, blocksparse_head_sliding_step)
+
+# merge attn states ops
+def merge_attn_states(output: torch.Tensor,
+                      prefix_output: torch.Tensor,
+                      prefix_lse: torch.Tensor,
+                      suffix_output: torch.Tensor,
+                      suffix_lse: torch.Tensor,
+                      output_lse: Optional[torch.Tensor] = None) -> None:
+    _kernels.merge_attn_states(output, output_lse, prefix_output,
+                                   prefix_lse, suffix_output, suffix_lse)
+
+# moe
+def moe_sum(input: torch.Tensor, output: torch.Tensor):
+    _kernels.moe_sum(input, output)
+
+
+def moe_align_block_size(topk_ids: torch.Tensor, num_experts: int,
+                         block_size: int, sorted_token_ids: torch.Tensor,
+                         experts_ids: torch.Tensor,
+                         num_tokens_post_pad: torch.Tensor) -> None:
+    _kernels.moe_align_block_size(topk_ids, num_experts, block_size,
+                                          sorted_token_ids, experts_ids,
+                                          num_tokens_post_pad)
+
+
+def sgl_moe_align_block_size(topk_ids: torch.Tensor, num_experts: int,
+                             block_size: int, sorted_token_ids: torch.Tensor,
+                             experts_ids: torch.Tensor,
+                             num_tokens_post_pad: torch.Tensor) -> None:
+    _kernels.sgl_moe_align_block_size(topk_ids, num_experts,
+                                              block_size, sorted_token_ids,
+                                              experts_ids, num_tokens_post_pad)
+
+
+def topk_softmax(topk_weights: torch.Tensor, topk_ids: torch.Tensor,
+                 token_expert_indicies: torch.Tensor,
+                 gating_output: float) -> None:
+    _kernels.topk_softmax(topk_weights, topk_ids,
+                                  token_expert_indicies, gating_output)
+
+
+def fused_grouped_topk(
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: int,
+    topk_group: int,
+    e_score_correction_bias: torch.Tensor):
+    if on_qy2():
+        num_tokens, _ = gating_output.shape
+        topk_weights = torch.empty((num_tokens, topk), device=gating_output.device, dtype=gating_output.dtype)
+        topk_ids = torch.empty((num_tokens, topk), device=gating_output.device, dtype=torch.int32)
+        _kernels.mp22_fused_grouped_topk(topk_weights, topk_ids, gating_output, e_score_correction_bias, num_expert_group, topk_group, topk, renormalize)
+    elif on_ph1():
+        # topk_weights in FP32
+        # topk_ids in Int32
+        topk_weights, topk_ids = _kernels.mp31_fused_grouped_topk(gating_output, e_score_correction_bias, num_expert_group, topk_group, topk, renormalize)
+    else:
+        raise ValueError("Fused grouped topk is not supported on current mtgpu arch.")
+
+    return topk_weights, topk_ids
+
+def fused_moe_mutlass(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        C: torch.Tensor,
+        A_scale: torch.Tensor,
+        B_scale: torch.Tensor, # must be fp32
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        sorted_token_ids: torch.Tensor,
+        expert_ids: torch.Tensor,
+        num_tokens_post_padded: torch.Tensor,
+        mul_routed_weight: bool,
+        topk: int,
+        use_int4_w4a16: bool,
+        BLOCK_SIZE_M: int,
+        BLOCK_SIZE_N: int,
+        BLOCK_SIZE_K: int,
+        GROUP_SIZE_M: int,
+        block_shape: Optional[List[int]] = None,
+    ):
+    # group_size = -1 when w16a16 and per-channel
+    group_size = block_shape[1] if block_shape else -1
+    _kernels.fused_moe_mutlass(
+        A,
+        B,
+        C,
+        A_scale,
+        B_scale,
+        topk_weights,
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        mul_routed_weight,
+        topk,
+        use_int4_w4a16,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+        BLOCK_SIZE_K,
+        GROUP_SIZE_M,
+        group_size,
+    )
+
+def fused_moe_gemv(
+        A: torch.Tensor,
+        B: torch.Tensor,
+        C: torch.Tensor,
+        A_scale: torch.Tensor,
+        B_scale: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        sorted_token_ids: torch.Tensor,
+        expert_ids: torch.Tensor,
+        num_tokens_post_padded: torch.Tensor,
+        mul_routed_weight: bool,
+        topk: int,
+        use_int4_w4a16: bool,
+        use_swigelu: bool,
+):
+    _kernels._fused_moe_gemv(
+        A,
+        B,
+        C,
+        A_scale,
+        B_scale,
+        topk_weights,
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        mul_routed_weight,
+        topk,
+        use_int4_w4a16,
+        use_swigelu,
+    )
+
+# activation
+def silu_and_mul(out: torch.Tensor, input: torch.Tensor) -> None:
+    _kernels.silu_and_mul(out, input)
+
+
+def fused_mul_add(
+    self: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    scale: Optional[float],
+    accurate: bool = True):
+    # if accurate == False, then we call inplace op: bias += (self * scale)
+    if not accurate:
+        bias.add_(self, alpha=scale)
+        return bias
+
+    # otherwise, we call custom outplace op, act: output = self * scale + bias
+    output = torch.empty_like(self)
+    _kernels._fused_mul_add(output, self, bias, scale)
+
+    return output
+
+
+def weak_ref_tensor(tensor: torch.Tensor) -> torch.Tensor:
+    return _kernels.weak_ref_tensor(tensor)
+
+
+def gemm_swiglu(
+        output: torch.Tensor,
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        alpha=1.0):
+    _kernels._gemm_swiglu(output, input, weight, alpha)
+
+
+def flash_attn_fwd(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        o: torch.Tensor,
+        cu_q_seqlen: torch.Tensor,
+        cu_kv_seqlen: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_kv: int,
+        causal: bool = True,
+        varlen: bool = True,
+        return_softmax_lse: bool = False
+):
+    lse = None
+    if return_softmax_lse:
+        lse = torch.empty((q.shape[1], q.shape[0]), device=q.device, dtype=torch.float32)
+    _kernels.flash_attn_fwd(q, k, v, o, lse, cu_q_seqlen, cu_kv_seqlen, max_seqlen_q, max_seqlen_kv, causal, varlen, return_softmax_lse)
+    return o, lse if return_softmax_lse else o
+
+
+def flash_attn_varlen_func(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        o: torch.Tensor,
+        cu_q_seqlen: torch.Tensor,
+        cu_kv_seqlen: torch.Tensor,
+        block_table: torch.Tensor,
+        sm_scale: float,
+        max_seqlen_q: int,
+        max_seqlen_kv: int,
+        causal: bool = True,
+        return_softmax_lse: bool = False
+):
+    lse = None
+    if return_softmax_lse:
+        s, h, d = q.shape
+        lse = torch.empty((h, s), device=q.device, dtype=torch.float32)
+    _kernels.mp22_flash_attn_varlen_func(q, k, v, o, lse, cu_q_seqlen, cu_kv_seqlen, block_table,  sm_scale, max_seqlen_q, max_seqlen_kv, causal, return_softmax_lse)
+    return o, lse if return_softmax_lse else o
+
+
+def gemm_swiglu_w4a16(output: torch.Tensor,
+                input: torch.Tensor,
+                weight: torch.Tensor,
+                scale: torch.Tensor,
+                alpha=1.0):
+    return _kernels._gemm_swiglu_w4a16(output, input, weight, scale, alpha)
+
+
+def gemm_w4a16(
+                input: torch.Tensor,
+                weight: torch.Tensor,
+                scale: torch.Tensor,
+                alpha=1.0):
+    output = torch.empty(input.shape[0], weight.shape[1], device=input.device, dtype=input.dtype)
+    _kernels._gemm_w4a16(output, input, weight, scale, alpha)
+    return output
+
+
+def batch_gemv(input: torch.Tensor,
+               weight: torch.Tensor,
+               alpha=1.0):
+    output = torch.empty(input.shape[0], weight.shape[1], device=input.device, dtype=input.dtype)
+    _kernels.batch_gemv(output, input, weight, alpha)
+    return output
+
+def fused_gemv(x: torch.Tensor,
+               qweight: torch.Tensor,
+               x_scales:  Optional[torch.Tensor] = None,
+               qweight_sacles:  Optional[torch.Tensor] = None,
+               use_swigelu: bool = False,
+               use_rms_norm: bool = False,
+               gamma: Optional[torch.Tensor] = None,
+               eps: float = 1e-6):
+    use_int4_w4a16 = False
+    out_shape = x.shape[:-1] + (qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, )
+    assert not (use_swigelu and use_rms_norm), "gemv only fused one activation (swigelu or rms_norm)!"
+
+    if use_rms_norm:
+        if gamma is None:
+            assert False, "rms_norm gamm is None!"
+
+    # fp8 grouped matmul
+    if qweight.dtype == torch.float8_e4m3fn:
+        # x: [m, k]
+        # qweight: [n, k]
+        # x_scales: [m, k / 128]
+        # qweight: [n / 128, k / 128]
+        # assert x_scales is not None, "FP8 grouped matmul x sacles is None!"
+        assert qweight.dtype == torch.float8_e4m3fn, "FP8 grouped matmul weight only support float8_e4m3fn!"
+        assert qweight_sacles is not None, "FP8 grouped matmul weight sacles is None!"
+        output = torch.empty(
+            out_shape,
+            device=x.device,
+            dtype=torch.bfloat16
+        )
+        _kernels._fused_gemv(x,
+                             qweight,
+                             output,
+                             x_scales,
+                             qweight_sacles,
+                             use_int4_w4a16,
+                             use_swigelu,
+                             use_rms_norm,
+                             gamma,
+                             eps)
+        return output
+    # w4a16 gemv
+    elif qweight_sacles is not None:
+        # qweight: [out, in/8]
+        # sacles: [out, in / group_size]
+        assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, "W4A16 gemv only support bfloat16 or float16!"
+        use_int4_w4a16 = True
+        out_shape = x.shape[:-1] + (qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, )
+        output = torch.empty(
+            out_shape,
+            device=x.device,
+            dtype=x.dtype
+        )
+        _kernels._fused_gemv(x,
+                             qweight,
+                             output,
+                             None,
+                             qweight_sacles,
+                             use_int4_w4a16,
+                             use_swigelu,
+                             use_rms_norm,
+                             gamma,
+                             eps)
+        return output
+    # general gemv
+    else:
+        output = torch.empty(
+            out_shape,
+            device=x.device,
+            dtype=x.dtype
+        )
+        _kernels._fused_gemv(x,
+                             qweight,
+                             output,
+                             None,
+                             None,
+                             use_int4_w4a16,
+                             use_swigelu,
+                             use_rms_norm,
+                             gamma,
+                             eps)
+        return output
+
+
+def group_scale_gemm(x: torch.Tensor,
+                     qweight: torch.Tensor,
+                     sacles: torch.Tensor):
+    # qweight: [out, in/8]
+    # sacles: [out, in / group_size]
+    out_shape = x.shape[:-1] + (qweight.shape[0], )
+    output = torch.empty(
+        out_shape,
+        device=x.device,
+        dtype=x.dtype
+    )
+    _kernels._group_scale_gemm(x, qweight, output, sacles)
+
+    return output
+
+
+# custom ar(all-reduce)
+def init_custom_ar(ipc_tensors: List[torch.Tensor], rank_data: torch.Tensor,
+                   rank: int, full_nvlink: bool) -> int:
+    return _kernels.init_custom_ar(ipc_tensors, rank_data, rank,
+                                                 full_nvlink)
+
+
+def all_reduce(fa: int, inp: torch.Tensor, out: torch.Tensor, reg_buffer: int,
+               reg_buffer_sz_bytes: int) -> None:
+    _kernels.all_reduce(fa, inp, out, reg_buffer,
+                                      reg_buffer_sz_bytes)
+
+
+def dispose(fa: int) -> None:
+    _kernels.dispose(fa)
+
+
+def meta_size() -> int:
+    return _kernels.meta_size()
+
+
+def register_buffer(fa: int, ipc_tensors: List[int]) -> None:
+    return _kernels.register_buffer(fa, ipc_tensors)
+
+
+def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
+    return _kernels.get_graph_buffer_ipc_meta(fa)
+
+
+def register_graph_buffers(fa: int, handles: List[List[int]],
+                           offsets: List[List[int]]) -> None:
+    _kernels.register_graph_buffers(fa, handles, offsets)
+
+
+def mudnn_w4a16_matmul(x: torch.Tensor, qweight: torch.Tensor,
+                       scales: torch.Tensor, zeros: torch.Tensor = None,
+                       bias: Optional[torch.Tensor] = None,
+                       split_k: bool = False):
+    out_shape = x.shape[:-1] + (qweight.shape[0], )
+    output = torch.empty(
+        out_shape,
+        device=x.device,
+        dtype=x.dtype
+    )
+    _kernels.mudnn_w4a16_matmul(
+        output.view(-1, out_shape[-1]),
+        x.view(-1, x.shape[-1]),
+        qweight.t(),
+        scales,
+        bias,
+        split_k
+    )
+
+    return output
+
+def mudnn_w8a8_scaled_mm(a: torch.Tensor,
+                    b: torch.Tensor,
+                    scale_a: torch.Tensor,
+                    scale_b: torch.Tensor,
+                    out_dtype: torch.dtype,
+                    bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """
+    `cutlass_scaled_mm` implements a fused version of
+        `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
+    where scale_a * a and scale_b * b are implemented using numpy-style
+    broadcasting.
+
+    In order to support blockwise scaling like found in DeepSeek V3 we also
+    support extended "group" broadcast rules. We extend the numpy-style
+    broadcasting rules with the following rule:
+        "if the extent of a dimension in the source shape is between 1 and
+        corresponding extent in the target shape we repeat each element along
+        that dimension  src_shape[dim] // target_shape[dim] times consecutively"
+    example if we have:
+            a = [[1, 2], and target_shape = (2, 4)
+                [3, 4]]
+    then we would expand a to:
+            a = [[1, 1, 2, 2],
+                [3, 3, 4, 4]]
+    currently we only support the case:
+        scale_a.shape * [1, 128] == a.shape
+        scale_b.shape * [128, 128] == b.shape
+    """
+    assert (b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0)
+    assert (out_dtype is torch.bfloat16 or out_dtype is torch.float16)
+    assert bias is None or bias.shape[0] == b.shape[
+        1] and bias.dtype == out_dtype
+
+    m = a.shape[0]
+    n = b.shape[0]
+
+    out = torch.empty((m, n), dtype=out_dtype, device=a.device)
+    _kernels.mudnn_w8a8_scaled_mm(out, a, b, scale_a, scale_b, bias)
+    return out
+
+def decode_mla(
+        q: torch.Tensor,
+        k_buffer: torch.Tensor,
+        v_buffer: torch.Tensor,
+        sm_scale: float,
+        num_splits_kv: int,
+        req_to_tokens: torch.Tensor,
+        b_seqlen: torch.Tensor,
+        att_out: torch.Tensor
+):
+    _kernels.decode_mla(q, k_buffer, v_buffer, sm_scale, num_splits_kv, req_to_tokens, b_seqlen, att_out)
+
+def mp31_decode_mha(
+        q: torch.Tensor,
+        k_buffer: torch.Tensor,
+        v_buffer: torch.Tensor,
+        sm_scale: float,
+        num_splits_kv: int,
+        req_to_tokens: torch.Tensor,
+        b_seqlen: torch.Tensor,
+        att_out: torch.Tensor
+):
+    _kernels.mp31_decode_mha(q, k_buffer, v_buffer, sm_scale, num_splits_kv, req_to_tokens, b_seqlen, att_out)
+
+def awq_dequantize(qweight: torch.Tensor, scales: torch.Tensor,
+                   zeros: torch.Tensor, split_k_iters: int, thx: int,
+                   thy: int) -> torch.Tensor:
+    from vllm.model_executor.layers.quantization.awq_triton import (
+        awq_dequantize_triton)
+    return awq_dequantize_triton(qweight, scales, zeros)
+
+def varlen_fa_seqlen_pad(
+    q_in,
+    k_in,
+    v_in,
+    q_out,
+    k_out,
+    v_out,
+    accum_lens_q,
+    accum_lens_k,
+    sum_seq,
+    max_seq,
+    batch_size):
+    _kernels.varlen_fa_seqlen_pad(
+                                q_in,
+                                k_in,
+                                v_in,
+                                q_out,
+                                k_out,
+                                v_out,
+                                accum_lens_q,
+                                accum_lens_k,
+                                sum_seq,
+                                max_seq,
+                                batch_size)
+
+def varlen_fa_seqlen_unpad(
+    attn_o,
+    out,
+    accum_lens,
+    sum_seq,
+    max_seq,
+    head_dim,
+    num_heads,
+    batch_size):
+    _kernels.varlen_fa_seqlen_unpad(
+        attn_o,
+        out,
+        accum_lens,
+        sum_seq,
+        max_seq,
+        head_dim,
+        num_heads,
+        batch_size)
+
+def varlen_fa_lse_seqlen_unpad(
+    lse,
+    out,
+    accum_lens,
+    sum_seq,
+    max_seq,
+    num_heads,
+    batch_size):
+    _kernels.varlen_fa_lse_seqlen_unpad(
+        lse,
+        out,
+        accum_lens,
+        sum_seq,
+        max_seq,
+        num_heads,
+        batch_size)
+
+# quantization
+def static_scaled_fp8_quant(out: torch.Tensor,
+                            input: torch.Tensor,
+                            scale: torch.Tensor):
+    return _kernels.static_scaled_fp8_quant(out, input, scale)
+
+def dynamic_scaled_fp8_quant(out: torch.Tensor,
+                             input: torch.Tensor,
+                             scale: torch.Tensor):
+    return _kernels.dynamic_scaled_fp8_quant(out, input, scale)
+
+def dynamic_per_token_scaled_fp8_quant(out: torch.Tensor,
+                                       input: torch.Tensor,
+                                       scale: torch.Tensor,
+                                       scale_ub: torch.Tensor):
+    return _kernels.dynamic_per_token_scaled_fp8_quant(out, input, scale, scale_ub)
+
+def musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
+    lhs: Tuple[torch.Tensor, torch.Tensor],
+    rhs: Tuple[torch.Tensor, torch.Tensor],
+    out: torch.Tensor,
+    m_indices: torch.Tensor,
+    block_m: int
+):
+    _kernels.musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(lhs, rhs, out, m_indices, block_m)
+
+def sgl_per_token_group_quant_fp8(
+    input: torch.Tensor,
+    output_q: torch.Tensor,
+    output_s: torch.Tensor,
+    group_size: int,
+    eps: float,
+    fp8_min: float,
+    fp8_max: float,
+    scale_ue8m0: bool,
+    ) -> None:
+    _kernels.sgl_per_token_group_quant_fp8(
+        input, output_q, output_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0
+    )
+
+# (B, N, P) -> (N, B, P)
+# decode_q_nope = decode_q_nope.transpose(0, 1)
+# Multiply (N, B, P) x (N, P, L) -> (N, B, L)
+# decode_ql_nope = torch.bmm(decode_q_nope, W_UK_T)
+# Convert from (N, B, L) to (B, N, L)
+# decode_ql_nope = decode_ql_nope.transpose(0, 1)
+def mla_mudnn_bmm(
+    c: torch.Tensor,    #(B, N, L)
+    a: torch.Tensor,    #(B, N, P)
+    b: torch.Tensor,):  #(N, P, L)
+    _kernels.mla_mudnn_bmm(c, a, b)
diff --git a/vllm_musa/vllm_musa/_version.py b/vllm_musa/vllm_musa/_version.py
new file mode 100644
index 000000000..530083781
--- /dev/null
+++ b/vllm_musa/vllm_musa/_version.py
@@ -0,0 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
+
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
+
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
+
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
+
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
+
+__version__ = version = '0.1.dev286+g9ec64ba.d20250815'
+__version_tuple__ = version_tuple = (0, 1, 'dev286', 'g9ec64ba.d20250815')
diff --git a/vllm_musa/vllm_musa/attention/__init__.py b/vllm_musa/vllm_musa/attention/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/attention/flash_attn.py b/vllm_musa/vllm_musa/attention/flash_attn.py
new file mode 100755
index 000000000..a4e6f9e92
--- /dev/null
+++ b/vllm_musa/vllm_musa/attention/flash_attn.py
@@ -0,0 +1,867 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Attention layer with FlashAttention."""
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, ClassVar, Optional
+
+import numpy as np
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
+                                              AttentionMetadata, AttentionType,
+                                              is_quantized_kv_cache)
+from vllm.attention.layer import Attention
+from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,
+                                           get_flash_attn_version,
+                                           is_flash_attn_varlen_func_available)
+
+if is_flash_attn_varlen_func_available():
+    from vllm.attention.utils.fa_utils import (flash_attn_varlen_func,
+                                               get_scheduler_metadata,
+                                               reshape_and_cache_flash)
+
+from vllm.config import VllmConfig, get_layers_from_vllm_config
+from vllm.logger import init_logger
+from vllm.utils import cdiv
+from vllm.v1.attention.backends.utils import (
+    AttentionMetadataBuilder, CommonAttentionMetadata, get_kv_cache_layout,
+    make_local_attention_virtual_batches)
+from vllm.v1.kv_cache_interface import AttentionSpec
+from vllm.v1.worker.block_table import BlockTable
+
+if TYPE_CHECKING:
+    from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+
+import vllm_musa._musa_custom_ops as ops
+
+logger = init_logger(__name__)
+
+# NOTE(woosuk): This is an arbitrary number. Tune it if needed.
+_DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH = 16
+
+
+class FlashAttentionBackend(AttentionBackend):
+
+    accept_output_buffer: bool = True
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        return [32, 64, 96, 128, 160, 192, 224, 256]
+
+    @classmethod
+    def validate_head_size(cls, head_size: int) -> None:
+        supported_head_sizes = cls.get_supported_head_sizes()
+        if head_size not in supported_head_sizes:
+            attn_type = cls.__name__.removesuffix("Backend")
+            raise ValueError(
+                f"Head size {head_size} is not supported by {attn_type}. "
+                f"Supported head sizes are: {supported_head_sizes}. "
+                "Set VLLM_ATTENTION_BACKEND=FLEX_ATTENTION to use "
+                "FlexAttention backend which supports all head sizes.")
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN_VLLM_V1"
+
+    @staticmethod
+    def get_impl_cls() -> type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> type["AttentionMetadata"]:
+        return FlashAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> type["FlashAttentionMetadataBuilder"]:
+        return FlashAttentionMetadataBuilder
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> tuple[int, ...]:
+        if block_size % 16 != 0:
+            raise ValueError("Block size must be a multiple of 16.")
+        return (2, num_blocks, block_size, num_kv_heads, head_size)
+
+    @staticmethod
+    def get_kv_cache_stride_order() -> tuple[int, ...]:
+        # `stride_order` indicates the permutation that gets
+        # us from `get_kv_cache_shape` to the actual memory layout we want.
+        cache_layout = get_kv_cache_layout()
+        if cache_layout == "NHD":
+            stride_order = (0, 1, 2, 3, 4)
+        elif cache_layout == "HND":
+            stride_order = (0, 1, 3, 2, 4)
+        else:
+            raise ValueError(f"Unknown cache layout format {cache_layout}.")
+        return stride_order
+
+
+@dataclass
+class FlashAttentionMetadata:
+    # NOTE(sang): Definition of context_len, query_len, and seq_len.
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seq_len ---------------------|
+    #                                   |-- query_len ---|
+
+    num_actual_tokens: int  # Number of tokens excluding padding.
+    max_query_len: int
+    query_start_loc: torch.Tensor
+    max_seq_len: int
+    seq_lens: torch.Tensor
+    block_table: torch.Tensor
+    slot_mapping: torch.Tensor
+
+    # For cascade attention.
+    use_cascade: bool
+    common_prefix_len: int
+    cu_prefix_query_lens: Optional[torch.Tensor]
+    prefix_kv_lens: Optional[torch.Tensor]
+    suffix_kv_lens: Optional[torch.Tensor]
+
+    # Optional aot scheduling
+    scheduler_metadata: Optional[torch.Tensor] = None
+    prefix_scheduler_metadata: Optional[torch.Tensor] = None
+    max_num_splits: int = 0
+
+    # for local attention
+    @dataclass
+    class LocalAttentionMetadata:
+        local_query_start_loc: torch.Tensor
+        local_seqused_k: torch.Tensor
+        local_block_table: torch.Tensor
+        local_max_query_len: int
+        local_max_seq_len: int
+        local_scheduler_metadata: Optional[torch.Tensor]
+
+    local_attn_metadata: Optional[LocalAttentionMetadata] = None
+
+
+def _get_sliding_window_configs(
+        vllm_config: VllmConfig) -> set[Optional[tuple[int, int]]]:
+    """Get the set of all sliding window configs used in the model."""
+    sliding_window_configs: set[Optional[tuple[int, int]]] = set()
+    layers = get_layers_from_vllm_config(vllm_config, Attention)
+    for layer in layers.values():
+        assert isinstance(layer.impl, FlashAttentionImpl)
+        sliding_window_configs.add(layer.impl.sliding_window)
+    return sliding_window_configs
+
+
+class FlashAttentionMetadataBuilder(
+        AttentionMetadataBuilder[FlashAttentionMetadata]):
+    full_cudagraph_supported: ClassVar[bool] = get_flash_attn_version() == 3
+
+    def __init__(self, runner: "GPUModelRunner", kv_cache_spec: AttentionSpec,
+                 block_table: BlockTable):
+        model_config = runner.model_config
+        compilation_config = runner.vllm_config.compilation_config
+
+        self.runner = runner
+        self.num_heads_q = model_config.get_num_attention_heads(
+            runner.parallel_config)
+        self.num_heads_kv = model_config.get_num_kv_heads(
+            runner.parallel_config)
+        self.headdim = model_config.get_head_size()
+        self.block_size = kv_cache_spec.block_size
+        self.kv_cache_spec = kv_cache_spec
+        self.block_table = block_table
+
+        self.max_num_splits = 0  # No upper bound on the number of splits.
+        self.aot_schedule = (get_flash_attn_version() == 3)
+        self.use_full_cuda_graph = compilation_config.full_cuda_graph
+        if self.use_full_cuda_graph:
+            # if not self.aot_schedule:
+            #     raise ValueError(
+            #         "AoT scheduling is required for full cuda graph.")
+            capture_sizes = compilation_config.cudagraph_capture_sizes
+            if not capture_sizes:
+                raise ValueError(
+                    "cudagraph_capture_sizes should not be None when "
+                    "full_cuda_graph is True.")
+            self.max_cudagraph_size = max(capture_sizes)
+            if self.max_cudagraph_size > 992:
+                # This condition derives from FA3's internal heuristic.
+                # TODO(woosuk): Support larger cudagraph sizes.
+                raise ValueError(
+                    "Capture size larger than 992 is not supported for "
+                    "full cuda graph.")
+
+            self.scheduler_metadata = torch.zeros(
+                self.runner.max_num_reqs + 1,
+                dtype=torch.int32,
+                device=self.runner.device,
+            )
+            # When using cuda graph, we need to set the upper bound of the
+            # number of splits so that large enough intermediate buffers are
+            # pre-allocated during capture.
+            self.max_num_splits = _DEFAULT_MAX_NUM_SPLITS_FOR_CUDA_GRAPH
+
+        # Sliding window size to be used with the AOT scheduler will be
+        # populated on first build() call.
+        self.aot_sliding_window: Optional[tuple[int, int]] = None
+
+    def build(
+        self, common_prefix_len: int,
+        common_attn_metadata: CommonAttentionMetadata
+    ) -> FlashAttentionMetadata:
+        num_reqs = common_attn_metadata.num_reqs
+        num_actual_tokens = common_attn_metadata.num_actual_tokens
+        max_query_len = common_attn_metadata.max_query_len
+
+        max_seq_len = int(self.runner.seq_lens_np[:num_reqs].max())
+        query_start_loc = common_attn_metadata.query_start_loc
+        seq_lens = common_attn_metadata.seq_lens
+        block_table = self.block_table
+        block_table_tensor = block_table.get_device_tensor()[:num_reqs]
+
+        block_table.slot_mapping[:num_actual_tokens].copy_(
+            block_table.slot_mapping_cpu[:num_actual_tokens],
+            non_blocking=True)
+        # Fill unused with -1. Needed for reshape_and_cache in full cuda graph
+        # mode.
+        block_table.slot_mapping[num_actual_tokens:].fill_(-1)
+
+        slot_mapping = block_table.slot_mapping[:num_actual_tokens]
+
+        if self.aot_sliding_window is None:
+            self.aot_sliding_window = (-1, -1)
+            # For the AOT scheduler we need the sliding window value to be
+            # constant for all layers to. We have to populate this on the first
+            # build() call so the layers are constructed (cannot populate)
+            # in __init__.
+            if self.aot_schedule:
+                sliding_window_configs = _get_sliding_window_configs(
+                    self.runner.vllm_config)
+                if len(sliding_window_configs) == 1:
+                    sliding_window_config = sliding_window_configs.pop()
+                    if sliding_window_config is not None:
+                        self.aot_sliding_window = sliding_window_config
+                elif len(sliding_window_configs) > 1:
+                    self.aot_schedule = False
+
+        def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
+                     max_seq_len, causal):
+            if self.aot_schedule:
+                return get_scheduler_metadata(
+                    batch_size=batch_size,
+                    max_seqlen_q=max_query_len,
+                    max_seqlen_k=max_seq_len,
+                    cache_seqlens=seqlens,
+                    num_heads_q=self.num_heads_q,
+                    num_heads_kv=self.num_heads_kv,
+                    headdim=self.headdim,
+                    page_size=self.block_size,
+                    cu_seqlens_q=cu_query_lens,
+                    causal=causal,
+                    window_size=self.aot_sliding_window,
+                    num_splits=self.max_num_splits,
+                )
+            return None
+
+        # for local attention
+        local_attn_metadata = None
+        if self.runner.attention_chunk_size is not None:
+            seqlens_q_local_np, virt_q_cu_seqlens_np, virt_k_seqlens_np, \
+                virt_block_table_tensor = make_local_attention_virtual_batches(
+                    self.runner.attention_chunk_size,
+                    self.runner.query_start_loc_np[:num_reqs + 1],
+                    self.runner.seq_lens_np[:num_reqs],
+                    block_table_tensor,
+                    self.block_size,
+                )
+            local_query_start_loc = torch.from_numpy(virt_q_cu_seqlens_np).to(
+                self.runner.device, non_blocking=True)
+            local_seqused_k = torch.from_numpy(virt_k_seqlens_np).to(
+                self.runner.device, non_blocking=True)
+            local_max_query_len = seqlens_q_local_np.max()
+            local_max_seq_len = virt_k_seqlens_np.max()
+            local_scheduler_metadata = schedule(
+                batch_size=local_query_start_loc.shape[0] - 1,
+                cu_query_lens=local_query_start_loc,
+                max_query_len=local_max_query_len,
+                seqlens=local_seqused_k,
+                max_seq_len=local_max_seq_len,
+                causal=True)
+
+            local_attn_metadata = FlashAttentionMetadata.LocalAttentionMetadata(
+                local_query_start_loc=local_query_start_loc,
+                local_seqused_k=local_seqused_k,
+                local_block_table=virt_block_table_tensor,
+                local_max_query_len=local_max_query_len,
+                local_max_seq_len=local_max_seq_len,
+                local_scheduler_metadata=local_scheduler_metadata,
+            )
+
+        use_cascade = common_prefix_len > 0
+
+        if use_cascade:
+            cu_prefix_query_lens = torch.tensor([0, num_actual_tokens],
+                                                dtype=torch.int32,
+                                                device=self.runner.device)
+            prefix_kv_lens = torch.tensor([common_prefix_len],
+                                          dtype=torch.int32,
+                                          device=self.runner.device)
+            suffix_kv_lens = (self.runner.seq_lens_np[:num_reqs] -
+                              common_prefix_len)
+            suffix_kv_lens = torch.from_numpy(suffix_kv_lens).to(
+                self.runner.device)
+            prefix_scheduler_metadata = schedule(
+                batch_size=1,
+                cu_query_lens=cu_prefix_query_lens,
+                max_query_len=num_actual_tokens,
+                seqlens=prefix_kv_lens,
+                max_seq_len=common_prefix_len,
+                causal=False)
+            scheduler_metadata = schedule(batch_size=num_reqs,
+                                          cu_query_lens=query_start_loc,
+                                          max_query_len=max_query_len,
+                                          seqlens=suffix_kv_lens,
+                                          max_seq_len=max_seq_len -
+                                          common_prefix_len,
+                                          causal=True)
+        else:
+            cu_prefix_query_lens = None
+            prefix_kv_lens = None
+            suffix_kv_lens = None
+            prefix_scheduler_metadata = None
+            scheduler_metadata = schedule(batch_size=num_reqs,
+                                          cu_query_lens=query_start_loc,
+                                          max_query_len=max_query_len,
+                                          seqlens=seq_lens,
+                                          max_seq_len=max_seq_len,
+                                          causal=True)
+
+        # if self.use_full_cuda_graph:
+        #     assert scheduler_metadata is not None
+        #     n = scheduler_metadata.shape[0]
+        #     self.scheduler_metadata[:n] = scheduler_metadata
+        #     # NOTE(woosuk): We should zero out the rest of the scheduler
+        #     # metadata to guarantee the correctness. Otherwise, some thread
+        #     # blocks may use the invalid scheduler metadata and overwrite the
+        #     # output buffer.
+        #     self.scheduler_metadata[n:] = 0
+        #     scheduler_metadata = self.scheduler_metadata[:n]
+
+        max_num_splits = 0
+        if (self.use_full_cuda_graph
+                and num_actual_tokens <= self.max_cudagraph_size):
+            # NOTE(woosuk): Setting num_splits > 1 may increase the memory
+            # usage, because the intermediate buffers of size [num_splits,
+            # num_heads, num_tokens, head_size] are allocated. Therefore,
+            # we only set num_splits when using cuda graphs.
+            max_num_splits = self.max_num_splits
+
+        attn_metadata = FlashAttentionMetadata(
+            num_actual_tokens=num_actual_tokens,
+            max_query_len=max_query_len,
+            query_start_loc=query_start_loc,
+            max_seq_len=max_seq_len,
+            seq_lens=seq_lens,
+            block_table=block_table_tensor,
+            slot_mapping=slot_mapping,
+            use_cascade=use_cascade,
+            common_prefix_len=common_prefix_len,
+            scheduler_metadata=scheduler_metadata,
+            cu_prefix_query_lens=cu_prefix_query_lens,
+            prefix_kv_lens=prefix_kv_lens,
+            suffix_kv_lens=suffix_kv_lens,
+            local_attn_metadata=local_attn_metadata,
+            prefix_scheduler_metadata=prefix_scheduler_metadata,
+            max_num_splits=max_num_splits,
+        )
+        return attn_metadata
+
+    def can_run_in_cudagraph(
+            self, common_attn_metadata: CommonAttentionMetadata) -> bool:
+        # Full CUDA Graph always supported (FA2 support checked separately)
+        return True
+
+    def use_cascade_attention(self, *args, **kwargs) -> bool:
+        return use_cascade_attention(*args, **kwargs)
+
+
+def ref_paged_attn(
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    query_lens: list[int],
+    kv_lens: list[int],
+    block_tables: torch.Tensor,
+    scale: float,
+    sliding_window: Optional[int] = None,
+    soft_cap: Optional[float] = None,
+    causal: bool = True,
+) -> torch.Tensor:
+    num_seqs = len(query_lens)
+    block_tables = block_tables.cpu().numpy()
+    _, block_size, num_kv_heads, head_size = key_cache.shape
+
+    outputs: list[torch.Tensor] = []
+    start_idx = 0
+    for i in range(num_seqs):
+        query_len = query_lens[i]
+        kv_len = kv_lens[i]
+        q = query[start_idx:start_idx + query_len]
+        q *= scale
+
+
+        num_kv_blocks = (kv_len + block_size - 1) // block_size
+        block_indices = block_tables[i, :num_kv_blocks]
+
+        k = key_cache[block_indices].view(-1, num_kv_heads, head_size)
+        k = k[:kv_len]
+        v = value_cache[block_indices].view(-1, num_kv_heads, head_size)
+        v = v[:kv_len]
+
+        if q.shape[1] != k.shape[1]:
+            k = torch.repeat_interleave(k, q.shape[1] // k.shape[1], dim=1)
+            v = torch.repeat_interleave(v, q.shape[1] // v.shape[1], dim=1)
+
+        attn = torch.einsum("qhd,khd->hqk", q, k).float()
+
+
+        if causal:
+            empty_mask = torch.ones(query_len, kv_len, device=query.device)
+            mask = torch.triu(empty_mask, diagonal=kv_len - query_len + 1).bool()
+            if sliding_window is not None:
+                sliding_window_mask = torch.triu(empty_mask,
+                                               diagonal=kv_len - (query_len + sliding_window) + 1).bool().logical_not()
+                mask |= sliding_window_mask
+            attn.masked_fill_(mask, float("-inf"))
+            attn[mask] = float("-inf")
+            # attn = torch.where(mask.unsqueeze(0), float("-inf"), attn)
+
+        if soft_cap is not None:
+            attn = soft_cap * torch.tanh(attn / soft_cap)
+
+        attn = torch.softmax(attn, dim=-1).to(v.dtype)
+        out = torch.einsum("hqk,khd->qhd", attn, v)
+
+
+        outputs.append(out)
+        start_idx += query_len
+
+    return torch.cat(outputs, dim=0)
+
+
+class FlashAttentionImpl(AttentionImpl):
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: Optional[list[float]],
+        sliding_window: Optional[int],
+        kv_cache_dtype: str,
+        blocksparse_params: Optional[dict[str, Any]] = None,
+        logits_soft_cap: Optional[float] = None,
+        attn_type: AttentionType = AttentionType.DECODER,
+        kv_sharing_target_layer_name: Optional[str] = None,
+        use_irope: bool = False,
+    ) -> None:
+        if blocksparse_params is not None:
+            raise ValueError(
+                "FlashAttention does not support block-sparse attention.")
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        if alibi_slopes is not None:
+            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
+        self.alibi_slopes = alibi_slopes
+        if sliding_window is None:
+            self.sliding_window = (-1, -1)
+        else:
+            self.sliding_window = (sliding_window - 1, 0)
+        self.kv_cache_dtype = kv_cache_dtype
+        if logits_soft_cap is None:
+            # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
+            logits_soft_cap = 0
+        self.logits_soft_cap = logits_soft_cap
+        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name
+
+        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+
+        FlashAttentionBackend.validate_head_size(head_size)
+
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError("Encoder self-attention and "
+                                      "encoder/decoder cross-attention "
+                                      "are not implemented for "
+                                      "FlashAttentionImpl")
+        self.use_irope = use_irope
+        self.vllm_flash_attn_version = get_flash_attn_version()
+        if is_quantized_kv_cache(self.kv_cache_dtype) \
+            and not flash_attn_supports_fp8():
+            raise NotImplementedError(
+                "FlashAttention does not support fp8 kv-cache on this device.")
+
+    def forward(
+        self,
+        layer: torch.nn.Module,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: FlashAttentionMetadata,
+        output: Optional[torch.Tensor] = None,
+        output_scale: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass with FlashAttention.
+
+        Args:
+            query: shape = [num_tokens, num_heads, head_size]
+            key: shape = [num_tokens, num_kv_heads, head_size]
+            value: shape = [num_tokens, num_kv_heads, head_size]
+            kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
+            attn_metadata: Metadata for attention.
+        Returns:
+            shape = [num_tokens, num_heads * head_size]
+        NOTE: FP8 quantization, flash-attn expect the size of
+              {q,k,v}_descale to be (num_sequences, num_kv_heads).
+              We use torch's .expand() to avoid duplicating values
+        """
+        assert output is not None, "Output tensor must be provided."
+
+        if output_scale is not None:
+            raise NotImplementedError(
+                "fused output quantization is not yet supported"
+                " for FlashAttentionImpl")
+
+        if attn_metadata is None:
+            # Profiling run.
+            return output
+
+        # IMPORTANT!
+        # NOTE(woosuk): With piece-wise CUDA graphs, this method is executed in
+        # eager-mode PyTorch. Thus, we need to be careful about any CPU overhead
+        # in this method. For example, `view` and `slice` (or `[:n]`) operations
+        # are surprisingly slow even in the case they do not invoke any GPU ops.
+        # Minimize the PyTorch ops in this method as much as possible.
+        # Whenever making a change in this method, please benchmark the
+        # performance to make sure it does not introduce any overhead.
+
+        num_actual_tokens = attn_metadata.num_actual_tokens
+        key_cache, value_cache = kv_cache.unbind(0)
+
+        if self.kv_sharing_target_layer_name is None:
+            # Reshape the input keys and values and store them in the cache.
+            # Skip this if sharing KV cache with an earlier attention layer.
+            # NOTE(woosuk): Here, key and value are padded while slot_mapping is
+            # not padded. However, we don't need to do key[:num_actual_tokens]
+            # and value[:num_actual_tokens] because the reshape_and_cache_flash
+            # op uses the slot_mapping's shape to determine the number of
+            # actual tokens.
+            torch.ops._C_cache_ops.reshape_and_cache_flash(
+                key,
+                value,
+                key_cache,
+                value_cache,
+                attn_metadata.slot_mapping,
+                self.kv_cache_dtype,
+                layer._k_scale,
+                layer._v_scale,
+            )
+
+        if self.kv_cache_dtype.startswith("fp8"):
+            assert False, "MUSA FLASH ATTENTION not support fp8 kvcache."
+            key_cache = key_cache.view(torch.float8_e4m3fn)
+            value_cache = value_cache.view(torch.float8_e4m3fn)
+            num_tokens, num_heads, head_size = query.shape
+            query, _ = ops.scaled_fp8_quant(
+                query.reshape(
+                    (num_tokens, num_heads * head_size)).contiguous(),
+                layer._q_scale)
+            query = query.reshape((num_tokens, num_heads, head_size))
+
+        # Compute attention and update output up to `num_actual_tokens`.
+        use_local_attn = \
+            (self.use_irope and attn_metadata.local_attn_metadata is not None)
+
+        if not attn_metadata.use_cascade or use_local_attn:
+            if use_local_attn:
+                assert attn_metadata.local_attn_metadata is not None
+                local_metadata = attn_metadata.local_attn_metadata
+                cu_seqlens_q = local_metadata.local_query_start_loc
+                seqused_k = local_metadata.local_seqused_k
+                max_seqlen_q = local_metadata.local_max_query_len
+                max_seqlen_k = local_metadata.local_max_seq_len
+                block_table = local_metadata.local_block_table
+                # scheduler_metadata = local_metadata.local_scheduler_metadata
+            else:
+                cu_seqlens_q = attn_metadata.query_start_loc
+                seqused_k = attn_metadata.seq_lens
+                max_seqlen_q = attn_metadata.max_query_len
+                max_seqlen_k = attn_metadata.max_seq_len
+                block_table = attn_metadata.block_table
+                # scheduler_metadata = attn_metadata.scheduler_metadata
+
+            ops.flash_attn_varlen_func(
+                query[:num_actual_tokens],
+                key_cache,
+                value_cache,
+                output[:num_actual_tokens],
+                cu_seqlens_q,
+                seqused_k,
+                block_table,
+                self.scale,
+                max_seqlen_q,
+                max_seqlen_k,
+                causal=True
+            )
+
+            # flash_attn_varlen_func(
+            #     q=query[:num_actual_tokens],
+            #     k=key_cache,
+            #     v=value_cache,
+            #     out=output[:num_actual_tokens],
+            #     cu_seqlens_q=cu_seqlens_q,
+            #     max_seqlen_q=max_seqlen_q,
+            #     seqused_k=seqused_k,
+            #     max_seqlen_k=max_seqlen_k,
+            #     softmax_scale=self.scale,
+            #     causal=True,
+            #     alibi_slopes=self.alibi_slopes,
+            #     window_size=self.sliding_window,
+            #     block_table=block_table,
+            #     softcap=self.logits_soft_cap,
+            #     scheduler_metadata=scheduler_metadata,
+            #     fa_version=self.vllm_flash_attn_version,
+            #     q_descale=layer._q_scale.expand(descale_shape),
+            #     k_descale=layer._k_scale.expand(descale_shape),
+            #     v_descale=layer._v_scale.expand(descale_shape),
+            #     num_splits=attn_metadata.max_num_splits,
+            # )
+            return output
+
+        assert not use_local_attn, (
+            "Cascade attention does not support local attention.")
+        # Cascade attention (rare case).
+        cascade_attention(
+            output[:num_actual_tokens],
+            query[:num_actual_tokens],
+            key_cache,
+            value_cache,
+            cu_query_lens=attn_metadata.query_start_loc,
+            max_query_len=attn_metadata.max_query_len,
+            cu_prefix_query_lens=attn_metadata.cu_prefix_query_lens,
+            prefix_kv_lens=attn_metadata.prefix_kv_lens,
+            suffix_kv_lens=attn_metadata.suffix_kv_lens,
+            max_kv_len=attn_metadata.max_seq_len,
+            softmax_scale=self.scale,
+            alibi_slopes=self.alibi_slopes,
+            sliding_window=self.sliding_window,
+            logits_soft_cap=self.logits_soft_cap,
+            block_table=attn_metadata.block_table,
+            common_prefix_len=attn_metadata.common_prefix_len,
+            fa_version=self.vllm_flash_attn_version,
+            prefix_scheduler_metadata=attn_metadata.prefix_scheduler_metadata,
+            suffix_scheduler_metadata=attn_metadata.scheduler_metadata,
+            q_descale=layer._q_scale,
+            k_descale=layer._k_scale,
+            v_descale=layer._v_scale,
+        )
+        return output
+
+
+def use_cascade_attention(
+    common_prefix_len: int,
+    query_lens: np.ndarray,
+    num_query_heads: int,
+    num_kv_heads: int,
+    use_alibi: bool,
+    use_sliding_window: bool,
+    num_sms: int,
+) -> bool:
+    """Decide whether to use cascade attention.
+
+    This function 1) checks whether cascade attention is supported with the
+    given configuration, and 2) heuristically decides whether using cascade
+    attention can improve performance.
+    """
+    # Too short common prefix. Probably not worth using cascade attention.
+    # We use an arbitrary threshold of 256 tokens. TODO: Tune this threshold.
+    # NOTE(woosuk): This is the common case. We should return False as soon as
+    # possible to avoid any unnecessary computation.
+    if common_prefix_len < 256:
+        return False
+    # Cascade attention is currently not supported with these variants.
+    if use_alibi or use_sliding_window:
+        return False
+    # Too few queries. Probably not worth using cascade attention.
+    # We use an arbitrary threshold of 8 queries. TODO: Tune this threshold.
+    num_reqs = len(query_lens)
+    if num_reqs < 8:
+        return False
+
+    # Heuristics to decide whether using cascade attention is beneficial.
+    # 1. When FlashDecoding is not used for normal attention, cascade attention
+    #    is likely to be faster since it saves memory bandwidth.
+    num_queries_per_kv = num_query_heads // num_kv_heads
+    # The criteria for using FlashDecoding can be found in the following link:
+    # https://github.com/vllm-project/flash-attention/blob/96266b1111111f3d11aabefaf3bacbab6a89d03c/csrc/flash_attn/flash_api.cpp#L535
+    use_flash_decoding = (num_queries_per_kv > 1 and not use_sliding_window
+                          and not use_alibi and np.all(query_lens == 1))
+    if not use_flash_decoding:
+        # Use cascade attention.
+        return True
+
+    # 2. When FlashDecoding is used for normal attention, it is not clear
+    #    whether cascade attention is beneficial, because FlashDecoding can
+    #    launch more CTAs than cascade attention.
+    #    We use a simple performance model to compare the two methods.
+    #    NOTE(woosuk): The performance model is very rough and may not be
+    #    accurate.
+    num_tokens = num_reqs
+    # NOTE(woosuk): These are default tile sizes. flash-attn might use
+    # different tile sizes (e.g., 64 or 256) depending on the configuration.
+    q_tile_size = 128
+    kv_tile_size = 128
+    num_prefix_tiles = cdiv(common_prefix_len, kv_tile_size)
+
+    cascade_ctas = num_query_heads * cdiv(num_tokens, q_tile_size)
+    cascade_waves = cdiv(cascade_ctas, num_sms)
+    cascade_time = cascade_waves * num_prefix_tiles
+
+    flash_decoding_ctas = (num_reqs * num_kv_heads *
+                           cdiv(num_queries_per_kv, q_tile_size))
+    flash_decoding_ctas *= num_prefix_tiles
+    flash_decoding_time = cdiv(flash_decoding_ctas, num_sms)
+
+    # Use cascade attention if it is faster than FlashDecoding.
+    return cascade_time < flash_decoding_time
+
+
+def cascade_attention(
+    output: torch.Tensor,
+    query: torch.Tensor,
+    key_cache: torch.Tensor,
+    value_cache: torch.Tensor,
+    cu_query_lens: torch.Tensor,
+    max_query_len: int,
+    cu_prefix_query_lens: torch.Tensor,
+    prefix_kv_lens: torch.Tensor,
+    suffix_kv_lens: torch.Tensor,
+    max_kv_len: int,
+    softmax_scale: float,
+    alibi_slopes: Optional[torch.Tensor],
+    sliding_window: tuple[int, int],
+    logits_soft_cap: float,
+    block_table: torch.Tensor,
+    common_prefix_len: int,
+    fa_version: int,
+    prefix_scheduler_metadata: Optional[torch.Tensor] = None,
+    suffix_scheduler_metadata: Optional[torch.Tensor] = None,
+    q_descale: Optional[torch.Tensor] = None,
+    k_descale: Optional[torch.Tensor] = None,
+    v_descale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    assert alibi_slopes is None, ("Cascade attention does not support ALiBi.")
+    # TODO: Support sliding window.
+    assert sliding_window == (-1, -1), (
+        "Cascade attention does not support sliding window.")
+
+    assert logits_soft_cap == None, (
+        "Cascade attention does not support logits soft cap.")
+
+    num_tokens = query.shape[0]
+    block_size = key_cache.shape[-3]
+    assert common_prefix_len % block_size == 0
+    num_common_kv_blocks = common_prefix_len // block_size
+    assert num_common_kv_blocks > 0
+    # descale_shape = (cu_prefix_query_lens.shape[0] - 1, key_cache.shape[-2])
+
+    # Process shared prefix.
+    prefix_output, prefix_lse = ops.flash_attn_varlen_func(
+                query,
+                key_cache,
+                value_cache,
+                output,
+                cu_prefix_query_lens,
+                prefix_kv_lens,
+                block_table[:1],
+                softmax_scale,
+                num_tokens,
+                common_prefix_len,
+                causal=False,
+                return_softmax_lse=True,
+            )
+    # prefix_output, prefix_lse = flash_attn_varlen_func(
+    #     q=query,
+    #     k=key_cache,
+    #     v=value_cache,
+    #     cu_seqlens_q=cu_prefix_query_lens,
+    #     seqused_k=prefix_kv_lens,
+    #     max_seqlen_q=num_tokens,
+    #     max_seqlen_k=common_prefix_len,
+    #     softmax_scale=softmax_scale,
+    #     causal=False,
+    #     window_size=sliding_window,
+    #     block_table=block_table[:1],
+    #     softcap=logits_soft_cap,
+    #     return_softmax_lse=True,
+    #     scheduler_metadata=prefix_scheduler_metadata,
+    #     fa_version=fa_version,
+    #     q_descale=q_descale.expand(descale_shape)
+    #     if q_descale is not None else None,
+    #     k_descale=k_descale.expand(descale_shape)
+    #     if k_descale is not None else None,
+    #     v_descale=v_descale.expand(descale_shape)
+    #     if v_descale is not None else None,
+    # )
+
+    # descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])
+
+    # Process suffix per query.
+
+    suffix_output, suffix_lse = ops.flash_attn_varlen_func(
+        query,
+        key_cache,
+        value_cache,
+        output,
+        cu_query_lens,
+        suffix_kv_lens,
+        block_table[:, num_common_kv_blocks:],
+        softmax_scale,
+        max_query_len,
+        max_kv_len - common_prefix_len,
+        causal=False,
+        return_softmax_lse=True,
+    )
+    # suffix_output, suffix_lse = flash_attn_varlen_func(
+    #     q=query,
+    #     k=key_cache,
+    #     v=value_cache,
+    #     cu_seqlens_q=cu_query_lens,
+    #     seqused_k=suffix_kv_lens,
+    #     max_seqlen_q=max_query_len,
+    #     max_seqlen_k=max_kv_len - common_prefix_len,
+    #     softmax_scale=softmax_scale,
+    #     causal=True,
+    #     window_size=sliding_window,
+    #     block_table=block_table[:, num_common_kv_blocks:],
+    #     softcap=logits_soft_cap,
+    #     return_softmax_lse=True,
+    #     scheduler_metadata=suffix_scheduler_metadata,
+    #     fa_version=fa_version,
+    #     q_descale=q_descale.expand(descale_shape)
+    #     if q_descale is not None else None,
+    #     k_descale=k_descale.expand(descale_shape)
+    #     if k_descale is not None else None,
+    #     v_descale=v_descale.expand(descale_shape)
+    #     if v_descale is not None else None,
+    # )
+
+    # Merge prefix and suffix outputs, and store the result in output.
+    ops.merge_attn_states(output, prefix_output, prefix_lse, suffix_output,
+                      suffix_lse)
diff --git a/vllm_musa/vllm_musa/distributed/__init__.py b/vllm_musa/vllm_musa/distributed/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/distributed/device_communicators/__init__.py b/vllm_musa/vllm_musa/distributed/device_communicators/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/distributed/device_communicators/musa_communicator.py b/vllm_musa/vllm_musa/distributed/device_communicators/musa_communicator.py
new file mode 100644
index 000000000..ff08026cd
--- /dev/null
+++ b/vllm_musa/vllm_musa/distributed/device_communicators/musa_communicator.py
@@ -0,0 +1,118 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from typing import Optional
+
+import torch
+from torch.distributed import ProcessGroup
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+
+from vllm.distributed.device_communicators.base_device_communicator import \
+    DeviceCommunicatorBase
+
+logger = init_logger(__name__)
+
+class MusaCommunicator(DeviceCommunicatorBase):
+
+    def __init__(self,
+                 cpu_group: ProcessGroup,
+                 device: Optional[torch.device] = None,
+                 device_group: Optional[ProcessGroup] = None,
+                 unique_name: str = ""):
+        super().__init__(cpu_group, device, device_group, unique_name)
+        # init device according to local rank
+        # local_rank = dist.get_rank(device_group)
+        # self.device = torch.device(f"musa:{local_rank}")
+        # Ref: https://github.com/vllm-project/vllm-ascend/commit/b64ee7d346511b6ea7a64b09db58c17aa1c915ef
+        self.device = torch.musa.current_device() # TODO(qn) should be deleted?
+        if "tp" not in unique_name:
+            # only tp uses custom allreduce
+            use_custom_allreduce = False
+        else:
+            from vllm.distributed.parallel_state import (
+                _ENABLE_CUSTOM_ALL_REDUCE)
+            use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
+
+        # ep does not use pynccl
+        use_pynccl = "ep" not in unique_name
+
+        self.use_pynccl = use_pynccl
+        self.use_custom_allreduce = use_custom_allreduce
+
+        # lazy import to avoid documentation build error
+        from vllm_musa.patch.custom_allreduce.patch_custom_allreduce import (
+            CustomAllreduce
+        )
+
+        self.ca_comm: Optional[CustomAllreduce] = None
+        if use_custom_allreduce and self.world_size > 1:
+            # Initialize a custom fast all-reduce implementation.
+            self.ca_comm = CustomAllreduce(
+                group=self.cpu_group,
+                device=self.device,
+            )
+
+        if self.use_all2all:
+            all2all_backend = envs.VLLM_ALL2ALL_BACKEND
+            if all2all_backend == "naive":
+                from vllm.distributed.device_communicators.all2all import NaiveAll2AllManager
+                self.all2all_manager = NaiveAll2AllManager(self.cpu_group)
+                logger.warning(
+                    "Using naive all2all manager, which is not"
+                    "efficient at all. The main purpose is for testing and"
+                    "debugging")
+            elif all2all_backend == "pplx":
+                raise ValueError(f"Unsupport all2all backend on musa platform: {all2all_backend}")
+            elif all2all_backend == "deepep_high_throughput":
+                raise ValueError(f"Unsupport all2all backend on musa platform: {all2all_backend}")
+            elif all2all_backend == "deepep_low_latency":
+                raise ValueError(f"Unsupport all2all backend on musa platform: {all2all_backend}")
+            else:
+                raise ValueError(f"Unknown all2all backend: {all2all_backend}")
+
+    def all_reduce(self, input_):
+        # always try custom allreduce first,
+        # and then pynccl.
+        ca_comm = self.ca_comm
+        if ca_comm is not None and not ca_comm.disabled and \
+            ca_comm.should_custom_ar(input_):
+            out = ca_comm.custom_all_reduce(input_)
+            assert out is not None
+            return out
+        torch.distributed.all_reduce(input_, group=self.device_group)
+        return input_
+
+    def destroy(self):
+        if self.ca_comm is not None:
+            self.ca_comm = None
+        if self.all2all_manager is not None:
+            self.all2all_manager.destroy()
+            self.all2all_manager = None
+
+    def dispatch(
+            self, hidden_states: torch.Tensor,
+            router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        assert self.all2all_manager is not None
+        hidden_states, router_logits = self.all2all_manager.dispatch(
+            hidden_states, router_logits)
+        return hidden_states, router_logits
+
+    def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        assert self.all2all_manager is not None
+        hidden_states = self.all2all_manager.combine(hidden_states)
+        return hidden_states
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/ops/__init__.py b/vllm_musa/vllm_musa/ops/__init__.py
new file mode 100644
index 000000000..41f6d7d66
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/__init__.py
@@ -0,0 +1,23 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import vllm_musa.ops.activation  # noqa
+import vllm_musa.ops.layernorm  # noqa
+import vllm_musa.ops.rotary_embedding  # noqa
+import vllm_musa.ops.triton_flash_attention  # noqa
+import vllm_musa.ops.quantization.gptq  # noqa
+import vllm_musa.ops.quantization.moe_wna16  # noqa
+import vllm_musa.ops.quantization.awq  # noqa
diff --git a/vllm_musa/vllm_musa/ops/activation.py b/vllm_musa/vllm_musa/ops/activation.py
new file mode 100644
index 000000000..d68fbd8a0
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/activation.py
@@ -0,0 +1,16 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/ops/deep_gemm.py b/vllm_musa/vllm_musa/ops/deep_gemm.py
new file mode 100644
index 000000000..11328b6f7
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/deep_gemm.py
@@ -0,0 +1,317 @@
+# Copy from https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py
+import functools
+from typing import Optional
+
+import torch
+
+import vllm.envs as envs
+from vllm.utils import round_up
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache, _fp8_quantize
+from vllm.model_executor.layers.fused_moe.deep_gemm_moe import _moe_permute
+from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import _moe_unpermute_and_reduce
+
+from vllm_musa import _musa_custom_ops as ops
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+def ceil_div(x: int, y: int) -> int:
+    """
+    Perform ceiling division of two integers.
+
+    Args:
+        x: the dividend.
+        y: the divisor.
+
+    Returns:
+        The result of the ceiling division.
+    """
+    return (x + y - 1) // y
+
+
+def get_tma_aligned_size(x: int, element_size: int) -> int:
+    """
+    Global memory address of TMA must be 16-byte aligned.
+    Since we use column-major layout for the LHS scaling tensor,
+        the M-axis of the LHS scaling tensor needs to be padded to a multiple of 16 bytes.
+
+    Arguments:
+        x: original M-axis shape of the LHS scaling tensor.
+        element_size: element size of the LHS scaling tensor.
+
+    Returns:
+        M-axis shape of the LHS scaling tensor after padding.
+    """
+    tma_alignment_bytes = 16
+    assert tma_alignment_bytes % element_size == 0
+    alignment = tma_alignment_bytes // element_size
+    return ceil_div(x, alignment) * alignment
+
+
+def get_col_major_tma_aligned_tensor(x: torch.Tensor) -> torch.Tensor:
+    """
+    Returns TMA-aligned transposed format of the input tensor. `torch.transpose` will be called if necessary.
+    If the input tensor is already column-major layout and 16-byte aligned along the M axis
+        (thus meets the requirement of LHS scaling tensor in DeepGEMM), this function will do nothing.
+
+    Arguments:
+        x: usually the LHS scaling tensor in GEMM.
+
+    Returns:
+        The LHS scaling tensor of TMA-aligned transposed format.
+    """
+    # NOTES: for the extreme performance, you may rewrite/fuse this function in CUDA
+    assert x.dim() in (2, 3)
+    remove_dim = False
+    m, n = x.shape[-2], x.shape[-1]
+    aligned_m = get_tma_aligned_size(m, x.element_size())
+    if x.dim() == 2:
+        if x.stride(0) == 1 and x.stride(1) == aligned_m:
+            return x
+        x, remove_dim = x.unsqueeze(0), True
+
+    b = x.shape[0]
+
+    # The last kernel gives a column-major TMA aligned layout
+    if x.stride(0) == aligned_m * n and x.stride(1) == 1 and x.stride(2) == aligned_m:
+        return x.squeeze(0) if remove_dim else x
+
+    # Normal layout requires transposing
+    aligned_x = torch.transpose(torch.empty((b, n, aligned_m), device=x.device, dtype=x.dtype), 1, 2)
+    aligned_x[:, :m, :] = x
+    aligned_x = aligned_x[:, :m, :]
+    return aligned_x.squeeze(0) if remove_dim else aligned_x
+
+def get_m_alignment_for_contiguous_layout():
+    """
+    When we do a grouped GEMM in contiguous format, LHS are grouped into several batches along the M axis.
+    Since we deal with exactly one sub-matrix of RHS for each GEMM block, batch sizes above should align well
+        with GEMM block shape.
+
+    Returns:
+        Group-level alignment requirement for grouped contiguous layout, which is always 128.
+    """
+    return 128
+
+@functools.cache
+def deep_gemm_block_shape() -> list[int]:
+    block = get_m_alignment_for_contiguous_layout()
+    return [block, block]
+
+
+def _valid_deep_gemm_shape(M: int, N: int, K: int):
+    align = deep_gemm_block_shape()[0]
+    return align <= M and N % align == 0 and K % align == 0
+
+
+def _valid_deep_gemm(hidden_states: torch.Tensor, w1: torch.Tensor,
+                     w2: torch.Tensor) -> bool:
+    """
+    Check if the given problem size is supported by the DeepGemm grouped
+    gemm kernel.  All of M, N, K and the quantization block_shape must be
+    aligned by `dg.get_m_alignment_for_contiguous_layout()`.
+    """
+
+    # the deep-gemm only has PH1 implementation.
+    device_prop = torch.cuda.get_device_properties(0)
+    if device_prop.major != 3 or device_prop.minor != 1:
+        return False
+
+    M = hidden_states.size(0)
+    _, K, N = w2.size()
+    if not _valid_deep_gemm_shape(M, N, K):
+        logger.debug("DeepGemm disabled: unalinged problem size.")
+        return False
+
+    if (w1.dtype != torch.float8_e4m3fn or w2.dtype != torch.float8_e4m3fn):
+        logger.debug("DeepGemm disabled: invalid weight dtype(s).")
+        return False
+
+    if (not hidden_states.is_contiguous() or not w1.is_contiguous()
+            or not w2.is_contiguous()):
+        logger.debug(
+            "DeepGemm disabled: weights or activations not contiguous.")
+        return False
+
+    return True
+
+def deep_gemm_moe_fp8(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    inplace: bool = False,
+    activation: str = "silu",
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    This function computes a a8w8-quantized Mixture of Experts (MoE) layer
+    using two sets of quantized weights, w1_q and w2_q, and top-k gating
+    mechanism. The matrix multiplications are implemented with DeepGemm
+    grouped gemm.
+
+    Parameters:
+    - hidden_states (torch.Tensor): The input tensor to the MoE layer.
+        Shape: [M, K]
+    - w1 (torch.Tensor): The first set of fp8 quantized expert weights.
+        Shape: [num_experts, K, 2N] (the weights are passed transposed)
+    - w2 (torch.Tensor): The second set of fp8 quantized expert weights.
+        Shape: [num_experts, N, K] (the weights are passed transposed)
+    - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
+        Shape: [num_experts] or [num_experts, 2N]
+    - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
+        Shape: [num_experts] or [num_experts, K]
+    - topk_weights (torch.Tensor): The weights of each token->expert mapping.
+    - topk_ids (torch.Tensor): The token->expert mapping for topk_weights.
+    - inplace (bool): If True, perform the operation in-place.
+        Defaults to False.
+    - activation (str): The activation function to apply after the first
+        MoE layer.
+    - global_num_experts (int): The total number of experts in the global
+        expert space.
+    - expert_map (Optional[torch.Tensor]):  A tensor mapping expert indices
+        from the global expert space to the local expert space of the expert
+        parallel shard.
+    - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
+        Shape: scalar or [M]
+    - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
+        quantize the intermediate result between the gemms.
+        Shape: scalar or [M]
+
+    Returns:
+    - torch.Tensor: The bfloat16 output tensor after applying the MoE layer.
+    """
+    # Lazy import to avoid CUDA initialization problems.
+    # import deep_gemm as dg
+
+    assert expert_map is None, "Expert maps not supported yet"
+
+    assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
+
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
+    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+    assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
+    assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
+    assert hidden_states.dtype in [
+        torch.float32, torch.float16, torch.bfloat16
+    ]
+    assert w1.dtype == torch.float8_e4m3fn
+    assert w2.dtype == torch.float8_e4m3fn
+    assert w1.shape[0] == w2.shape[0], "Expert number mismatch"
+    assert w1.shape[0] == w1_scale.shape[0], "w1 scales expert number mismatch"
+    assert w1.shape[0] == w2_scale.shape[0], "w2 scales expert number mismatch"
+    assert a1_scale is None or a1_scale.dim(
+    ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[
+        0] == hidden_states.shape[0], "Input scale shape mismatch"
+    assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch"  # noqa: E501
+
+    num_tokens, _ = hidden_states.shape
+    E, N, _ = w1.shape
+    K = w2.shape[1]
+    if global_num_experts == -1:
+        global_num_experts = E
+
+    # We execute the fused_moe kernel in chunks to circumvent this issue:
+    # https://github.com/vllm-project/vllm/issues/5938
+    CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
+
+    # assert _valid_deep_gemm(hidden_states, w1, w2, expert_map)
+
+    if inplace:
+        out_hidden_states = hidden_states
+    else:
+        out_hidden_states = torch.empty_like(hidden_states)
+
+    block_m = get_m_alignment_for_contiguous_layout()
+    block_shape = [block_m, block_m]
+
+    assert w1_scale is not None
+    assert w2_scale is not None
+
+    # We attempt to transpose and align offline in Fp8MoEMethod, in which
+    # case these calls will be nops.  Otherwise, they'll be performed every
+    # time the layer is executed.
+    # w1_scale = get_col_major_tma_aligned_tensor(w1_scale).contiguous()
+    # w2_scale = get_col_major_tma_aligned_tensor(w2_scale).contiguous()
+
+    M_sum = topk_ids.numel() + global_num_experts * (block_m - 1)
+    M_sum = round_up(M_sum, block_m)
+
+    num_chunks = (num_tokens // CHUNK_SIZE) + 1
+
+    # We can reuse the memory between cache1 and cache3 because by the time
+    # we need cache3, we're done with cache1
+    workspace13 = torch.empty(M_sum * max(N, K),
+                              device=hidden_states.device,
+                              dtype=hidden_states.dtype)
+
+    workspace1 = workspace13[:M_sum * N].view(M_sum, N)
+    workspace2 = torch.empty((M_sum, N // 2),
+                             device=hidden_states.device,
+                             dtype=hidden_states.dtype)
+    workspace3 = workspace13[:M_sum * K].view(M_sum, K)
+
+    for chunk in range(num_chunks):
+        begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
+                                          min((chunk + 1) * CHUNK_SIZE,
+                                              num_tokens))
+        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+        tokens_in_chunk, _ = curr_hidden_states.shape
+
+        if tokens_in_chunk == 0:
+            break
+
+        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+        a1q_scale: Optional[torch.Tensor] = None
+
+        qcurr_hidden_states, a1q_scale = _fp8_quantize(curr_hidden_states,
+                                                       a1_scale,
+                                                       per_act_token=False,
+                                                       block_shape=block_shape)
+
+        (qcurr_hidden_states, a1q_scale, sorted_token_ids, expert_ids,
+         inv_perm) = _moe_permute(qcurr_hidden_states, a1q_scale,
+                                  curr_topk_ids, global_num_experts,
+                                  expert_map, block_m)
+
+
+        # Adjust the intermediate cache size and config for the last chunk.
+        # Note that in most cases we only have one chunk so the cache size
+        # and config are already set correctly and do not need to be adjusted.
+        if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
+            curr_M = sorted_token_ids.numel()
+            workspace1 = _resize_cache(workspace1, (curr_M, N))
+            workspace2 = _resize_cache(workspace2, (curr_M, N // 2))
+            workspace3 = _resize_cache(workspace3, (curr_M, K))
+
+        ops.musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous((qcurr_hidden_states, a1q_scale), (w1, w1_scale), workspace1, expert_ids, block_m)
+
+        if activation == "silu":
+            torch.ops._C.silu_and_mul(workspace2, workspace1.view(-1, N))
+        elif activation == "gelu":
+            torch.ops._C.gelu_and_mul(workspace2, workspace1.view(-1, N))
+        else:
+            raise ValueError(f"Unsupported FusedMoe activation: {activation}")
+
+        a2q_scale: Optional[torch.Tensor] = None
+
+
+        qworkspace2, a2q_scale = _fp8_quantize(workspace2,
+                                               a2_scale,
+                                               per_act_token=False,
+                                               block_shape=block_shape)
+        ops.musa_m_grouped_gemm_fp8_fp8_bf16_nt_contiguous((qworkspace2, a2q_scale), (w2, w2_scale), workspace3, expert_ids, block_m)
+
+        _moe_unpermute_and_reduce(
+            out_hidden_states[begin_chunk_idx:end_chunk_idx],
+            workspace3.view(*workspace3.shape), inv_perm, curr_topk_weights, apply_router_weight_on_input=False)
+
+    return out_hidden_states
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/ops/layernorm.py b/vllm_musa/vllm_musa/ops/layernorm.py
new file mode 100644
index 000000000..1e7e20781
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/layernorm.py
@@ -0,0 +1,105 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm_musa import _musa_custom_ops as ops
+
+
+def __init__(
+        self,
+        hidden_size: int,
+        eps: float = 1e-6,
+        var_hidden_size: Optional[int] = None,
+        has_weight: bool = True,
+        dtype: Optional[torch.dtype] = None,
+        fuse_cast: Optional[bool] = False,
+    ) -> None:
+        super(RMSNorm, self).__init__()
+
+        self.hidden_size = hidden_size
+        self.variance_epsilon = eps
+        self.variance_size_override = (None if var_hidden_size == hidden_size
+                                       else var_hidden_size)
+        self.has_weight = has_weight
+        if dtype is not None:
+            self.weight = torch.ones(hidden_size, dtype=dtype)
+        else:
+            self.weight = torch.ones(hidden_size)
+        if self.has_weight:
+            self.weight = torch.nn.Parameter(self.weight)
+        self.fuse_cast = fuse_cast
+
+
+def rmsnorm_cast(
+  x: torch.Tensor, weight:torch.Tensor,  variance_epsilon: float
+) -> tuple[torch.Tensor, torch.Tensor]:
+
+    out, out_scale = ops.rmsnorm_cast(x, weight, variance_epsilon, quant_size=128)  # TODO: Hardcode quant_size to 128
+    return (out, out_scale)
+
+
+def rmsnorm_add_cast(
+    x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+    variance_epsilon: float
+)-> tuple[torch.Tensor, torch.Tensor]:
+
+    out_fp8, out_scale = ops.rmsnorm_add_cast(x, residual, weight, variance_epsilon, quant_size=128) # TODO: Hardcode quant_size to 128
+    return (out_fp8, out_scale), x     # x = x + residual, out is bf16 precison output
+
+
+def forward_oot(
+        self,
+        x: torch.Tensor,
+        residual: Optional[torch.Tensor] = None,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if not self.fuse_cast:
+            if residual is not None:
+                ops.musa_fused_add_rms_norm(
+                    x,
+                    residual,
+                    self.weight, # added by qn: self.weight.data
+                    self.variance_epsilon,
+                )
+                return x, residual
+            x = F.rms_norm(x, (self.hidden_size, ), self.weight, self.variance_epsilon)
+            # out = torch.empty_like(x)
+            # ops.rms_norm(
+            #     out,
+            #     x,
+            #     self.weight,  # added by qn: self.weight.data
+            #     self.variance_epsilon,
+            # )
+            return x
+        else:
+             if residual is not None:
+                  return rmsnorm_add_cast(x, residual, self.weight, self.variance_epsilon)
+                    # ops.fused_add_rms_norm(
+                    #     x,
+                    #     residual,
+                    #     self.weight, # added by qn: self.weight.data
+                    #     self.variance_epsilon,
+                    # )
+                    # return x, residual
+             else:
+                  return rmsnorm_cast(x, self.weight, self.variance_epsilon)
+
+RMSNorm.__init__ = __init__
+RMSNorm.forward = forward_oot
diff --git a/vllm_musa/vllm_musa/ops/moe_align_block_size.py b/vllm_musa/vllm_musa/ops/moe_align_block_size.py
new file mode 100644
index 000000000..a944cd52e
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/moe_align_block_size.py
@@ -0,0 +1,243 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from typing import Optional
+
+import torch
+
+import vllm.envs as envs
+from vllm_musa import _musa_custom_ops as ops
+from vllm.triton_utils import tl, triton
+from vllm.utils import round_up
+
+
+def ceil_div(a, b):
+    return (a + b - 1) // b
+
+
+@triton.jit
+def moe_align_block_size_stage1(
+    topk_ids_ptr,
+    tokens_cnts_ptr,
+    num_experts: tl.constexpr,
+    numel: tl.constexpr,
+    tokens_per_thread: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    start_idx = pid * tokens_per_thread
+
+    off_c = (pid + 1) * num_experts
+
+    for i in range(tokens_per_thread):
+        if start_idx + i < numel:
+            idx = tl.load(topk_ids_ptr + start_idx + i)
+            token_cnt = tl.load(tokens_cnts_ptr + off_c + idx)
+            tl.store(tokens_cnts_ptr + off_c + idx, token_cnt + 1)
+
+
+@triton.jit
+def moe_align_block_size_stage2(
+    tokens_cnts_ptr,
+    num_experts: tl.constexpr,
+):
+    pid = tl.program_id(0)
+
+    last_cnt = 0
+    for i in range(1, num_experts + 1):
+        token_cnt = tl.load(tokens_cnts_ptr + i * num_experts + pid)
+        last_cnt = last_cnt + token_cnt
+        tl.store(tokens_cnts_ptr + i * num_experts + pid, last_cnt)
+
+
+@triton.jit
+def moe_align_block_size_stage3(
+    total_tokens_post_pad_ptr,
+    tokens_cnts_ptr,
+    cumsum_ptr,
+    num_experts: tl.constexpr,
+    block_size: tl.constexpr,
+):
+    last_cumsum = 0
+    off_cnt = num_experts * num_experts
+    for i in range(1, num_experts + 1):
+        token_cnt = tl.load(tokens_cnts_ptr + off_cnt + i - 1)
+        last_cumsum = last_cumsum + tl.cdiv(token_cnt, block_size) * block_size
+        tl.store(cumsum_ptr + i, last_cumsum)
+    tl.store(total_tokens_post_pad_ptr, last_cumsum)
+
+
+@triton.jit
+def moe_align_block_size_stage4(
+    topk_ids_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    tokens_cnts_ptr,
+    cumsum_ptr,
+    num_experts: tl.constexpr,
+    block_size: tl.constexpr,
+    numel: tl.constexpr,
+    tokens_per_thread: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    start_idx = tl.load(cumsum_ptr + pid)
+    end_idx = tl.load(cumsum_ptr + pid + 1)
+
+    for i in range(start_idx, end_idx, block_size):
+        tl.store(expert_ids_ptr + i // block_size, pid)
+
+    start_idx = pid * tokens_per_thread
+    off_t = pid * num_experts
+
+    for i in range(start_idx, tl.minimum(start_idx + tokens_per_thread,
+                                         numel)):
+        expert_id = tl.load(topk_ids_ptr + i)
+        token_cnt = tl.load(tokens_cnts_ptr + off_t + expert_id)
+        rank_post_pad = token_cnt + tl.load(cumsum_ptr + expert_id)
+        tl.store(sorted_token_ids_ptr + rank_post_pad, i)
+        tl.store(tokens_cnts_ptr + off_t + expert_id, token_cnt + 1)
+
+
+# Triton implementation based on:
+# https://github.com/sgl-project/sglang/commit/ba5112ff691d791a9e38c6c71f59324a5fcb49d0
+def moe_align_block_size_triton(
+    topk_ids: torch.Tensor,
+    num_experts: int,
+    block_size: int,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_pad: torch.Tensor,
+) -> None:
+    numel = topk_ids.numel()
+    grid = (num_experts, )
+    tokens_cnts = torch.zeros((num_experts + 1, num_experts),
+                              dtype=torch.int32,
+                              device=topk_ids.device)
+    cumsum = torch.zeros((num_experts + 1, ),
+                         dtype=torch.int32,
+                         device=topk_ids.device)
+    tokens_per_thread = ceil_div(numel, num_experts)
+
+    moe_align_block_size_stage1[grid](
+        topk_ids,
+        tokens_cnts,
+        num_experts,
+        numel,
+        tokens_per_thread,
+    )
+    moe_align_block_size_stage2[grid](
+        tokens_cnts,
+        num_experts,
+    )
+    moe_align_block_size_stage3[(1, )](
+        num_tokens_post_pad,
+        tokens_cnts,
+        cumsum,
+        num_experts,
+        block_size,
+    )
+    moe_align_block_size_stage4[grid](
+        topk_ids,
+        sorted_token_ids,
+        expert_ids,
+        tokens_cnts,
+        cumsum,
+        num_experts,
+        block_size,
+        numel,
+        tokens_per_thread,
+    )
+
+
+def moe_align_block_size(
+    topk_ids: torch.Tensor,
+    block_size: int,
+    num_experts: int,
+    expert_map: Optional[torch.Tensor] = None,
+    pad_sorted_ids: bool = False
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Aligns the token distribution across experts to be compatible with block
+    size for matrix multiplication.
+
+    Parameters:
+    - topk_ids: A tensor of shape [total_tokens, top_k] representing the
+        top-k expert indices for each token.
+    - block_size: The block size used in block matrix multiplication.
+    - num_experts: The total number of experts.
+    - expert_map: A tensor of shape [num_experts] that maps the expert index
+        from the global space to the local index space of the current
+        expert parallel shard. If the expert is not in the current expert
+        parallel shard, the mapping is set to -1.
+    - pad_sorted_ids: A flag indicating whether the sorted_token_ids length
+      should be padded to a multiple of block_size,
+
+    Returns:
+    - sorted_token_ids: A tensor containing the sorted token indices according
+        to their allocated expert.
+    - expert_ids: A tensor indicating the assigned expert index for each block.
+    - num_tokens_post_padded: The total number of tokens after padding,
+        ensuring divisibility by block_size.
+
+    This function pads the number of tokens that each expert needs to process
+    so that it is divisible by block_size.
+    Padding ensures that during block matrix multiplication, the dimensions
+    align correctly.
+
+    Example:
+    Given topk_ids = [[2, 3, 4], [1, 2, 4], [1, 3, 4], [1, 2, 3]],
+    block_size = 4, and num_experts = 4:
+    - We initially have 12 tokens (after repeating 'top_k' times) and 4 experts,
+        with each expert needing to process 3 tokens.
+    - As block_size is 4, we pad 1 token for each expert.
+    - First, flatten topk_ids to [2, 3, 4, 1, 2, 4, 1, 3, 4, 1, 2, 3].
+    - Then append padding tokens [12, 12, 12, 12] for each block.
+    - After sorting by expert index, we obtain token_ids
+        [3, 6, 9, 12, 0, 4, 10, 12, 1, 7, 11, 12, 2, 5, 8, 12].
+        Tokens 12 are non-existent (padding) and are ignored in
+        the subsequent matrix multiplication.
+    - The padding ensures that the total number of tokens is now divisible
+        by block_size for proper block matrix operations.
+    """
+    max_num_tokens_padded = topk_ids.numel() + num_experts * (block_size - 1)
+    if pad_sorted_ids:
+        max_num_tokens_padded = round_up(max_num_tokens_padded, block_size)
+    sorted_ids = torch.empty((max_num_tokens_padded, ),
+                             dtype=torch.int32,
+                             device=topk_ids.device)
+    sorted_ids.fill_(topk_ids.numel())
+    max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size)
+    # Expert ids must be zeroed out to prevent index out of bounds error while
+    # mapping global expert ids to local expert ids in expert parallelism.
+    expert_ids = torch.zeros((max_num_m_blocks, ),
+                             dtype=torch.int32,
+                             device=topk_ids.device)
+    num_tokens_post_pad = torch.empty((1),
+                                      dtype=torch.int32,
+                                      device=topk_ids.device)
+    if num_experts >= 224:
+        if num_experts != 256:
+            moe_align_block_size_triton(
+                topk_ids,
+                num_experts,
+                block_size,
+                sorted_ids,
+                expert_ids,
+                num_tokens_post_pad,
+            )
+        else:
+            # Currently requires num_experts=256
+            ops.sgl_moe_align_block_size(
+                topk_ids,
+                num_experts,
+                block_size,
+                sorted_ids,
+                expert_ids,
+                num_tokens_post_pad,
+            )
+    else:
+        ops.moe_align_block_size(topk_ids, num_experts, block_size, sorted_ids,
+                                 expert_ids, num_tokens_post_pad)
+    if expert_map is not None:
+        expert_ids = expert_map[expert_ids]
+
+    return sorted_ids, expert_ids, num_tokens_post_pad
diff --git a/vllm_musa/vllm_musa/ops/quantization/__init__.py b/vllm_musa/vllm_musa/ops/quantization/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/ops/quantization/awq.py b/vllm_musa/vllm_musa/ops/quantization/awq.py
new file mode 100644
index 000000000..e847ebb99
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/quantization/awq.py
@@ -0,0 +1,51 @@
+from typing import Any, Dict, List, Optional
+
+import torch
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
+from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
+                                               UnquantizedLinearMethod)
+from vllm.model_executor.layers.quantization.awq import AWQConfig, AWQLinearMethod, is_layer_skipped_awq
+
+from vllm_musa import _musa_custom_ops as ops
+
+
+def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional["LinearMethodBase"]:
+    if isinstance(layer, LinearBase):
+        if is_layer_skipped_awq(prefix, self.modules_to_not_convert):
+            return UnquantizedLinearMethod()
+        return AWQLinearMethod(self)
+    # copy from awq_marlin.py
+    elif isinstance(layer, FusedMoE):
+        # if layer.num_experts > 32:
+        # For MoEs with many experts the moe_wna16 kernel is faster
+        return MoeWNA16Config.from_config(
+            self).get_quant_method(layer, prefix)
+
+    return None
+
+
+def apply(self,
+          layer: torch.nn.Module,
+          x: torch.Tensor,
+          bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+    qweight = layer.qweight
+    scales = layer.scales
+    qzeros = layer.qzeros
+    pack_factor = self.quant_config.pack_factor
+    out_shape = (x.shape[:-1] + (qweight.shape[-1] * pack_factor, ))
+    reshaped_x = x.reshape(-1, x.shape[-1])
+
+    # num_tokens >= threshold
+    # FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
+
+    out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
+    out = torch.matmul(reshaped_x, out)
+
+    if bias is not None:
+        out.add_(bias)
+    return out.reshape(out_shape)
+
+AWQConfig.get_quant_method = get_quant_method
+AWQLinearMethod.apply = apply
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/ops/quantization/gptq.py b/vllm_musa/vllm_musa/ops/quantization/gptq.py
new file mode 100644
index 000000000..aee7a94af
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/quantization/gptq.py
@@ -0,0 +1,219 @@
+import torch
+from torch.nn import Parameter
+from typing import Optional, List
+import numpy as np
+
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
+from vllm.model_executor.layers.quantization.gptq import (get_linear_quant_method,
+                                                          ExllamaState,
+                                                          GPTQConfig,
+                                                          GPTQLinearMethod)
+from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
+                                           GroupQuantScaleParameter,
+                                           PackedColumnParameter,
+                                           PackedvLLMParameter,
+                                           RowvLLMParameter)
+from vllm_musa import _musa_custom_ops as ops
+
+
+def get_quant_method(self, layer: torch.nn.Module,
+                     prefix: str) -> Optional["GPTQLinearMethod"]:
+
+    # copy from gptq_marlin.py
+    if isinstance(layer, FusedMoE):
+        # if layer.local_num_experts > 32:
+        # For MoEs with many experts the moe_wna16 kernel is faster
+        return MoeWNA16Config.from_config(
+            self).get_quant_method(layer, prefix)
+    return get_linear_quant_method(self, layer, prefix, GPTQLinearMethod)
+
+
+def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+    # per-group qweight shape: [in/8, out]
+    if self.quant_config.group_size != -1:
+        layer.qzeros = Parameter(layer.qzeros.data, requires_grad=False)
+        layer.qweight = Parameter(layer.qweight.data.t().contiguous(), requires_grad=False)
+        layer.g_idx = Parameter(layer.g_idx.data, requires_grad=False)
+        layer.scales = Parameter(layer.scales.data.t().contiguous(), requires_grad=False)
+    # per-channel mudnn qweight shape: [out, in/8]
+    else:
+        layer.qzeros = Parameter(layer.qzeros.data, requires_grad=False)
+        layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
+        layer.g_idx = Parameter(layer.g_idx.data, requires_grad=False)
+        layer.scales = Parameter(layer.scales.data.float(), requires_grad=False)
+        layer.qweight_t = Parameter(layer.qweight.data.T.contiguous(), requires_grad=False)
+
+
+def create_weights(
+    self,
+    layer: torch.nn.Module,
+    input_size_per_partition: int,
+    output_partition_sizes: List[int],
+    input_size: int,
+    output_size: int,
+    params_dtype: torch.dtype,
+    **extra_weight_attrs,
+):
+    del output_size  # Unused.
+    weight_loader = extra_weight_attrs.get("weight_loader")
+    if input_size_per_partition % self.quant_config.group_size != 0:
+        raise ValueError(
+            "The input size is not aligned with the quantized "
+            "weight shape. This can be caused by too large "
+            "tensor parallel size.")
+    output_size_per_partition = sum(output_partition_sizes)
+    if (output_size_per_partition % self.quant_config.pack_factor.numerator
+            != 0):
+        raise ValueError(
+            "The output size is not aligned with the quantized "
+            "weight shape. This can be caused by too large "
+            "tensor parallel size.")
+
+    if self.quant_config.group_size != -1:
+        group_size = self.quant_config.group_size
+    else:
+        group_size = input_size
+    exllama_state = ExllamaState.UNINITIALIZED
+    scale_and_zero_size = input_size // group_size
+    scale_and_zero_input_dim = None
+    if (input_size != input_size_per_partition
+            and self.quant_config.group_size != -1):
+        # For act-order models, we cannot use Exllama for row parallel layer
+        if self.quant_config.desc_act:
+            exllama_state = ExllamaState.UNUSED
+        else:
+            # we need to partition qzeros and scales for exllama kernel
+            scale_and_zero_size = input_size_per_partition // group_size
+            scale_and_zero_input_dim = 0
+
+    if self.quant_config.group_size != -1:
+        qweight = PackedvLLMParameter(
+            data=torch.empty(
+                input_size_per_partition // self.quant_config.pack_factor,
+                output_size_per_partition,
+                dtype=torch.int32,
+            ),
+            input_dim=0,
+            output_dim=1,
+            packed_dim=0,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader)
+    else:
+        qweight = PackedvLLMParameter(
+            data=torch.empty(
+                output_size_per_partition,
+                input_size_per_partition // self.quant_config.pack_factor,
+                dtype=torch.int32,
+            ),
+            input_dim=1,
+            output_dim=0,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            weight_loader=weight_loader)
+
+    g_idx = RowvLLMParameter(data=torch.tensor(
+        [
+            i // self.quant_config.group_size
+            for i in range(input_size_per_partition)
+        ],
+        dtype=torch.int32,
+    ),
+                                input_dim=0,
+                                weight_loader=weight_loader)
+    qzeros_args = {
+        "data":
+        torch.empty(
+            scale_and_zero_size,
+            output_size_per_partition // self.quant_config.pack_factor,
+            dtype=torch.int32,
+        ),
+        "weight_loader":
+        weight_loader
+    }
+    weight_scale_args = {
+        "data":
+        torch.empty(
+            scale_and_zero_size,
+            output_size_per_partition,
+            dtype=params_dtype,
+        ),
+        "weight_loader":
+        weight_loader
+    }
+    if scale_and_zero_input_dim is None:
+        scales = ChannelQuantScaleParameter(output_dim=1,
+                                            **weight_scale_args)
+        qzeros = PackedColumnParameter(
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            **qzeros_args)
+
+    else:
+        scales = GroupQuantScaleParameter(output_dim=1,
+                                            input_dim=0,
+                                            **weight_scale_args)
+        qzeros = PackedvLLMParameter(
+            input_dim=0,
+            output_dim=1,
+            packed_dim=1,
+            packed_factor=self.quant_config.pack_factor,
+            **qzeros_args)
+
+    layer.register_parameter("qweight", qweight)
+    layer.register_parameter("g_idx", g_idx)
+    layer.register_parameter("qzeros", qzeros)
+    layer.register_parameter("scales", scales)
+
+    layer.exllama_state = exllama_state
+
+
+def apply(self,
+          layer: torch.nn.Module,
+          x: torch.Tensor,
+          bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+    if self.quant_config.group_size != -1:
+        batch = x.shape[0]
+        if batch >= 64:
+            output = ops.group_scale_gemm(
+                x,
+                layer.qweight,
+                layer.scales,
+            )
+        else:
+            output = ops.fused_gemv(
+                x,
+                layer.qweight,
+                None,
+                layer.scales,
+            )
+
+        if bias is not None:
+            output.add_(bias)
+    else:
+        batch = x.shape[0]
+        if batch >= 64:
+            output = ops.mudnn_w4a16_matmul(
+                x,
+                layer.qweight,
+                layer.scales,
+                bias=bias,
+            )
+        else:
+            output = ops.gemm_w4a16(
+                x,
+                layer.qweight_t,
+                layer.scales
+            )
+            if bias is not None:
+                output.add_(bias)
+
+
+    return output
+
+
+GPTQConfig.get_quant_method = get_quant_method
+GPTQLinearMethod.process_weights_after_loading = process_weights_after_loading
+GPTQLinearMethod.apply = apply
+GPTQLinearMethod.create_weights = create_weights
diff --git a/vllm_musa/vllm_musa/ops/quantization/moe_wna16.py b/vllm_musa/vllm_musa/ops/quantization/moe_wna16.py
new file mode 100644
index 000000000..a08c4c817
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/quantization/moe_wna16.py
@@ -0,0 +1,398 @@
+import torch
+from typing import Any, Callable, Dict, List, Optional
+
+from vllm.distributed import get_tensor_model_parallel_rank, get_tp_group
+from vllm.model_executor.layers.quantization.moe_wna16 import (
+    MoeWNA16Config, MoeWNA16Method)
+from vllm.model_executor.layers.quantization.gptq import GPTQConfig
+from vllm.model_executor.layers.quantization.awq import AWQConfig
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.platforms import current_platform
+from vllm.model_executor.layers.fused_moe.layer import (
+    FusedMoE, FusedMoeWeightScaleSupported)
+from vllm_musa import _musa_custom_ops as ops
+
+
+def __init__(self, linear_quant_method: str, weight_bits: int,
+                group_size: int, has_zp: bool, lm_head_quantized: bool,
+                modules_to_not_convert: Optional[List[str]],
+                full_config: Dict[str, Any]) -> None:
+    super(MoeWNA16Config, self).__init__()
+    self.weight_bits = weight_bits
+    self.group_size = group_size
+    self.has_zp = has_zp
+    self.bit8_pack_factor = 8 // self.weight_bits
+    self.lm_head_quantized = lm_head_quantized
+    self.linear_quant_method = linear_quant_method
+    self.full_config = full_config
+    self.use_marlin = False
+    # Avoid circular import
+    from vllm.model_executor.layers.quantization.awq import AWQConfig
+    from vllm.model_executor.layers.quantization.awq_marlin import (
+        AWQMarlinConfig)
+    if self.linear_quant_method == "gptq":
+        # self.use_marlin = GPTQMarlinConfig.is_gptq_marlin_compatible(
+        #     full_config)
+        self.use_marlin = False
+    elif self.linear_quant_method == "awq":
+        # capability_tuple = current_platform.get_device_capability()
+        # device_capability = (-1 if capability_tuple is None else
+        #                         capability_tuple.to_int())
+        # awq_min_capability = AWQConfig.get_min_capability()
+        # if device_capability < awq_min_capability:
+        #     raise ValueError(
+        #         "The quantization method moe_wna16 + awq is not supported "
+        #         "for the current GPU. "
+        #         f"Minimum capability: {awq_min_capability}. "
+        #         f"Current capability: {device_capability}.")
+        # self.use_marlin = AWQMarlinConfig.is_awq_marlin_compatible(
+        #     full_config)
+        self.use_marlin = False
+    else:
+        raise ValueError("moe_wna16 only support gptq and awq.")
+
+    if modules_to_not_convert is None:
+        self.modules_to_not_convert = []
+    else:
+        self.modules_to_not_convert = modules_to_not_convert
+
+
+@classmethod
+def from_config(cls, config: Dict[str, Any]) -> "MoeWNA16Config":
+    # accept config as GPTQConfig and AWQConfig
+    if isinstance(config, GPTQConfig):
+        linear_quant_method = "gptq"
+        has_zp = False
+        lm_head_quantized = config.lm_head_quantized
+        modules_to_not_convert = []
+        return cls(linear_quant_method, config.weight_bits, config.group_size, has_zp,
+                lm_head_quantized, modules_to_not_convert, config)
+    if isinstance(config, AWQConfig):
+        linear_quant_method = "awq"
+        has_zp = config.zero_point
+        lm_head_quantized = False
+        modules_to_not_convert = config.modules_to_not_convert
+        return cls(linear_quant_method, config.weight_bits, config.group_size, has_zp,
+                lm_head_quantized, modules_to_not_convert, config)
+
+    linear_quant_method = cls.get_from_keys(config, ["quant_method"])
+    weight_bits = cls.get_from_keys(config, ["bits"])
+    group_size = cls.get_from_keys(config, ["group_size"])
+    lm_head_quantized = cls.get_from_keys_or(config, ["lm_head"],
+                                                default=False)
+    if linear_quant_method == "gptq":
+        has_zp = not cls.get_from_keys(config, ["sym"])
+        modules_to_not_convert = []
+    elif linear_quant_method == "awq":
+        has_zp = cls.get_from_keys(config, ["zero_point"])
+        modules_to_not_convert = cls.get_from_keys_or(
+            config, ["modules_to_not_convert"], None)
+    else:
+        raise ValueError("moe_wna16 only support gptq and awq.")
+
+    return cls(linear_quant_method, weight_bits, group_size, has_zp,
+                lm_head_quantized, modules_to_not_convert, config)
+
+
+def create_weights(self, layer: torch.nn.Module, num_experts: int,
+                       hidden_size: int, intermediate_size_per_partition: int,
+                       params_dtype: torch.dtype, **extra_weight_attrs):
+
+    layer.quant_config = self.quant_config
+    bit8_pack_factor = self.quant_config.bit8_pack_factor
+    bit32_pack_factor = bit8_pack_factor * 4
+    group_size = self.quant_config.group_size
+
+    if group_size != -1:
+        group_size_div_factor = 1
+
+        # make intermediate_size and hidden_size diviable by group_size
+        # we reduce the group size to ensure that
+        # and we would repeat the loaded_weight later
+        while intermediate_size_per_partition % group_size or \
+                hidden_size % group_size:
+            group_size = group_size // 2
+            group_size_div_factor *= 2
+            assert group_size >= 32
+        strategy = FusedMoeWeightScaleSupported.GROUP.value
+    else:
+        group_size = hidden_size
+        group_size_div_factor = 1
+        strategy = FusedMoeWeightScaleSupported.CHANNEL.value
+
+    layer.group_size = group_size
+    layer.group_size_div_factor = group_size_div_factor
+
+    extra_weight_attrs.update({
+        "quant_method": strategy,
+        "is_transposed": False
+    })
+
+    assert 'weight_loader' in extra_weight_attrs
+    weight_loader = extra_weight_attrs['weight_loader']
+    wrapped_weight_loader = MoeWNA16Method.get_weight_loader(
+        layer, weight_loader)
+    extra_weight_attrs['weight_loader'] = wrapped_weight_loader
+
+    # Fused gate_up_proj (column parallel)
+    w13_pack_factor = bit8_pack_factor if self.quant_config.group_size != -1 else bit32_pack_factor
+    w13_dtype = torch.uint8 if self.quant_config.group_size != -1 else torch.int32
+    w13_qweight = torch.nn.Parameter(torch.empty(
+        num_experts,
+        2 * intermediate_size_per_partition,
+        hidden_size // w13_pack_factor,
+        dtype=w13_dtype),
+                                        requires_grad=False)
+    layer.register_parameter("w13_qweight", w13_qweight)
+    set_weight_attrs(w13_qweight, extra_weight_attrs)
+
+    # down_proj (row parallel)
+    w2_pack_factor = bit8_pack_factor if self.quant_config.group_size != -1 else bit32_pack_factor
+    w2_dtype = torch.uint8 if self.quant_config.group_size != -1 else torch.int32
+    w2_qweight = torch.nn.Parameter(torch.empty(
+        num_experts,
+        hidden_size,
+        intermediate_size_per_partition // w2_pack_factor,
+        dtype=w2_dtype),
+                                    requires_grad=False)
+    layer.register_parameter("w2_qweight", w2_qweight)
+    set_weight_attrs(w2_qweight, extra_weight_attrs)
+
+    w13_group_size = group_size if self.quant_config.group_size != -1 else hidden_size
+    w13_scales = torch.nn.Parameter(torch.zeros(
+        num_experts,
+        2 * intermediate_size_per_partition,
+        hidden_size // w13_group_size,
+        dtype=torch.float32),
+                                    requires_grad=False)
+    layer.register_parameter("w13_scales", w13_scales)
+    set_weight_attrs(w13_scales, extra_weight_attrs)
+
+    w2_group_size = group_size if self.quant_config.group_size != -1 else intermediate_size_per_partition
+    w2_scales = torch.nn.Parameter(torch.zeros(
+        num_experts,
+        hidden_size,
+        intermediate_size_per_partition // w2_group_size,
+        dtype=torch.float32),
+                                    requires_grad=False)
+    layer.register_parameter("w2_scales", w2_scales)
+    set_weight_attrs(w2_scales, extra_weight_attrs)
+
+    if self.quant_config.has_zp:
+        w13_qzeros = torch.nn.Parameter(torch.zeros(
+            num_experts,
+            2 * intermediate_size_per_partition // bit8_pack_factor,
+            hidden_size // w13_group_size,
+            dtype=torch.uint8),
+                                        requires_grad=False)
+        layer.register_parameter("w13_qzeros", w13_qzeros)
+        set_weight_attrs(w13_qzeros, extra_weight_attrs)
+
+        w2_qzeros = torch.nn.Parameter(torch.zeros(
+            num_experts,
+            hidden_size // bit8_pack_factor,
+            intermediate_size_per_partition // w2_group_size,
+            dtype=torch.uint8),
+                                        requires_grad=False)
+        layer.register_parameter("w2_qzeros", w2_qzeros)
+        set_weight_attrs(w2_qzeros, extra_weight_attrs)
+
+    if self.quant_config.linear_quant_method == "gptq":
+        # some param are unused, but we need to init them in order to
+        # load weights
+        invalid_param_keys = ["w13_g_idx", "w2_g_idx"]
+        if not self.quant_config.has_zp:
+            invalid_param_keys += ["w13_qzeros", "w2_qzeros"]
+        for key in invalid_param_keys:
+            param = torch.nn.Parameter(torch.empty((0, ),
+                                                    dtype=torch.int32),
+                                        requires_grad=False)
+            layer.register_parameter(key, param)
+            set_weight_attrs(param, extra_weight_attrs)
+
+
+def apply(
+    self,
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    top_k: int,
+    renormalize: bool,
+    use_grouped_topk: bool = False,
+    topk_group: Optional[int] = None,
+    num_expert_group: Optional[int] = None,
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
+    custom_routing_function: Optional[Callable] = None,
+    scoring_func: str = "softmax",
+    e_score_correction_bias: Optional[torch.Tensor] = None,
+    apply_router_weight_on_input: bool = False,
+    activation: str = "silu",
+    enable_eplb: bool = False,
+    expert_load_view: Optional[torch.Tensor] = None,
+    logical_to_physical_map: Optional[torch.Tensor] = None,
+    logical_replica_count: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    if enable_eplb:
+        raise NotImplementedError(
+            "EPLB not supported for `MoeWNA16Method` yet.")
+    topk_weights, topk_ids = FusedMoE.select_experts(
+        hidden_states=x,
+        router_logits=router_logits,
+        use_grouped_topk=use_grouped_topk,
+        top_k=top_k,
+        renormalize=renormalize,
+        topk_group=topk_group,
+        num_expert_group=num_expert_group,
+        custom_routing_function=custom_routing_function,
+        scoring_func=scoring_func,
+        e_score_correction_bias=e_score_correction_bias)
+
+    from vllm.model_executor.layers.fused_moe import fused_experts
+    weight_bits = self.quant_config.weight_bits
+    has_zp = self.quant_config.has_zp
+
+    return fused_experts(
+        x,
+        layer.w13_qweight,
+        layer.w2_qweight,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        inplace=True,
+        use_int4_w4a16=weight_bits == 4,
+        use_int8_w8a16=weight_bits == 8,
+        global_num_experts=global_num_experts,
+        apply_router_weight_on_input=apply_router_weight_on_input,
+        expert_map=expert_map,
+        w1_scale=layer.w13_scales,
+        w2_scale=layer.w2_scales,
+        w1_zp=layer.w13_qzeros if has_zp else None,
+        w2_zp=layer.w2_qzeros if has_zp else None,
+        block_shape=[0, layer.group_size])
+
+
+@staticmethod
+def get_weight_loader(layer, weight_loader):
+
+    def convert_awq_tensor(tensor, tensor_type):
+        # convert awq qweight/qzeros to a standard format (assume int4)
+        # qweight: (k, n // pack_factor_bit32) -> (n, k // pack_factor_bit8)
+        # qzeros: (k // group_size, n // pack_factor_bit32) ->
+        #         (n // pack_factor_bit8, k // group_size)
+        # pack_factor_bit32 = 32 // weight_bits
+        # pack_factor_bit8 = 8 // weight_bits
+
+        # 0. suppose origin shape (a, b), dtype int32
+        # 1. convert to uint8, shape (a, b) -> (a, 4 * b)
+        size0 = tensor.size(0)
+        tensor = tensor.view(torch.uint8)
+
+        # 2. unpack to uint4 (only when weight_bits == 4)
+        #    shape (a, 4 * b) -> (a, 4 * b, 2)
+        shifter = torch.tensor([0, 4],
+                                dtype=torch.uint8,
+                                device=tensor.device)
+        tensor = (tensor[:, :, None] >> shifter) & 0xF
+
+        # 3. change order, see
+        # https://github.com/casper-hansen/AutoAWQ/blob/v0.2.8/awq/utils/quant_utils.py
+        # shape -> (a, 4 * b * pack_factor_bit8)
+        reverse_awq_pack_order = [0, 4, 1, 5, 2, 6, 3, 7]
+        tensor = tensor.view(-1, 8)[:, reverse_awq_pack_order]
+        tensor = tensor.view(size0, -1)
+
+        # 4. transpose, shape -> (4 * b * pack_factor_bit8, a)
+        tensor = tensor.T.contiguous()
+
+        # 5. repack (only when weight_bits == 4)
+        # qweight shape -> (4 * b * pack_factor_bit8, a // pack_factor_bit8)
+        # qzeros shape -> (4 * b, a)
+
+        if tensor_type == "qweight":
+            tensor = tensor[:, 1::2] * 16 + tensor[:, ::2]
+        elif tensor_type == "qzeros":
+            tensor = tensor[1::2, :] * 16 + tensor[::2, :]
+        return tensor
+
+    def convert_gptq_int4_qzeros(tensor):
+        tensor = tensor.view(torch.uint8)
+        shifter = torch.tensor([0, 4],
+                                dtype=torch.uint8,
+                                device=tensor.device)
+        tensor = (tensor[:, :, None] >> shifter) & 0xF
+        tensor = tensor + 1
+        tensor = tensor[:, :, 0] + tensor[:, :, 1] * 16
+        return tensor
+
+    def moe_wna16_weight_loader(param: torch.nn.Parameter,
+                                loaded_weight: torch.Tensor,
+                                weight_name: str, shard_id: str,
+                                expert_id: int):
+        if "g_idx" in weight_name:
+            return
+        if not layer.quant_config.has_zp and "qzeros" in weight_name:
+            return
+
+        device = get_tp_group().device
+        tp_rank = get_tensor_model_parallel_rank()
+        loaded_weight = loaded_weight.to(device)
+        shard_size = layer.intermediate_size_per_partition
+
+        # convert gptq and awq weight to a standard format
+        if layer.quant_config.linear_quant_method == "awq":
+            assert layer.quant_config.weight_bits == 4
+            if "weight" in weight_name:
+                loaded_weight = convert_awq_tensor(loaded_weight,
+                                                    "qweight")
+            elif "zeros" in weight_name:
+                loaded_weight = convert_awq_tensor(loaded_weight, "qzeros")
+            else:
+                loaded_weight = loaded_weight.T
+        elif layer.quant_config.linear_quant_method == "gptq":
+            assert layer.quant_config.weight_bits in [4, 8]
+            if "weight" in weight_name:
+                if layer.quant_config.group_size != -1:
+                    loaded_weight = loaded_weight.T.contiguous().view(
+                        torch.uint8)
+            elif "zeros" in weight_name:
+                # add 1 to gptq qzeros to align with awq
+                loaded_weight = loaded_weight.view(torch.uint8)
+                if layer.quant_config.weight_bits == 4:
+                    loaded_weight = convert_gptq_int4_qzeros(
+                        loaded_weight).T
+                else:
+                    loaded_weight = loaded_weight.T + 1
+            else:
+                loaded_weight = loaded_weight.T
+
+        # repeat the qzeros/scales to fit new group size
+        if layer.group_size_div_factor > 1 and \
+                "qzeros" in weight_name or "scales" in weight_name:
+            loaded_weight = loaded_weight.repeat_interleave(
+                layer.group_size_div_factor, 1)
+
+        if "scales" in weight_name:
+            loaded_weight = loaded_weight.to(torch.float32)
+
+        if "w13_qzeros" in weight_name:
+            tensor = loaded_weight.view(layer.tp_size, -1,
+                                        loaded_weight.size(1))[tp_rank]
+            if shard_id == "w1":
+                param.data[expert_id, :shard_size // 2] = tensor
+            else:
+                param.data[expert_id, shard_size // 2:] = tensor
+        elif "w2_qzeros" in weight_name:
+            param.data[expert_id] = loaded_weight.view(
+                loaded_weight.size(0), layer.tp_size, -1)[:, tp_rank]
+        else:
+            weight_loader(param, loaded_weight, weight_name, shard_id,
+                            expert_id)
+
+    return moe_wna16_weight_loader
+
+
+MoeWNA16Config.__init__ = __init__
+MoeWNA16Config.from_config = from_config
+MoeWNA16Method.create_weights = create_weights
+MoeWNA16Method.apply = apply
+MoeWNA16Method.get_weight_loader = get_weight_loader
diff --git a/vllm_musa/vllm_musa/ops/rotary_embedding.py b/vllm_musa/vllm_musa/ops/rotary_embedding.py
new file mode 100644
index 000000000..54a3ac5b8
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/rotary_embedding.py
@@ -0,0 +1,77 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from typing import Optional, Tuple
+
+import torch
+
+from vllm_musa import _musa_custom_ops as ops
+
+def rope_forward_oot(
+    self,
+    positions: torch.Tensor,
+    query: torch.Tensor,
+    key: Optional[torch.Tensor] = None,
+    offsets: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+
+    # __setattr__ in nn.Module (called by `self.cos_sin_cache = ...`)
+    # is expensive, so avoid calling it if possible
+    if self.cos_sin_cache.device != query.device or \
+        self.cos_sin_cache.dtype != query.dtype:
+        self.cos_sin_cache = self.cos_sin_cache.to(query.device,
+                                                    dtype=query.dtype)
+
+    # ops.rotary_embedding()/batched_rotary_embedding()
+    # are in-place operations that update the query and key tensors.
+    if offsets is not None:
+        ops.batched_rotary_embedding(positions, query, key, self.head_size,
+                                        self.cos_sin_cache,
+                                        self.is_neox_style, self.rotary_dim,
+                                        offsets)
+    else:
+        ops.rotary_embedding(positions, query, key, self.head_size,
+                                self.cos_sin_cache, self.is_neox_style)
+    return query, key
+
+def ds_rope_forward_oot(
+        self,
+        positions: torch.Tensor,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        offsets: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        if self.cos_sin_cache.device != query.device or \
+            self.cos_sin_cache.dtype != query.dtype:
+            self.cos_sin_cache = self.cos_sin_cache.to(query.device,
+                                                    dtype=query.dtype)
+        # ops.rotary_embedding()/batched_rotary_embedding()
+        # are in-place operations that update the query and key tensors.
+        if offsets is not None:
+            ops.batched_rotary_embedding_contiguous(positions, query, key, self.head_size,
+                                         self.cos_sin_cache,
+                                         self.is_neox_style, self.rotary_dim,
+                                         offsets)
+        else:
+            ops.rotary_embedding_contiguous(positions, query, key, self.head_size,
+                                 self.cos_sin_cache, self.is_neox_style)
+        return query, key
+
+import vllm.model_executor.layers.rotary_embedding
+vllm.model_executor.layers.rotary_embedding.RotaryEmbedding.forward = rope_forward_oot
+vllm.model_executor.layers.rotary_embedding.DeepseekScalingRotaryEmbedding.forward = ds_rope_forward_oot
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/ops/triton_flash_attention.py b/vllm_musa/vllm_musa/ops/triton_flash_attention.py
new file mode 100644
index 000000000..194746f2e
--- /dev/null
+++ b/vllm_musa/vllm_musa/ops/triton_flash_attention.py
@@ -0,0 +1,736 @@
+#!/usr/bin/env python
+# SPDX-License-Identifier: Apache-2.0
+"""
+Fused Attention
+===============
+
+This is a Triton implementation of the Flash Attention v2 algorithm from Tri Dao
+(https://tridao.me/publications/flash2/flash2.pdf)
+Credits: OpenAI kernel team, AMD ML Frameworks Triton team
+
+Features supported:
+
+1) Fwd with causal masking
+2) Any sequence lengths without padding (currently fwd kernel only)
+3) Support for different sequence lengths for q and k
+4) Nested tensor API currently does not support dropout or bias.
+
+Not currently supported:
+
+1) Non power of two head dims
+
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+torch_dtype: tl.constexpr = torch.float16
+
+
+@triton.jit
+def cdiv_fn(x, y):
+    return (x + y - 1) // y
+
+
+@triton.jit
+def max_fn(x, y):
+    return tl.math.max(x, y)
+
+
+@triton.jit
+def dropout_offsets(philox_seed, philox_offset, dropout_p, m, n, stride):
+    ms = tl.arange(0, m)
+    ns = tl.arange(0, n)
+    return philox_offset + ms[:, None] * stride + ns[None, :]
+
+
+@triton.jit
+def dropout_rng(philox_seed, philox_offset, dropout_p, m, n, stride):
+    rng_offsets = dropout_offsets(philox_seed, philox_offset, dropout_p, m, n,
+                                  stride).to(tl.uint32)
+    # TODO: use tl.randint for better performance
+    return tl.rand(philox_seed, rng_offsets)
+
+
+@triton.jit
+def dropout_mask(philox_seed, philox_offset, dropout_p, m, n, stride):
+    rng_output = dropout_rng(philox_seed, philox_offset, dropout_p, m, n,
+                             stride)
+    rng_keep = rng_output > dropout_p
+    return rng_keep
+
+
+@triton.jit
+def load_fn(block_ptr, first, second, pad):
+    if first and second:
+        tensor = tl.load(block_ptr, boundary_check=(0, 1), padding_option=pad)
+    elif first:
+        tensor = tl.load(block_ptr, boundary_check=(0, ), padding_option=pad)
+    elif second:
+        tensor = tl.load(block_ptr, boundary_check=(1, ), padding_option=pad)
+    else:
+        tensor = tl.load(block_ptr)
+    return tensor
+
+
+@triton.jit
+def _attn_fwd_inner(
+    acc,
+    l_i,
+    m_i,
+    q,
+    K_block_ptr,
+    V_block_ptr,
+    start_m,
+    actual_seqlen_k,
+    dropout_p,
+    philox_seed,
+    batch_philox_offset,
+    encoded_softmax_block_ptr,
+    block_min,
+    block_max,
+    offs_n_causal,
+    masked_blocks,
+    n_extra_tokens,
+    bias_ptr,
+    IS_CAUSAL: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    OFFS_M: tl.constexpr,
+    OFFS_N: tl.constexpr,
+    PRE_LOAD_V: tl.constexpr,
+    MASK_STEPS: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    RETURN_ENCODED_SOFTMAX: tl.constexpr,
+    PADDED_HEAD: tl.constexpr,
+):
+    # loop over k, v, and update accumulator
+    for start_n in range(block_min, block_max, BLOCK_N):
+        # For padded blocks, we will overrun the tensor size if
+        # we load all BLOCK_N. For others, the blocks are all within range.
+        k = load_fn(
+            K_block_ptr,
+            PADDED_HEAD,
+            MASK_STEPS and (n_extra_tokens != 0),
+            "zero",
+        )
+        if PRE_LOAD_V:
+            v = load_fn(
+                V_block_ptr,
+                MASK_STEPS and (n_extra_tokens != 0),
+                PADDED_HEAD,
+                "zero",
+            )
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        # We start from end of seqlen_k so only the first iteration would need
+        # to be checked for padding if it is not a multiple of block_n
+        # TODO: This can be optimized to only be true for the padded block.
+        if MASK_STEPS:  # noqa: SIM102
+            # If this is the last block / iteration, we want to
+            # mask if the sequence length is not a multiple of block size
+            # a solution is to always do BLOCK_M // BLOCK_N + 1 steps
+            # if not is_modulo_mn. last step might get wasted but that is okay.
+            # check if this masking works for that case.
+            if (start_n + BLOCK_N == block_max) and (n_extra_tokens != 0):
+                boundary_m = tl.full([BLOCK_M],
+                                     actual_seqlen_k,
+                                     dtype=tl.int32)
+                size_n = start_n + OFFS_N[None, :]
+                mask = size_n < boundary_m[:, None]
+                qk = tl.where(mask, qk, float("-inf"))
+        if IS_CAUSAL:
+            causal_boundary = start_n + offs_n_causal
+            causal_mask = OFFS_M[:, None] >= causal_boundary[None, :]
+            qk = tl.where(causal_mask, qk, float("-inf"))
+        # -- compute qk ----
+        qk += tl.dot(q, k)
+        if bias_ptr is not None:
+            bias = load_fn(bias_ptr, False, MASK_STEPS
+                           and (n_extra_tokens != 0), "zero")
+            # While bias is added after multiplying qk with sm_scale, our
+            # optimization to use 2^x instead of e^x results in an additional
+            # scale factor of log2(e) which we must also multiply the bias with.
+            qk += bias * 1.44269504089
+        m_ij = tl.maximum(m_i, tl.max(qk, 1))
+        qk = qk - m_ij[:, None]
+        p = tl.math.exp2(qk)
+
+        # CAVEAT: Must update l_ij before applying dropout
+        l_ij = tl.sum(p, 1)
+        if ENABLE_DROPOUT:
+            philox_offset = (batch_philox_offset +
+                             start_m * BLOCK_M * actual_seqlen_k + start_n -
+                             BLOCK_N)
+            keep = dropout_mask(
+                philox_seed,
+                philox_offset,
+                dropout_p,
+                BLOCK_M,
+                BLOCK_N,
+                actual_seqlen_k,
+            )
+            if RETURN_ENCODED_SOFTMAX:
+                tl.store(
+                    encoded_softmax_block_ptr,
+                    tl.where(keep, p,
+                             -p).to(encoded_softmax_block_ptr.type.element_ty),
+                )
+            p = tl.where(keep, p, 0.0)
+        elif RETURN_ENCODED_SOFTMAX:
+            tl.store(
+                encoded_softmax_block_ptr,
+                p.to(encoded_softmax_block_ptr.type.element_ty),
+            )
+        # -- update output accumulator --
+        alpha = tl.math.exp2(m_i - m_ij)
+        acc = acc * alpha[:, None]
+        if not PRE_LOAD_V:
+            v = load_fn(
+                V_block_ptr,
+                MASK_STEPS and (n_extra_tokens != 0),
+                PADDED_HEAD,
+                "zero",
+            )
+        # -- update m_i and l_i
+        l_i = l_i * alpha + l_ij
+        # update m_i and l_i
+        m_i = m_ij
+        acc += tl.dot(p.to(V_block_ptr.type.element_ty), v)
+        V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0))
+        K_block_ptr = tl.advance(K_block_ptr, (0, BLOCK_N))
+        if bias_ptr is not None:
+            bias_ptr = tl.advance(bias_ptr, (0, BLOCK_N))
+        if RETURN_ENCODED_SOFTMAX:
+            encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
+                                                   (0, BLOCK_N))
+    return acc, l_i, m_i
+
+
+@triton.autotune(
+    configs=[
+        triton.Config(
+            {
+                "BLOCK_M": 32,
+                "BLOCK_N": 32,
+                "PRE_LOAD_V": False,
+            },
+            num_warps=8,
+        ),
+    ],
+    key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'],
+)
+@triton.jit
+def attn_fwd(
+    Q,
+    K,
+    V,
+    bias,
+    sm_scale,
+    L,
+    Out,
+    stride_qz,
+    stride_qh,
+    stride_qm,
+    stride_qk,
+    stride_kz,
+    stride_kh,
+    stride_kn,
+    stride_kk,
+    stride_vz,
+    stride_vh,
+    stride_vk,
+    stride_vn,
+    stride_oz,
+    stride_oh,
+    stride_om,
+    stride_on,
+    stride_bz,
+    stride_bh,
+    stride_bm,
+    stride_bn,
+    cu_seqlens_q,
+    cu_seqlens_k,
+    dropout_p,
+    philox_seed,
+    philox_offset_base,
+    encoded_softmax,
+    HQ: tl.constexpr,
+    HK: tl.constexpr,
+    ACTUAL_BLOCK_DMODEL: tl.constexpr,
+    MAX_SEQLENS_Q: tl.constexpr,
+    MAX_SEQLENS_K: tl.constexpr,
+    VARLEN: tl.constexpr,
+    IS_CAUSAL: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    PRE_LOAD_V: tl.constexpr,
+    BIAS_TYPE: tl.constexpr,
+    ENABLE_DROPOUT: tl.constexpr,
+    RETURN_ENCODED_SOFTMAX: tl.constexpr,
+):
+    start_m = tl.program_id(0)
+    off_h_q = tl.program_id(1)
+    off_z = tl.program_id(2)
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = tl.arange(0, BLOCK_N)
+    if VARLEN:
+        cu_seqlens_q_start = tl.load(cu_seqlens_q + off_z)
+        cu_seqlens_q_end = tl.load(cu_seqlens_q + off_z + 1)
+        seqlen_q = cu_seqlens_q_end - cu_seqlens_q_start
+        # We have a one-size-fits-all grid in id(0). Some seqlens might be too
+        # small for all start_m so for those we return early.
+        if start_m * BLOCK_M > seqlen_q:
+            return
+        cu_seqlens_k_start = tl.load(cu_seqlens_k + off_z)
+        cu_seqlens_k_end = tl.load(cu_seqlens_k + off_z + 1)
+        seqlen_k = cu_seqlens_k_end - cu_seqlens_k_start
+    else:
+        cu_seqlens_q_start = 0
+        cu_seqlens_k_start = 0
+        seqlen_q = MAX_SEQLENS_Q
+        seqlen_k = MAX_SEQLENS_K
+
+    # Now we compute whether we need to exit early due to causal masking.
+    # This is because for seqlen_q > seqlen_k, M rows of the attn scores
+    # are completely masked, resulting in 0s written to the output, and
+    # inf written to LSE. We don't need to do any GEMMs in this case.
+    # This block of code determines what N is, and if this WG is operating
+    # on those M rows.
+    n_blocks = cdiv_fn(seqlen_k, BLOCK_N)
+    if IS_CAUSAL:
+        # If seqlen_q == seqlen_k, the attn scores are a square matrix.
+        # If seqlen_q != seqlen_k, attn scores are rectangular which means
+        # the causal mask boundary is bottom right aligned, and ends at either
+        # the top edge (seqlen_q < seqlen_k) or left edge.
+        # This captures the decrease in n_blocks if we have a rectangular attn
+        # matrix
+        n_blocks_seqlen = cdiv_fn(
+            (start_m + 1) * BLOCK_M + seqlen_k - seqlen_q, BLOCK_N)
+        # This is what adjusts the block_max for the current WG, only
+        # if IS_CAUSAL. Otherwise we want to always iterate through all n_blocks
+        n_blocks = min(n_blocks, n_blocks_seqlen)
+        # If we have no blocks after adjusting for seqlen deltas, this WG is
+        # part of the blocks that are all 0. We exit early.
+        if n_blocks <= 0:
+            o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om +
+                        off_h_q * stride_oh)
+            O_block_ptr = tl.make_block_ptr(
+                base=Out + o_offset,
+                shape=(seqlen_q, BLOCK_DMODEL),
+                strides=(stride_om, stride_on),
+                offsets=(start_m * BLOCK_M, 0),
+                block_shape=(BLOCK_M, BLOCK_DMODEL),
+                order=(1, 0),
+            )
+            acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=Out.type.element_ty)
+            # We still need to write 0s to the result
+            # tl.store(O_block_ptr,
+            # acc.to(Out.type.element_ty), boundary_check=(0,1))
+            # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q
+            #          + offs_m
+            # We store inf to LSE, not -inf because in the bwd pass,
+            # we subtract this
+            # from qk which makes it -inf, such that exp(qk - inf) = 0
+            # for these masked blocks.
+            # l = tl.full([BLOCK_M], value=float("inf"), dtype=tl.float32)
+            # tl.store(l_ptrs, l)
+            # TODO: Should dropout and return encoded softmax be handled here?
+            return
+
+    # If MQA / GQA, set the K and V head offsets appropriately.
+    GROUP_SIZE: tl.constexpr = HQ // HK
+    off_h_k = off_h_q // GROUP_SIZE if GROUP_SIZE != 1 else off_h_q
+
+    n_extra_tokens = 0
+    if seqlen_k < BLOCK_N:
+        n_extra_tokens = BLOCK_N - seqlen_k
+    elif seqlen_k % BLOCK_N:
+        n_extra_tokens = seqlen_k % BLOCK_N
+    padded_head = ACTUAL_BLOCK_DMODEL != BLOCK_DMODEL
+
+    # Compute pointers for all the tensors used in this kernel.
+    q_offset = (off_z * stride_qz + off_h_q * stride_qh +
+                cu_seqlens_q_start * stride_qm)
+    Q_block_ptr = tl.make_block_ptr(
+        base=Q + q_offset,
+        shape=(seqlen_q, ACTUAL_BLOCK_DMODEL),
+        strides=(stride_qm, stride_qk),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    k_offset = (off_z * stride_kz + off_h_k * stride_kh +
+                cu_seqlens_k_start * stride_kn)
+    K_block_ptr = tl.make_block_ptr(
+        base=K + k_offset,
+        shape=(ACTUAL_BLOCK_DMODEL, seqlen_k),
+        strides=(stride_kk, stride_kn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_DMODEL, BLOCK_N),
+        order=(0, 1),
+    )
+    v_offset = (off_z * stride_vz + off_h_k * stride_vh +
+                cu_seqlens_k_start * stride_vk)
+    V_block_ptr = tl.make_block_ptr(
+        base=V + v_offset,
+        shape=(seqlen_k, ACTUAL_BLOCK_DMODEL),
+        strides=(stride_vk, stride_vn),
+        offsets=(0, 0),
+        block_shape=(BLOCK_N, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    if BIAS_TYPE != 0:
+        bias_ptr = tl.make_block_ptr(
+            base=bias + off_h_q * stride_bh,
+            shape=(seqlen_q, seqlen_k),
+            strides=(stride_bm, stride_bn),
+            offsets=(start_m * BLOCK_M, 0),
+            block_shape=(BLOCK_M, BLOCK_N),
+            order=(1, 0),
+        )
+    else:
+        bias_ptr = None
+    if ENABLE_DROPOUT:
+        batch_philox_offset = philox_offset_base \
+                              + (off_z * HQ + off_h_q) \
+                              * seqlen_q * seqlen_k
+    else:
+        batch_philox_offset = 0
+    # We can ask to return the dropout mask without actually doing any dropout.
+    # In this case, we return an invalid pointer so indicate the mask is not i
+    # valid.
+    # TODO: Fix encoded softmax. It currently uses just h_q in the base offset.
+    if RETURN_ENCODED_SOFTMAX:
+        encoded_softmax_block_ptr = tl.make_block_ptr(
+            base=encoded_softmax + off_h_q * seqlen_q * seqlen_k,
+            shape=(seqlen_q, seqlen_k),
+            strides=(seqlen_k, 1),
+            offsets=(start_m * BLOCK_M, 0),
+            block_shape=(BLOCK_M, BLOCK_N),
+            order=(1, 0),
+        )
+    else:
+        encoded_softmax_block_ptr = 0
+    # initialize pointer to m and l
+    m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)
+    l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+    # scale sm_scale by log_2(e) and use 2^x in the loop as we do not
+    # have native e^x support in HW.
+    qk_scale = sm_scale * 1.44269504089
+    # Q is loaded once at the beginning and shared by all N blocks.
+    q = load_fn(Q_block_ptr, True, padded_head, "zero")
+    q = (q * qk_scale).to(Q_block_ptr.type.element_ty)
+
+    # Here we compute how many full and masked blocks we have.
+    padded_block_k = n_extra_tokens != 0
+    is_modulo_mn = not padded_block_k and (seqlen_q % BLOCK_M == 0)
+    if IS_CAUSAL:
+        # There are always at least BLOCK_M // BLOCK_N masked blocks.
+        # Additionally there might be one more due to dissimilar seqlens.
+        masked_blocks = BLOCK_M // BLOCK_N + (not is_modulo_mn)
+    else:
+        # Padding on Q does not need to be masked in the FA loop.
+        masked_blocks = padded_block_k
+    # if IS_CAUSAL, not is_modulo_mn does not always result in an additional
+    # block. In this case we might exceed n_blocks so pick the min.
+    masked_blocks = min(masked_blocks, n_blocks)
+    n_full_blocks = n_blocks - masked_blocks
+    block_min = 0
+    block_max = n_blocks * BLOCK_N
+    # Compute for full blocks. Here we set causal to false regardless of its
+    # value because there is no masking. Similarly we do not need padding.
+    if n_full_blocks > 0:
+        block_max = (n_blocks - masked_blocks) * BLOCK_N
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            K_block_ptr,
+            V_block_ptr,
+            start_m,
+            seqlen_k,
+            dropout_p,
+            philox_seed,
+            batch_philox_offset,
+            encoded_softmax_block_ptr,
+            # _, _, offs_n_causal, masked_blocks, n_extra_tokens, _
+            block_min,
+            block_max,
+            0,
+            0,
+            0,
+            bias_ptr,
+            # IS_CAUSAL, ....
+            False,
+            BLOCK_M,
+            BLOCK_DMODEL,
+            BLOCK_N,
+            offs_m,
+            offs_n,
+            # _, MASK_STEPS, ...
+            PRE_LOAD_V,
+            False,
+            ENABLE_DROPOUT,
+            RETURN_ENCODED_SOFTMAX,
+            padded_head,
+        )
+        block_min = block_max
+        block_max = n_blocks * BLOCK_N
+
+    tl.debug_barrier()
+    # Remaining blocks, if any, are full / not masked.
+    if masked_blocks > 0:
+        offs_n_causal = offs_n + (seqlen_q - seqlen_k) if IS_CAUSAL else 0
+        K_block_ptr = tl.advance(K_block_ptr, (0, n_full_blocks * BLOCK_N))
+        V_block_ptr = tl.advance(V_block_ptr, (n_full_blocks * BLOCK_N, 0))
+        if bias_ptr is not None:
+            bias_ptr = tl.advance(bias_ptr, (0, n_full_blocks * BLOCK_N))
+        if RETURN_ENCODED_SOFTMAX:
+            encoded_softmax_block_ptr = tl.advance(encoded_softmax_block_ptr,
+                                                   (0, n_full_blocks))
+        acc, l_i, m_i = _attn_fwd_inner(
+            acc,
+            l_i,
+            m_i,
+            q,
+            K_block_ptr,
+            V_block_ptr,
+            start_m,
+            seqlen_k,
+            dropout_p,
+            philox_seed,
+            batch_philox_offset,
+            encoded_softmax_block_ptr,
+            block_min,
+            block_max,
+            offs_n_causal,
+            masked_blocks,
+            n_extra_tokens,
+            bias_ptr,
+            IS_CAUSAL,
+            BLOCK_M,
+            BLOCK_DMODEL,
+            BLOCK_N,
+            offs_m,
+            offs_n,
+            # _, MASK_STEPS, ...
+            PRE_LOAD_V,
+            True,
+            ENABLE_DROPOUT,
+            RETURN_ENCODED_SOFTMAX,
+            padded_head,
+        )
+    # epilogue
+    acc = acc / l_i[:, None]
+    if ENABLE_DROPOUT:
+        acc = acc / (1 - dropout_p)
+    # If seqlen_q > seqlen_k but the delta is not a multiple of BLOCK_M,
+    # then we have one block with a row of all NaNs which come from computing
+    # softmax over a row of all -infs (-inf - inf = NaN). We check for that here
+    # and store 0s where there are NaNs as these rows should've been zeroed out.
+    end_m_idx = (start_m + 1) * BLOCK_M
+    start_m_idx = start_m * BLOCK_M
+    causal_start_idx = seqlen_q - seqlen_k
+    acc = acc.to(Out.type.element_ty)
+    if IS_CAUSAL:  # noqa: SIM102
+        if causal_start_idx > start_m_idx and causal_start_idx < end_m_idx:
+            out_mask_boundary = tl.full((BLOCK_DMODEL, ),
+                                        causal_start_idx,
+                                        dtype=tl.int32)
+            mask_m_offsets = start_m_idx + tl.arange(0, BLOCK_M)
+            out_ptrs_mask = (mask_m_offsets[:, None]
+                             >= out_mask_boundary[None, :])
+            z = 0.0
+            acc = tl.where(out_ptrs_mask, acc, z.to(acc.type.element_ty))
+    # write back LSE
+    # l_ptrs = L + off_z * HQ * MAX_SEQLENS_Q + off_h_q * MAX_SEQLENS_Q + offs_m
+    # If seqlen_q not multiple of BLOCK_M, we need to mask out the last
+    # few rows. This is only true for the last M block. For others,
+    # overflow_size will be -ve
+    # overflow_size = end_m_idx - seqlen_q
+    # if overflow_size > 0:
+    #    boundary = tl.full((BLOCK_M,), BLOCK_M - overflow_size, dtype=tl.int32)
+    #    # This is a > check because mask being 0 blocks the store.
+    #    l_ptrs_mask = boundary > tl.arange(0, BLOCK_M)
+    #    tl.store(l_ptrs, m_i + tl.math.log2(l_i), mask=l_ptrs_mask)
+    # else:
+    #    tl.store(l_ptrs, m_i + tl.math.log2(l_i))
+
+    # write back O
+    o_offset = (off_z * stride_oz + cu_seqlens_q_start * stride_om +
+                off_h_q * stride_oh)
+    O_block_ptr = tl.make_block_ptr(
+        base=Out + o_offset,
+        shape=(seqlen_q, ACTUAL_BLOCK_DMODEL),
+        strides=(stride_om, stride_on),
+        offsets=(start_m * BLOCK_M, 0),
+        block_shape=(BLOCK_M, BLOCK_DMODEL),
+        order=(1, 0),
+    )
+    # Need boundary check on this to make sure the padding from the
+    # Q and KV tensors in both dims are not part of what we store back.
+    # TODO: Do the boundary check optionally.
+    tl.store(O_block_ptr, acc, boundary_check=(0, 1))
+
+
+def check_args(
+    q,
+    k,
+    v,
+    o,
+    varlen=True,
+    max_seqlens=None,
+    cu_seqlens_q=None,
+    cu_seqlens_k=None,
+):
+    assert q.dim() == k.dim() and q.dim() == v.dim()
+    if varlen:
+        assert q.dim() == 3
+        total_q, nheads_q, head_size = q.shape
+        total_k, nheads_k, _ = k.shape
+        assert cu_seqlens_q is not None
+        assert cu_seqlens_k is not None
+        assert len(cu_seqlens_q) == len(cu_seqlens_k)
+    else:
+        assert q.dim() == 4
+        batch, nheads_q, seqlen_q, head_size = q.shape
+        _, nheads_k, seqlen_k, _ = k.shape
+        assert max_seqlens > 0
+    assert k.shape == v.shape
+    assert q.shape[-1] == k.shape[-1] and q.shape[-1] == v.shape[-1]
+    # TODO: Change assert if we support qkl f8 and v f16
+    assert q.dtype == k.dtype and q.dtype == v.dtype
+    assert head_size <= 256
+    assert o.shape == q.shape
+    assert (nheads_q % nheads_k) == 0
+
+
+class _attention(torch.autograd.Function):
+
+    @staticmethod
+    def forward(
+        ctx,
+        q,
+        k,
+        v,
+        o,
+        cu_seqlens_q,
+        cu_seqlens_k,
+        max_seqlens_q,
+        max_seqlens_k,
+        causal=False,
+        sm_scale=1.0,
+        bias=None,
+    ):
+        if o is None:
+            o = torch.empty_like(q, dtype=v.dtype)
+
+        check_args(
+            q,
+            k,
+            v,
+            o,
+            varlen=True,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+        )
+        if True:  # varlen
+            total_q, nheads_q, head_size = q.shape
+            total_k, nheads_k, _ = k.shape
+            batch = len(cu_seqlens_q) - 1
+            q_strides = (0, q.stride(1), q.stride(0), q.stride(2))
+            k_strides = (0, k.stride(1), k.stride(0), k.stride(2))
+            v_strides = (0, v.stride(1), v.stride(0), v.stride(2))
+            o_strides = (0, o.stride(1), o.stride(0), o.stride(2))
+        else:
+            batch, seqlen_q, nheads_q, head_size = q.shape
+            _, seqlen_k, nheads_k, _ = k.shape
+            q_strides = (q.stride(0), q.stride(2), q.stride(1), q.stride(3))
+            k_strides = (k.stride(0), k.stride(2), k.stride(1), k.stride(3))
+            v_strides = (v.stride(0), v.stride(2), v.stride(1), v.stride(3))
+            o_strides = (o.stride(0), o.stride(2), o.stride(1), o.stride(3))
+
+        # Get closest power of 2 over or equal to 32.
+        unpadded_head_dims = {32, 64, 128, 256}
+        if head_size not in unpadded_head_dims:
+            padded_d_model = None
+            for i in unpadded_head_dims:
+                if i > head_size:
+                    padded_d_model = i
+                    break
+            assert padded_d_model is not None
+        else:
+            padded_d_model = head_size
+
+        grid = lambda META: (
+            triton.cdiv(max_seqlens_q, META["BLOCK_M"]),
+            nheads_q,
+            batch,
+        )
+
+        encoded_softmax = None
+
+        # Seed the RNG so we get reproducible results for testing.
+        philox_seed = 0x1BF52
+        philox_offset = 0x1D4B42
+
+        if bias is not None:
+            bias_strides = (
+                bias.stride(0),
+                bias.stride(1),
+                bias.stride(2),
+                bias.stride(3),
+            )
+        else:
+            bias_strides = (0, 0, 0, 0)
+
+        attn_fwd[grid](
+            q,
+            k,
+            v,
+            bias,
+            sm_scale,
+            None,
+            o,
+            *q_strides,
+            *k_strides,
+            *v_strides,
+            *o_strides,
+            *bias_strides,
+            cu_seqlens_q,
+            cu_seqlens_k,
+            dropout_p=0.0,
+            philox_seed=philox_seed,
+            philox_offset_base=philox_offset,
+            encoded_softmax=encoded_softmax,
+            HQ=nheads_q,
+            HK=nheads_k,
+            ACTUAL_BLOCK_DMODEL=head_size,
+            MAX_SEQLENS_Q=max_seqlens_q,
+            MAX_SEQLENS_K=max_seqlens_k,
+            IS_CAUSAL=causal,
+            VARLEN=True,
+            BLOCK_DMODEL=padded_d_model,
+            BIAS_TYPE=0 if bias is None else 1,
+            ENABLE_DROPOUT=False,
+            RETURN_ENCODED_SOFTMAX=False,
+        )
+
+        ctx.grid = grid
+        ctx.sm_scale = sm_scale
+        ctx.BLOCK_DMODEL = head_size
+        ctx.causal = causal
+        ctx.dropout_p = 0.0
+        ctx.philox_seed = philox_seed
+        ctx.philox_offset = philox_offset
+        ctx.encoded_softmax = encoded_softmax
+        ctx.return_encoded_softmax = False
+        return o, encoded_softmax
+
+
+triton_attention = _attention.apply
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/__init__.py b/vllm_musa/vllm_musa/patch/__init__.py
new file mode 100644
index 000000000..47b1de6d2
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/__init__.py
@@ -0,0 +1,60 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+import vllm.envs as envs
+import vllm_musa.patch.patch_vllm_utils
+import vllm_musa.patch.patch_ray
+import vllm_musa.patch.patch_musa_custom_ops
+import vllm_musa.patch.patch_rotary_embedding
+import vllm_musa.patch.moe.patch_utils
+if envs.VLLM_USE_V1:
+    import vllm_musa.patch.v1.spec_decode.patch_eagle
+    import vllm_musa.patch.v1.simple_graph.patch_decorators
+    import vllm_musa.patch.v1.patch_gpu_model_runner
+    import vllm_musa.patch.v1.patch_sampler
+if envs.VLLM_USE_V1:
+    import vllm_musa.patch.v1.patch_attention_layer
+else:
+    import vllm_musa.patch.patch_attention_layer
+import vllm_musa.patch.moe.patch_fused_moe
+import vllm_musa.patch.moe.patch_fused_moe_layer
+import vllm_musa.patch.moe.patch_grouped_topk
+import vllm_musa.patch.patch_vocab_parallel_embedding
+import vllm_musa.patch.patch_linear
+import vllm_musa.patch.models.deepseek_v2
+import vllm_musa.patch.patch_parallel_state
+if envs.VLLM_USE_V1:
+    import vllm_musa.patch.v1.patch_rocm
+import vllm_musa.patch.patch_mla_common
+import vllm_musa.patch.patch_triton_mla
+import vllm_musa.patch.patch_triton_decode_attention
+if envs.VLLM_USE_V1:
+    import vllm_musa.patch.v1.patch_merge_attn_states
+    import vllm_musa.patch.v1.patch_mla_common
+    import vllm_musa.patch.v1.patch_triton_mla
+    import vllm_musa.patch.v1.patch_triton_decode_attention
+    import vllm_musa.patch.v1.patch_moe_layer
+# NOTICE: custom all-reduce no need patch to vllm, will be called from MusaCommunicator.
+# patch vllm/distributed/device_communicators/cuda_wrapper.py
+# import vllm_musa.patch.custom_allreduce.patch_cuda_wrapper
+# patch vllm/distributed/device_communicators/custom_all_reduce.py
+# import vllm_musa.patch.custom_allreduce.patch_custom_allreduce
+import vllm_musa.patch.patch_sampler
+import vllm_musa.patch.quantization.patch_fp8_utils
+import vllm_musa.patch.quantization.patch_fp8
+import vllm_musa.patch.quantization.patch_quant_utils
+import vllm_musa.patch.patch_triton_decode_attention
+import vllm_musa.patch.patch_metrics
diff --git a/vllm_musa/vllm_musa/patch/custom_allreduce/__init__.py b/vllm_musa/vllm_musa/patch/custom_allreduce/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/patch/custom_allreduce/patch_cuda_wrapper.py b/vllm_musa/vllm_musa/patch/custom_allreduce/patch_cuda_wrapper.py
new file mode 100644
index 000000000..862fb3238
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/custom_allreduce/patch_cuda_wrapper.py
@@ -0,0 +1,132 @@
+import ctypes
+from typing import Any, Dict, Optional
+
+import vllm.envs as envs
+
+from vllm.distributed.device_communicators.cuda_wrapper import (
+    cudaError_t,
+    cudaMemcpyKind,
+    cudaIpcMemHandle_t,
+    Function,
+    find_loaded_library)
+
+class MusaRTLibrary:
+    exported_functions = [
+        # cudaError_t cudaSetDevice ( int  device )
+        Function("musaSetDevice", cudaError_t, [ctypes.c_int]),
+        # cudaError_t 	cudaDeviceSynchronize ( void )
+        Function("musaDeviceSynchronize", cudaError_t, []),
+        # cudaError_t cudaDeviceReset ( void )
+        Function("musaDeviceReset", cudaError_t, []),
+
+        # const char* 	cudaGetErrorString ( cudaError_t error )
+        Function("musaGetErrorString", ctypes.c_char_p, [cudaError_t]),
+
+        # cudaError_t 	cudaMalloc ( void** devPtr, size_t size )
+        Function("musaMalloc", cudaError_t,
+                 [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t]),
+        # cudaError_t 	cudaFree ( void* devPtr )
+        Function("musaFree", cudaError_t, [ctypes.c_void_p]),
+        # cudaError_t cudaMemset ( void* devPtr, int  value, size_t count )
+        Function("musaMemset", cudaError_t,
+                 [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]),
+        # cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa
+        Function("musaMemcpy", cudaError_t, [
+            ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, cudaMemcpyKind
+        ]),
+
+        # cudaError_t cudaIpcGetMemHandle ( cudaIpcMemHandle_t* handle, void* devPtr ) # noqa
+        Function("musaIpcGetMemHandle", cudaError_t,
+                 [ctypes.POINTER(cudaIpcMemHandle_t), ctypes.c_void_p]),
+        # cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int  flags ) # noqa
+        Function("musaIpcOpenMemHandle", cudaError_t, [
+            ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint
+        ]),
+    ]
+
+    # class attribute to store the mapping from the path to the library
+    # to avoid loading the same library multiple times
+    path_to_library_cache: Dict[str, Any] = {}
+
+    # class attribute to store the mapping from library path
+    #  to the corresponding dictionary
+    path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}
+
+    def __init__(self, so_file: Optional[str] = None):
+        if so_file is None:
+            so_file = find_loaded_library("libmusart")
+            if so_file is None:
+                so_file = envs.VLLM_CUDART_SO_PATH  # fallback to env var
+            assert so_file is not None, \
+                (
+                    "libcudart is not loaded in the current process, "
+                    "try setting VLLM_CUDART_SO_PATH"
+                )
+        if so_file not in MusaRTLibrary.path_to_library_cache:
+            lib = ctypes.CDLL(so_file)
+            MusaRTLibrary.path_to_library_cache[so_file] = lib
+        self.lib = MusaRTLibrary.path_to_library_cache[so_file]
+
+        if so_file not in MusaRTLibrary.path_to_dict_mapping:
+            _funcs = {}
+            for func in MusaRTLibrary.exported_functions:
+                f = getattr(self.lib, func.name)
+                f.restype = func.restype
+                f.argtypes = func.argtypes
+                _funcs[func.name] = f
+            MusaRTLibrary.path_to_dict_mapping[so_file] = _funcs
+        self.funcs = MusaRTLibrary.path_to_dict_mapping[so_file]
+
+    def CUDART_CHECK(self, result: cudaError_t) -> None:
+        if result != 0:
+            error_str = self.cudaGetErrorString(result)
+            raise RuntimeError(f"CUDART error: {error_str}")
+
+    def cudaGetErrorString(self, error: cudaError_t) -> str:
+        return self.funcs["musaGetErrorString"](error).decode("utf-8")
+
+    def cudaSetDevice(self, device: int) -> None:
+        self.CUDART_CHECK(self.funcs["musaSetDevice"](device))
+
+    def cudaDeviceSynchronize(self) -> None:
+        self.CUDART_CHECK(self.funcs["musaDeviceSynchronize"]())
+
+    def cudaDeviceReset(self) -> None:
+        self.CUDART_CHECK(self.funcs["musaDeviceReset"]())
+
+    def cudaMalloc(self, size: int) -> ctypes.c_void_p:
+        devPtr = ctypes.c_void_p()
+        self.CUDART_CHECK(self.funcs["musaMalloc"](ctypes.byref(devPtr), size))
+        return devPtr
+
+    def cudaFree(self, devPtr: ctypes.c_void_p) -> None:
+        self.CUDART_CHECK(self.funcs["musaFree"](devPtr))
+
+    def cudaMemset(self, devPtr: ctypes.c_void_p, value: int,
+                   count: int) -> None:
+        self.CUDART_CHECK(self.funcs["musaMemset"](devPtr, value, count))
+
+    def cudaMemcpy(self, dst: ctypes.c_void_p, src: ctypes.c_void_p,
+                   count: int) -> None:
+        cudaMemcpyDefault = 4
+        kind = cudaMemcpyDefault
+        self.CUDART_CHECK(self.funcs["musaMemcpy"](dst, src, count, kind))
+
+    def cudaIpcGetMemHandle(self,
+                            devPtr: ctypes.c_void_p) -> cudaIpcMemHandle_t:
+        handle = cudaIpcMemHandle_t()
+        self.CUDART_CHECK(self.funcs["musaIpcGetMemHandle"](
+            ctypes.byref(handle), devPtr))
+        return handle
+
+    def cudaIpcOpenMemHandle(self,
+                             handle: cudaIpcMemHandle_t) -> ctypes.c_void_p:
+        cudaIpcMemLazyEnablePeerAccess = 1
+        devPtr = ctypes.c_void_p()
+        self.CUDART_CHECK(self.funcs["musaIpcOpenMemHandle"](
+            ctypes.byref(devPtr), handle, cudaIpcMemLazyEnablePeerAccess))
+        return devPtr
+
+# NOTICE: custom all-reduce no need patch to vllm, will be called from MusaCommunicator.
+# import vllm.distributed.device_communicators.cuda_wrapper
+# vllm.distributed.device_communicators.cuda_wrapper.CudaRTLibrary = MusaRTLibrary
diff --git a/vllm_musa/vllm_musa/patch/custom_allreduce/patch_custom_allreduce.py b/vllm_musa/vllm_musa/patch/custom_allreduce/patch_custom_allreduce.py
new file mode 100644
index 000000000..570201c15
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/custom_allreduce/patch_custom_allreduce.py
@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import ctypes
+from contextlib import contextmanager
+from typing import List, Optional, Union
+
+import torch
+import torch_musa
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from vllm_musa import _musa_custom_ops as ops
+from vllm_musa.patch.custom_allreduce.patch_cuda_wrapper import MusaRTLibrary
+from vllm.logger import init_logger
+
+try:
+    ops.meta_size()
+    custom_ar = True
+except Exception:
+    custom_ar = False
+
+logger = init_logger(__name__)
+
+def is_weak_contiguous(inp: torch.Tensor):
+    return inp.is_contiguous() or (inp.storage().nbytes() -
+                                   inp.storage_offset() * inp.element_size()
+                                   == inp.numel() * inp.element_size())
+
+class CustomAllreduce:
+
+    _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
+
+    # max_size: max supported allreduce size
+    def __init__(self,
+                 group: ProcessGroup,
+                 device: Union[int, str, torch.device],
+                 max_size=8192 * 1024) -> None:
+        """
+        Args:
+            group: the process group to work on. If None, it will use the
+                default process group.
+            device: the device to bind the CustomAllreduce to. If None,
+                it will be bind to f"musa:{local_rank}".
+        It is the caller's responsibility to make sure each communicator
+        is bind to a unique device, and all communicators in this group
+        are in the same node.
+        """
+        self._IS_CAPTURING = False
+        self.disabled = True
+
+        if not custom_ar:
+            # disable because of missing custom allreduce library
+            # e.g. in a non-musa environment
+            return
+
+        self.group = group
+        assert dist.get_backend(group) != dist.Backend.MCCL, (
+            "CustomAllreduce should be attached to a non-MCCL group.")
+
+        rank = dist.get_rank(group=self.group)
+        world_size = dist.get_world_size(group=self.group)
+        if world_size == 1:
+            # No need to initialize custom allreduce for single GPU case.
+            return
+        if world_size not in CustomAllreduce._SUPPORTED_WORLD_SIZES:
+            logger.warning(
+                "Custom allreduce is disabled due to an unsupported world"
+                " size: %d. Supported world sizes: %s. To silence this "
+                "warning, specify disable_custom_all_reduce=True explicitly.",
+                world_size, str(CustomAllreduce._SUPPORTED_WORLD_SIZES))
+            return
+        if isinstance(device, int):
+            device = torch.device(f"musa:{device}")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        # now `device` is a `torch.device` object
+        assert isinstance(device, torch.device)
+        self.device = device
+
+        self.disabled = False
+        # Buffers memory are owned by this Python class and passed to C++.
+        # Meta data composes of two parts: meta data for synchronization and a
+        # temporary buffer for storing intermediate allreduce results.
+        self.meta_ptrs = self.create_shared_buffer(ops.meta_size() + max_size,
+                                                   group=group)
+        # This is a pre-registered IPC buffer. In eager mode, input tensors
+        # are first copied into this buffer before allreduce is performed
+        self.buffer_ptrs = self.create_shared_buffer(max_size, group=group)
+        # This is a buffer for storing the tuples of pointers pointing to
+        # IPC buffers from all ranks. Each registered tuple has size of
+        # 8*world_size bytes where world_size is at most 8. Allocating 8MB
+        # is enough for 131072 such tuples. The largest model I've seen only
+        # needs less than 10000 of registered tuples.
+        self.rank_data = torch.empty(8 * 1024 * 1024,
+                                     dtype=torch.uint8,
+                                     device=self.device)
+        self.max_size = max_size
+        self.rank = rank
+        self.world_size = world_size
+        self.full_nvlink = True
+        self._ptr = ops.init_custom_ar(self.meta_ptrs, self.rank_data, rank,
+                                       self.full_nvlink)
+        ops.register_buffer(self._ptr, self.buffer_ptrs)
+
+    @staticmethod
+    def create_shared_buffer(
+            size_in_bytes: int,
+            group: Optional[ProcessGroup] = None) -> List[int]:
+        """
+        Creates a shared buffer and returns a list of pointers
+        representing the buffer on all processes in the group.
+        """
+        lib = MusaRTLibrary()
+        pointer = lib.cudaMalloc(size_in_bytes)
+        lib.cudaMemset(pointer, 0, size_in_bytes)
+        handle = lib.cudaIpcGetMemHandle(pointer)
+        world_size = dist.get_world_size(group=group)
+        rank = dist.get_rank(group=group)
+        handles = [None] * world_size
+        dist.all_gather_object(handles, handle, group=group)
+
+        pointers: List[int] = []
+        for i, h in enumerate(handles):
+            if i == rank:
+                pointers.append(pointer.value)  # type: ignore
+            else:
+                pointers.append(
+                    lib.cudaIpcOpenMemHandle(h).value)  # type: ignore
+
+        return pointers
+
+    @staticmethod
+    def free_shared_buffer(pointers: List[int],
+                           group: Optional[ProcessGroup] = None) -> None:
+        rank = dist.get_rank(group=group)
+        lib = MusaRTLibrary()
+        lib.cudaFree(ctypes.c_void_p(pointers[rank]))
+
+    @contextmanager
+    def capture(self):
+        """
+        The main responsibility of this context manager is the
+        `register_graph_buffers` call at the end of the context.
+        It records all the buffer addresses used in the MUSA graph.
+        """
+        try:
+            self._IS_CAPTURING = True
+            yield
+        finally:
+            self._IS_CAPTURING = False
+            if not self.disabled:
+                self.register_graph_buffers()
+
+
+    def register_graph_buffers(self):
+        handle, offset = ops.get_graph_buffer_ipc_meta(self._ptr)
+        logger.info("Registering %d musa graph addresses", len(offset))
+        # We cannot directly use `dist.all_gather_object` here
+        # because it is incompatible with `gloo` backend under inference mode.
+        # see https://github.com/pytorch/pytorch/issues/126032 for details.
+        all_data = [[None, None]
+                    for _ in range(dist.get_world_size(group=self.group))]
+        all_data[self.rank] = [handle, offset]
+        ranks = sorted(dist.get_process_group_ranks(group=self.group))
+        for i, rank in enumerate(ranks):
+            dist.broadcast_object_list(all_data[i],
+                                       src=rank,
+                                       group=self.group,
+                                       device="cpu")
+        # Unpack list of tuples to tuple of lists.
+        handles = [d[0] for d in all_data]  # type: ignore
+        offsets = [d[1] for d in all_data]  # type: ignore
+        ops.register_graph_buffers(self._ptr, handles, offsets)
+
+    def should_custom_ar(self, inp: torch.Tensor):
+        if self.disabled:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # custom allreduce requires input byte size to be multiples of 16
+        if inp_size % 16 != 0:
+            return False
+        if not is_weak_contiguous(inp):
+            return False
+        # for 4 or more non NVLink-capable GPUs, custom allreduce provides
+        # little performance improvement over NCCL.
+        if self.world_size >= 2 or self.full_nvlink:
+            return inp_size < self.max_size
+        return False
+
+    def all_reduce(self,
+                   inp: torch.Tensor,
+                   *,
+                   out: torch.Tensor = None,
+                   registered: bool = False):
+        """Performs an out-of-place all reduce.
+
+        If registered is True, this assumes inp's pointer is already
+        IPC-registered. Otherwise, inp is first copied into a pre-registered
+        buffer.
+        """
+        if out is None:
+            out = torch.empty_like(inp)
+        if registered:
+            ops.all_reduce(self._ptr, inp, out, 0, 0)
+        else:
+            ops.all_reduce(self._ptr, inp, out, self.buffer_ptrs[self.rank],
+                           self.max_size)
+        return out
+
+    def custom_all_reduce(self, input: torch.Tensor) -> Optional[torch.Tensor]:
+        """The main allreduce API that provides support for musa graph."""
+        # When custom allreduce is disabled, this will be None.
+        if self.disabled or not self.should_custom_ar(input):
+            return None
+        if self._IS_CAPTURING:
+            # try:
+            #     return self.all_reduce(input, registered=True)
+            # except Exception as exc:
+            #     print(exc)
+            #     return torch.empty_like(input)
+            if torch.musa.is_current_stream_capturing():
+                return self.all_reduce(input, registered=True)
+            else:
+                # If warm up, mimic the allocation pattern since custom
+                # allreduce is out-of-place.
+                return torch.empty_like(input)
+        else:
+            # Note: outside of musa graph context, custom allreduce incurs a
+            # cost of musaMemcpy, which should be small (<=1% of overall
+            # latency) compared to the performance gain of using custom kernels
+            return self.all_reduce(input, registered=False)
+
+    def close(self):
+        if not self.disabled and self._ptr:
+            ops.dispose(self._ptr)
+            self._ptr = 0
+            self.free_shared_buffer(self.meta_ptrs)
+            self.free_shared_buffer(self.buffer_ptrs)
+
+    def __del__(self):
+        self.close()
+
+# NOTICE: custom all-reduce no need patch to vllm, will be called from MusaCommunicator.
+# import vllm.distributed.device_communicators.custom_all_reduce
+# vllm.distributed.device_communicators.custom_all_reduce.CustomAllreduce = CustomAllreduce
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/models/__init__.py b/vllm_musa/vllm_musa/patch/models/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/patch/models/deepseek_v2.py b/vllm_musa/vllm_musa/patch/models/deepseek_v2.py
new file mode 100644
index 000000000..b40dfdcd8
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/models/deepseek_v2.py
@@ -0,0 +1,245 @@
+from typing import Any, Optional, Union
+
+import torch
+from torch import nn
+
+from transformers import PretrainedConfig
+
+import vllm
+from vllm.config import (CacheConfig, ModelConfig,
+                         get_current_vllm_config)
+from vllm.distributed import (get_tensor_model_parallel_world_size,
+                              get_ep_group,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.fused_moe import FusedMoE
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (ReplicatedLinear,
+                                               RowParallelLinear)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.fp8 import (Fp8Config, Fp8LinearMethod)
+from vllm.model_executor.models.deepseek_v2 import (DeepseekV2DecoderLayer,
+                                                    DeepseekV2MLAAttention)
+
+from vllm_musa._musa_custom_ops import (fused_mul_add, gemm_swiglu, gemm_swiglu_w4a16, fused_gemv)
+from vllm_musa.patch.patch_linear import MergedColumnParallelLinear
+from vllm_musa.patch.quantization.patch_fp8_utils import apply_w8a8_block_fp8_linear
+
+FUSED_RMS_NROM = False
+
+class CustomDeepseekV2MLP(nn.Module):
+
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        hidden_act: str,
+        quant_config: Optional[QuantizationConfig] = None,
+        reduce_results: bool = True,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.gate_up_proj = MergedColumnParallelLinear(
+            hidden_size, [intermediate_size] * 2,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.gate_up_proj",
+            save_transpose_weight=True)
+        self.down_proj = RowParallelLinear(intermediate_size,
+                                           hidden_size,
+                                           bias=False,
+                                           quant_config=quant_config,
+                                           reduce_results=reduce_results,
+                                           prefix=f"{prefix}.down_proj")
+        if hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {hidden_act}. "
+                             "Only silu is supported for now.")
+        self.act_fn = torch.nn.SwishGLU()
+
+    def forward(self, x):
+        batch = x.shape[0]
+        if (batch >= 4):
+            # TODO: 此部分考虑后续并发较高的情况，根据算子性能保留项，但是此处会增加一份weight，后续可能会删除此部分推理代码
+            gate_up, _ = self.gate_up_proj(x)
+            x = self.act_fn(gate_up)
+            x, _ = self.down_proj(x)
+        else:
+            if isinstance(self.gate_up_proj.quant_method, Fp8LinearMethod):
+                out = apply_w8a8_block_fp8_linear(x,
+                                                  self.gate_up_proj.weight,
+                                                  self.gate_up_proj.quant_method.quant_config.weight_block_size,
+                                                  self.gate_up_proj.weight_scale_inv,
+                                                  use_swigelu=True)
+                x, _ = self.down_proj(out)
+            elif (hasattr(self.gate_up_proj, "qweight")): # w4a16
+                if self.gate_up_proj.quant_method.quant_config.group_size != -1:
+                    x = fused_gemv(
+                        x,
+                        self.gate_up_proj.qweight,
+                        None,
+                        self.gate_up_proj.scales,
+                        use_swigelu=True)
+                    x, _ = self.down_proj(x)
+                else:
+                    out_dim = self.gate_up_proj.qweight.shape[0] // 2
+                    out = torch.empty(batch, out_dim, device=x.device, dtype=x.dtype)
+                    gemm_swiglu_w4a16(out, x, self.gate_up_proj.qweight_t, self.gate_up_proj.scales)
+                    x, _ = self.down_proj(out)
+            else:   #bf16
+                if  (hasattr(self.gate_up_proj, "fused_swiglu_weight")):
+                    out_dim = self.gate_up_proj.weight.shape[0] // 2
+                    out = torch.empty(batch, out_dim, device=x.device, dtype=x.dtype)
+                    gemm_swiglu(out, x, self.gate_up_proj.fused_swiglu_weight, alpha=1.0)
+                    x, _ = self.down_proj(out)
+                else:
+                    gate_up, _ = self.gate_up_proj(x)
+                    x = self.act_fn(gate_up)
+                    x, _ = self.down_proj(x)
+        return x
+
+class CustomDeepseekV2MoE(nn.Module):
+
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        quant_config: Optional[QuantizationConfig] = None,
+        prefix: str = "",
+        enable_eplb: bool = False,
+    ):
+        super().__init__()
+        self.tp_size = get_tensor_model_parallel_world_size()
+        self.routed_scaling_factor = config.routed_scaling_factor
+
+        self.ep_group = get_ep_group().device_group
+        self.ep_rank = self.ep_group.rank()
+        self.ep_size = self.ep_group.size()
+        self.n_routed_experts: int = config.n_routed_experts
+        self.n_shared_experts: int = config.n_shared_experts
+
+        if config.hidden_act != "silu":
+            raise ValueError(f"Unsupported activation: {config.hidden_act}. "
+                             "Only silu is supported for now.")
+
+        self.gate = ReplicatedLinear(config.hidden_size,
+                                     config.n_routed_experts,
+                                     bias=False,
+                                     quant_config=None,
+                                     prefix=f"{prefix}.gate")
+        if config.topk_method == "noaux_tc":
+            self.gate.e_score_correction_bias = nn.Parameter(
+                torch.empty(config.n_routed_experts))
+        else:
+            self.gate.e_score_correction_bias = None
+
+        # Load balancing settings.
+        vllm_config = get_current_vllm_config()
+        parallel_config = vllm_config.parallel_config
+        self.enable_eplb = enable_eplb
+
+        self.n_redundant_experts = parallel_config.num_redundant_experts
+        self.n_logical_experts = self.n_routed_experts
+        self.n_physical_experts = (self.n_logical_experts +
+                                   self.n_redundant_experts)
+        self.n_local_physical_experts = self.n_physical_experts // self.ep_size
+
+        self.physical_expert_start = (self.ep_rank *
+                                      self.n_local_physical_experts)
+        self.physical_expert_end = (self.physical_expert_start +
+                                    self.n_local_physical_experts)
+
+        self.experts = FusedMoE(
+            num_experts=config.n_routed_experts,
+            top_k=config.num_experts_per_tok,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.moe_intermediate_size,
+            reduce_results=False,
+            renormalize=config.norm_topk_prob,
+            quant_config=quant_config,
+            use_grouped_topk=True,
+            num_expert_group=config.n_group,
+            topk_group=config.topk_group,
+            prefix=f"{prefix}.experts",
+            scoring_func=config.scoring_func,
+            e_score_correction_bias=self.gate.e_score_correction_bias,
+            enable_eplb=self.enable_eplb,
+            num_redundant_experts=self.n_redundant_experts)
+
+        if config.n_shared_experts is not None:
+            intermediate_size = (config.moe_intermediate_size *
+                                 config.n_shared_experts)
+            self.shared_experts = CustomDeepseekV2MLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                reduce_results=self.experts.must_reduce_shared_expert_outputs(
+                ),
+                prefix=f"{prefix}.shared_experts",
+            )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        num_tokens, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        shared_output = self.shared_experts(hidden_states)
+        # router_logits: (num_tokens, n_experts)
+        router_logits, _ = self.gate(hidden_states)
+        final_hidden_states = self.experts(
+            hidden_states=hidden_states,
+            router_logits=router_logits)
+        final_hidden_states = fused_mul_add(
+            final_hidden_states,
+            shared_output,
+            self.routed_scaling_factor,
+            accurate=False  # set to true if meets precision problem
+        )
+        if self.tp_size > 1:
+            final_hidden_states = tensor_model_parallel_all_reduce(
+                final_hidden_states)
+
+        return final_hidden_states.view(num_tokens, hidden_dim)
+
+class CustomDeepseekV2MLAAttention(DeepseekV2MLAAttention):
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.q_lora_rank is not None:
+            if isinstance(self.q_a_proj.quant_method, Fp8LinearMethod) and FUSED_RMS_NROM:
+                q_c = apply_w8a8_block_fp8_linear(
+                    hidden_states,
+                    self.q_a_proj.weight,
+                    self.q_a_proj.quant_method.quant_config.weight_block_size,
+                    self.q_a_proj.weight_scale_inv,
+                    use_rms_norm=True,
+                    gamma=self.q_a_layernorm.weight,
+                    eps=self.q_a_layernorm.variance_epsilon,
+                )
+            else:
+                q_c = self.q_a_proj(hidden_states)[0]
+                q_c = self.q_a_layernorm(q_c)
+            q = self.q_b_proj(q_c)[0]
+        else:
+            q = self.q_proj(hidden_states)[0]
+        kv_c, k_pe = self.kv_a_proj_with_mqa(hidden_states)[0].split(
+            [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        kv_c_normed = self.kv_a_layernorm(kv_c.contiguous())
+
+        q = q.view(-1, self.num_local_heads, self.qk_head_dim)
+        # Add head dim of 1 to k_pe
+        k_pe = k_pe.unsqueeze(1)
+
+        q[..., self.qk_nope_head_dim:], k_pe = self.rotary_emb(
+            positions, q[..., self.qk_nope_head_dim:], k_pe)
+
+        attn_out = self.mla_attn(
+            q,
+            kv_c_normed,
+            k_pe,
+            output_shape=(hidden_states.shape[0],
+                          self.num_local_heads * self.v_head_dim))
+        return self.o_proj(attn_out)[0]
+
+
+vllm.model_executor.models.deepseek_v2.DeepseekV2MLP = CustomDeepseekV2MLP
+vllm.model_executor.models.deepseek_v2.DeepseekV2MoE = CustomDeepseekV2MoE
+vllm.model_executor.models.deepseek_v2.DeepseekV2MLAAttention = CustomDeepseekV2MLAAttention
diff --git a/vllm_musa/vllm_musa/patch/moe/__init__.py b/vllm_musa/vllm_musa/patch/moe/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/patch/moe/patch_fused_moe.py b/vllm_musa/vllm_musa/patch/moe/patch_fused_moe.py
new file mode 100644
index 000000000..c3fff628c
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/moe/patch_fused_moe.py
@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+import torch
+import functools
+import vllm.envs as envs
+import triton.language as tl
+import torch.nn.functional as F
+from typing import Any, List, Optional, Sequence, Dict
+from vllm_musa import _musa_custom_ops as ops
+from vllm.model_executor.layers.fused_moe.fused_moe import (get_config_dtype_str,
+                                                            get_config_quant_dtype,
+                                                            try_get_optimal_moe_config,
+                                                            invoke_fused_moe_kernel)
+from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input
+from vllm_musa.ops.moe_align_block_size import moe_align_block_size
+from vllm_musa.ops.deep_gemm import (deep_gemm_moe_fp8,
+                                     _valid_deep_gemm)
+from vllm_musa.platforms.musa import on_qy2, on_ph1
+
+
+def should_use_gemv_fused_moe(
+    use_fp8_w8a8: bool,
+    num_tokens: int,
+    expert_map: Optional[Any],
+    enable_triton_moe: bool
+) -> bool:
+    if enable_triton_moe:
+        return False
+
+    if use_fp8_w8a8:
+        return True
+
+    if on_qy2():
+        return num_tokens < 32 or expert_map is not None
+
+    if on_ph1():
+        return True
+
+    return True
+
+def should_moe_wna16_use_cuda(num_valid_tokens: int, group_size: int,
+                              num_experts: int, bit: int):
+    return False
+
+def fused_experts_impl(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    inplace: bool = False,
+    activation: str = "silu",
+    apply_router_weight_on_input: bool = False,
+    use_fp8_w8a8: bool = False,
+    use_int8_w8a8: bool = False,
+    use_int8_w8a16: bool = False,
+    use_int4_w4a16: bool = False,
+    per_channel_quant: bool = False,
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
+    w1_scale: Optional[torch.Tensor] = None,
+    w2_scale: Optional[torch.Tensor] = None,
+    w1_zp: Optional[torch.Tensor] = None,
+    w2_zp: Optional[torch.Tensor] = None,
+    a1_scale: Optional[torch.Tensor] = None,
+    a2_scale: Optional[torch.Tensor] = None,
+    block_shape: Optional[List[int]] = None,
+    use_gemv_fused_moe: bool = False
+) -> torch.Tensor:
+    if (use_fp8_w8a8
+            and _valid_deep_gemm(hidden_states, w1, w2)):
+        assert apply_router_weight_on_input is False
+        return deep_gemm_moe_fp8(
+            hidden_states=hidden_states,
+            w1=w1,
+            w2=w2,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            inplace=inplace,
+            activation=activation,
+            global_num_experts=global_num_experts,
+            expert_map=expert_map,
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+        )
+
+    # Check constraints.
+    if use_int4_w4a16:
+        assert hidden_states.shape[1] // 2 == w1.shape[
+            2], "Hidden size mismatch"
+    else:
+        assert hidden_states.shape[1] == w1.shape[2], (
+            f"Hidden size mismatch {hidden_states.shape[1]} != {w1.shape[2]}")
+
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
+    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+    assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
+    assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
+    assert hidden_states.dtype in [
+        torch.float32, torch.float16, torch.bfloat16
+    ]
+
+    num_tokens = hidden_states.shape[0]
+    E, N, _ = w1.shape
+    K = w2.shape[1]
+    if global_num_experts == -1:
+        global_num_experts = E
+    top_k_num = topk_ids.shape[1]
+    # We execute the fused_moe kernel in chunks to circumvent this issue:
+    # https://github.com/vllm-project/vllm/issues/5938
+    CHUNK_SIZE = envs.VLLM_FUSED_MOE_CHUNK_SIZE
+    M = min(num_tokens, CHUNK_SIZE)
+    config_dtype = get_config_dtype_str(use_fp8_w8a8=use_fp8_w8a8,
+                                        use_int8_w8a16=use_int8_w8a16,
+                                        use_int4_w4a16=use_int4_w4a16,
+                                        dtype=hidden_states.dtype)
+
+    qtype = get_config_quant_dtype(use_fp8_w8a8=use_fp8_w8a8,
+                                   use_int8_w8a8=use_int8_w8a8,
+                                   use_int8_w8a16=use_int8_w8a16,
+                                   use_int4_w4a16=use_int4_w4a16)
+
+    get_config_func = functools.partial(
+        try_get_optimal_moe_config,
+        w1.shape,
+        w2.shape,
+        top_k_num,
+        config_dtype,
+        block_shape=block_shape,
+    )
+
+    config = get_config_func(M)
+
+    # We can reuse the memory between these because by the time we need
+    # cache3, we're done with cache1
+    cache13 = torch.empty(M * top_k_num * max(N, K),
+                          device=hidden_states.device,
+                          dtype=hidden_states.dtype)
+    intermediate_cache1 = cache13[:M * top_k_num * N].view(M, top_k_num, N)
+    intermediate_cache3 = cache13[:M * top_k_num * K].view(M, top_k_num, K)
+
+    # This needs separate memory since it's used concurrently with cache1
+    intermediate_cache2 = torch.empty((M * top_k_num, N // 2),
+                                      device=hidden_states.device,
+                                      dtype=hidden_states.dtype)
+
+    if hidden_states.dtype == torch.bfloat16:
+        compute_type = tl.bfloat16
+    elif hidden_states.dtype == torch.float16:
+        compute_type = tl.float16
+    elif hidden_states.dtype == torch.float32:
+        compute_type = tl.float32
+    else:
+        raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
+
+    if inplace:
+        out_hidden_states = hidden_states
+    else:
+        out_hidden_states = torch.empty_like(hidden_states) # TODO or torch.zeros_like?
+
+    for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+        begin_chunk_idx, end_chunk_idx = (chunk * CHUNK_SIZE,
+                                          min((chunk + 1) * CHUNK_SIZE,
+                                              num_tokens))
+        curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+        tokens_in_chunk, _ = curr_hidden_states.shape
+
+        if tokens_in_chunk == 0:
+            break
+
+        if tokens_in_chunk < CHUNK_SIZE and chunk > 0:
+            # Adjust the intermediate cache size and config for the last
+            # chunk. Note that in most cases we only have one chunk
+            # so the cache size and config are already set correctly and
+            # do not need to be adjusted.
+            intermediate_cache1 = intermediate_cache1[:tokens_in_chunk]
+            intermediate_cache2 = intermediate_cache2[:tokens_in_chunk *
+                                                      topk_ids.shape[1]]
+            intermediate_cache3 = intermediate_cache3[:tokens_in_chunk]
+            config = get_config_func(tokens_in_chunk)
+
+        config = dict()
+        config["BLOCK_SIZE_M"] = 32
+        config["BLOCK_SIZE_N"] = 32
+        config["BLOCK_SIZE_K"] = 64
+        config["GROUP_SIZE_M"] = 1
+
+        curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+        curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+
+        # 3 fused_moe kernels for performance concern:
+        #   * gemv
+        #   * mutlass
+        #   * triton
+        ENABLE_TRITON_MOE = bool(int(os.getenv("VLLM_MUSA_ENABLE_MOE_TRITON", "0")))
+        use_gemv_fused_moe = should_use_gemv_fused_moe(
+            use_fp8_w8a8=use_fp8_w8a8,
+            num_tokens=num_tokens,
+            expert_map=expert_map,
+            enable_triton_moe=ENABLE_TRITON_MOE
+        )
+        if not use_gemv_fused_moe:
+            qcurr_hidden_states, a1q_scale = moe_kernel_quantize_input(
+                A=curr_hidden_states,
+                A_scale=a1_scale,
+                quant_dtype=qtype,
+                per_act_token_quant=per_channel_quant,
+                block_shape=block_shape)
+
+        if use_gemv_fused_moe:
+            sorted_token_ids = torch.empty(0, dtype=curr_topk_ids.dtype)
+            expert_ids = torch.empty(0, dtype=curr_topk_ids.dtype)
+            num_tokens_post_padded = torch.empty(0, dtype=curr_topk_ids.dtype)
+            if expert_map is not None: # For EP
+                curr_topk_ids = expert_map[curr_topk_ids]
+        else: # use triton fused moe or mutlass fused moe
+            sorted_token_ids, expert_ids, num_tokens_post_padded = (
+                moe_align_block_size(curr_topk_ids, config['BLOCK_SIZE_M'],
+                                     global_num_experts, expert_map))
+
+        if not ENABLE_TRITON_MOE:
+            # w16a16, w4a16 per-channel and w4a16 per-group musa gemv kernel
+            if use_gemv_fused_moe:
+                ops.fused_moe_gemv(curr_hidden_states,
+                                    w1,
+                                    intermediate_cache2,
+                                    None,
+                                    w1_scale,
+                                    curr_topk_weights,
+                                    curr_topk_ids,
+                                    sorted_token_ids,
+                                    expert_ids,
+                                    num_tokens_post_padded,
+                                    apply_router_weight_on_input,
+                                    top_k_num,
+                                    use_int4_w4a16,
+                                    use_swigelu=True)
+            # w16a16, w4a16 per-channel and w4a16 per-group mutlass gemm kernel
+            else:
+                ops.fused_moe_mutlass(
+                                    curr_hidden_states,
+                                    w1,
+                                    intermediate_cache1,
+                                    None,
+                                    w1_scale,
+                                    curr_topk_weights,
+                                    curr_topk_ids,
+                                    sorted_token_ids,
+                                    expert_ids,
+                                    num_tokens_post_padded,
+                                    apply_router_weight_on_input,
+                                    top_k_num,
+                                    use_int4_w4a16,
+                                    config['BLOCK_SIZE_M'],
+                                    config['BLOCK_SIZE_N'],
+                                    config['BLOCK_SIZE_K'],
+                                    config['GROUP_SIZE_M'],
+                                    block_shape,
+                                    )
+        else:
+            invoke_fused_moe_kernel(qcurr_hidden_states,
+                                    w1,
+                                    intermediate_cache1,
+                                    a1q_scale,
+                                    w1_scale,
+                                    w1_zp,
+                                    curr_topk_weights,
+                                    sorted_token_ids,
+                                    expert_ids,
+                                    num_tokens_post_padded,
+                                    apply_router_weight_on_input,
+                                    topk_ids.shape[1],
+                                    config,
+                                    compute_type=compute_type,
+                                    use_fp8_w8a8=use_fp8_w8a8,
+                                    use_int8_w8a8=use_int8_w8a8,
+                                    use_int8_w8a16=use_int8_w8a16,
+                                    use_int4_w4a16=use_int4_w4a16,
+                                    per_channel_quant=per_channel_quant,
+                                    block_shape=block_shape)
+
+        if not use_gemv_fused_moe:
+            intermediate_cache2 = F.swish_glu(intermediate_cache1.view(-1, N))
+            qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
+                A=intermediate_cache2,
+                A_scale=a2_scale,
+                quant_dtype=qtype,
+                per_act_token_quant=per_channel_quant,
+                block_shape=block_shape)
+
+        if not ENABLE_TRITON_MOE:
+            # w16a16, w4a16 per-channel and w4a16 per-group musa gemv kernel
+            if use_gemv_fused_moe:
+                ops.fused_moe_gemv(intermediate_cache2,
+                                    w2,
+                                    intermediate_cache3,
+                                    None,
+                                    w2_scale,
+                                    curr_topk_weights,
+                                    curr_topk_ids,
+                                    sorted_token_ids,
+                                    expert_ids,
+                                    num_tokens_post_padded,
+                                    not apply_router_weight_on_input,
+                                    1,
+                                    use_int4_w4a16,
+                                    use_swigelu=False)
+            # w16a16, w4a16 per-channel and w4a16 per-group mutlass gemm kernel
+            else:
+                ops.fused_moe_mutlass(
+                                    intermediate_cache2,
+                                    w2,
+                                    intermediate_cache3,
+                                    None,
+                                    w2_scale,
+                                    curr_topk_weights,
+                                    curr_topk_ids,
+                                    sorted_token_ids,
+                                    expert_ids,
+                                    num_tokens_post_padded,
+                                    not apply_router_weight_on_input,
+                                    1,
+                                    use_int4_w4a16,
+                                    config['BLOCK_SIZE_M'],
+                                    config['BLOCK_SIZE_N'],
+                                    config['BLOCK_SIZE_K'],
+                                    config['GROUP_SIZE_M'],
+                                    block_shape,
+                                    )
+        else:
+            invoke_fused_moe_kernel(qintermediate_cache2,
+                                    w2,
+                                    intermediate_cache3,
+                                    a2q_scale,
+                                    w2_scale,
+                                    w2_zp,
+                                    curr_topk_weights,
+                                    sorted_token_ids,
+                                    expert_ids,
+                                    num_tokens_post_padded,
+                                    not apply_router_weight_on_input,
+                                    1,
+                                    config,
+                                    compute_type=compute_type,
+                                    use_fp8_w8a8=use_fp8_w8a8,
+                                    use_int8_w8a8=use_int8_w8a8,
+                                    use_int8_w8a16=use_int8_w8a16,
+                                    use_int4_w4a16=use_int4_w4a16,
+                                    per_channel_quant=per_channel_quant,
+                                    block_shape=block_shape)
+
+        ops.moe_sum(intermediate_cache3.view(*intermediate_cache3.shape),
+                    out_hidden_states[begin_chunk_idx:end_chunk_idx])
+
+    return out_hidden_states
+
+
+def inplace_fused_experts(hidden_states: torch.Tensor,
+                          w1: torch.Tensor,
+                          w2: torch.Tensor,
+                          topk_weights: torch.Tensor,
+                          topk_ids: torch.Tensor,
+                          activation: str = "silu",
+                          apply_router_weight_on_input: bool = False,
+                          use_fp8_w8a8: bool = False,
+                          use_int8_w8a8: bool = False,
+                          use_int8_w8a16: bool = False,
+                          use_int4_w4a16: bool = False,
+                          per_channel_quant: bool = False,
+                          global_num_experts: int = -1,
+                          expert_map: Optional[torch.Tensor] = None,
+                          w1_scale: Optional[torch.Tensor] = None,
+                          w2_scale: Optional[torch.Tensor] = None,
+                          w1_zp: Optional[torch.Tensor] = None,
+                          w2_zp: Optional[torch.Tensor] = None,
+                          a1_scale: Optional[torch.Tensor] = None,
+                          a2_scale: Optional[torch.Tensor] = None,
+                          block_shape: Optional[List[int]] = None) -> None:
+    fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
+                       activation, apply_router_weight_on_input, use_fp8_w8a8,
+                       use_int8_w8a8, use_int8_w8a16, use_int4_w4a16,
+                       per_channel_quant, global_num_experts, expert_map,
+                       w1_scale, w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
+                       block_shape)
+
+
+def outplace_fused_experts(
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        activation: str = "silu",
+        apply_router_weight_on_input: bool = False,
+        use_fp8_w8a8: bool = False,
+        use_int8_w8a8: bool = False,
+        use_int8_w8a16: bool = False,
+        use_int4_w4a16: bool = False,
+        per_channel_quant: bool = False,
+        global_num_experts: int = -1,
+        expert_map: Optional[torch.Tensor] = None,
+        w1_scale: Optional[torch.Tensor] = None,
+        w2_scale: Optional[torch.Tensor] = None,
+        w1_zp: Optional[torch.Tensor] = None,
+        w2_zp: Optional[torch.Tensor] = None,
+        a1_scale: Optional[torch.Tensor] = None,
+        a2_scale: Optional[torch.Tensor] = None,
+        block_shape: Optional[List[int]] = None) -> torch.Tensor:
+    return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,
+                              False, activation, apply_router_weight_on_input,
+                              use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16,
+                              use_int4_w4a16, per_channel_quant,
+                              global_num_experts, expert_map, w1_scale,
+                              w2_scale, w1_zp, w2_zp, a1_scale, a2_scale,
+                              block_shape)
+
+def torch_vllm_inplace_fused_experts(**kwargs) -> torch.Tensor:
+    inplace_fused_experts(**kwargs)
+    hidden_states = kwargs['hidden_states']
+    return hidden_states
+
+
+def torch_vllm_outplace_fused_experts(**kwargs) -> torch.Tensor:
+    return outplace_fused_experts(**kwargs)
+
+torch.ops.vllm.inplace_fused_experts = inplace_fused_experts
+torch.ops.vllm.outplace_fused_experts = outplace_fused_experts
+
+
+import vllm.model_executor.layers.fused_moe.fused_moe
+vllm.model_executor.layers.fused_moe.fused_moe.should_moe_wna16_use_cuda = should_moe_wna16_use_cuda
+vllm.model_executor.layers.fused_moe.fused_moe.fused_experts_impl = fused_experts_impl
+vllm.model_executor.layers.fused_moe.fused_moe.torch_vllm_inplace_fused_experts = torch_vllm_inplace_fused_experts
+vllm.model_executor.layers.fused_moe.fused_moe.torch_vllm_outplace_fused_experts = torch_vllm_outplace_fused_experts
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/moe/patch_fused_moe_layer.py b/vllm_musa/vllm_musa/patch/moe/patch_fused_moe_layer.py
new file mode 100644
index 000000000..75b0b645d
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/moe/patch_fused_moe_layer.py
@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Callable, Optional
+import torch
+
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+from vllm.model_executor.layers.fused_moe.layer import FusedMoE
+
+def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+    pass
+
+def forward_oot(
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    use_grouped_topk: bool,
+    top_k: int,
+    router_logits: torch.Tensor,
+    renormalize: bool,
+    topk_group: Optional[int] = None,
+    num_expert_group: Optional[int] = None,
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
+    custom_routing_function: Optional[Callable] = None,
+    scoring_func: str = "softmax",
+    e_score_correction_bias: Optional[torch.Tensor] = None,
+    apply_router_weight_on_input: bool = False,
+    activation: str = "silu",
+    enable_eplb: bool = False,
+    expert_load_view: Optional[torch.Tensor] = None,
+    logical_to_physical_map: Optional[torch.Tensor] = None,
+    logical_replica_count: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        if enable_eplb:
+            assert expert_load_view is not None
+            assert logical_to_physical_map is not None
+            assert logical_replica_count is not None
+            assert isinstance(layer, FusedMoE)
+
+        topk_weights, topk_ids = FusedMoE.select_experts(
+            hidden_states=x,
+            router_logits=router_logits,
+            top_k=top_k,
+            use_grouped_topk=use_grouped_topk,
+            renormalize=renormalize,
+            topk_group=topk_group,
+            num_expert_group=num_expert_group,
+            custom_routing_function=custom_routing_function,
+            scoring_func=scoring_func,
+            e_score_correction_bias=e_score_correction_bias,
+            # indices_type=self.topk_indices_dtype # TODO: For DeepEP or pplx
+            enable_eplb=enable_eplb,
+            expert_map=expert_map,
+            expert_load_view=expert_load_view,
+            logical_to_physical_map=logical_to_physical_map,
+            logical_replica_count=logical_replica_count,
+        )
+
+        return fused_experts(
+            hidden_states=x,
+            w1=layer.w13_weight,
+            w2=layer.w2_weight,
+            topk_weights=topk_weights,
+            topk_ids=topk_ids,
+            inplace=True,
+            activation=activation,
+            apply_router_weight_on_input=apply_router_weight_on_input,
+            global_num_experts=global_num_experts,
+            expert_map=expert_map,
+        )
+
+def apply(
+    self,
+    layer: torch.nn.Module,
+    x: torch.Tensor,
+    router_logits: torch.Tensor,
+    top_k: int,
+    renormalize: bool,
+    use_grouped_topk: bool = False,
+    topk_group: Optional[int] = None,
+    num_expert_group: Optional[int] = None,
+    global_num_experts: int = -1,
+    expert_map: Optional[torch.Tensor] = None,
+    custom_routing_function: Optional[Callable] = None,
+    scoring_func: str = "softmax",
+    e_score_correction_bias: Optional[torch.Tensor] = None,
+    apply_router_weight_on_input: bool = False,
+    activation: str = "silu",
+    enable_eplb: bool = False,
+    expert_load_view: Optional[torch.Tensor] = None,
+    logical_to_physical_map: Optional[torch.Tensor] = None,
+    logical_replica_count: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+
+    return forward_oot(
+        x=x,
+        layer=layer,
+        router_logits=router_logits,
+        top_k=top_k,
+        renormalize=renormalize,
+        use_grouped_topk=use_grouped_topk,
+        topk_group=topk_group,
+        num_expert_group=num_expert_group,
+        global_num_experts=global_num_experts,
+        expert_map=expert_map,
+        custom_routing_function=custom_routing_function,
+        scoring_func=scoring_func,
+        e_score_correction_bias=e_score_correction_bias,
+        activation=activation,
+        apply_router_weight_on_input=apply_router_weight_on_input,
+        enable_eplb=enable_eplb,
+        expert_load_view=expert_load_view,
+        logical_to_physical_map=logical_to_physical_map,
+        logical_replica_count=logical_replica_count)
+
+import vllm.model_executor.layers.fused_moe.layer
+vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod.process_weights_after_loading = process_weights_after_loading
+vllm.model_executor.layers.fused_moe.layer.UnquantizedFusedMoEMethod.apply = apply
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/moe/patch_grouped_topk.py b/vllm_musa/vllm_musa/patch/moe/patch_grouped_topk.py
new file mode 100644
index 000000000..db05b4440
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/moe/patch_grouped_topk.py
@@ -0,0 +1,140 @@
+import torch
+from typing import Callable, Optional
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+            fused_topk, grouped_topk)
+
+def biased_grouped_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    correction_bias: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: int = 0,
+    topk_group: int = 0,
+):
+    assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
+
+    scores = gating_output.sigmoid()
+    num_token = scores.shape[0]
+    scores_for_choice = scores.view(num_token, -1) + correction_bias.unsqueeze(0)
+    group_scores = (
+        scores_for_choice.view(num_token, num_expert_group, -1)
+        .topk(2, dim=-1)[0]
+        .sum(dim=-1)
+    )  # [n, n_group]
+    group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
+        1
+    ]  # [n, top_k_group]
+    group_mask = torch.zeros_like(group_scores)  # [n, n_group]
+    group_mask.scatter_(1, group_idx, 1)  # [n, n_group]
+    score_mask = (
+        group_mask.unsqueeze(-1)
+        .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group)
+        .reshape(num_token, -1)
+    )  # [n, e]
+    tmp_scores = scores_for_choice.masked_fill(~score_mask.bool(), 0.0)  # [n, e]
+    _, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)
+    topk_weights = scores.gather(1, topk_ids)
+
+    if renormalize:
+        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
+
+    return topk_weights.to(torch.float32), topk_ids.to(torch.int32)
+
+def _fused_grouped_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: int = 0,
+    topk_group: int = 0,
+    e_score_correction_bias: Optional[torch.Tensor] = None):
+
+    from vllm_musa import _musa_custom_ops as ops
+
+    topk_weights, topk_ids = ops.fused_grouped_topk(gating_output, topk, renormalize, num_expert_group, topk_group, e_score_correction_bias)
+
+    return topk_weights, topk_ids.to(torch.int32)
+
+
+def fused_grouped_topk(
+    hidden_states: torch.Tensor,
+    gating_output: torch.Tensor,
+    topk: int,
+    renormalize: bool,
+    num_expert_group: int = 0,
+    topk_group: int = 0,
+    e_score_correction_bias: Optional[torch.Tensor] = None
+):
+    if topk_group % 2 == 0:
+        return _fused_grouped_topk(hidden_states, gating_output, topk, renormalize, num_expert_group, topk_group, e_score_correction_bias)
+
+    return biased_grouped_topk(hidden_states, gating_output, e_score_correction_bias, topk, renormalize, num_expert_group, topk_group)
+
+def select_experts(
+        hidden_states: torch.Tensor,
+        router_logits: torch.Tensor,
+        top_k: int,
+        use_grouped_topk: bool,
+        renormalize: bool,
+        topk_group: Optional[int] = None,
+        num_expert_group: Optional[int] = None,
+        custom_routing_function: Optional[Callable] = None,
+        scoring_func: str = "softmax",
+        e_score_correction_bias: Optional[torch.Tensor] = None,
+        indices_type: Optional[torch.dtype] = None,
+        enable_eplb: bool = False, # need more effort to work
+        expert_map: Optional[torch.Tensor] = None,
+        expert_load_view: Optional[torch.Tensor] = None,
+        logical_to_physical_map: Optional[torch.Tensor] = None,
+        logical_replica_count: Optional[torch.Tensor] = None):
+    # DeekSeekv2 uses grouped_top_k
+    if use_grouped_topk:
+        assert topk_group is not None
+        assert num_expert_group is not None
+        if e_score_correction_bias is None:
+            topk_weights, topk_ids = grouped_topk(
+                hidden_states=hidden_states,
+                gating_output=router_logits,
+                topk=top_k,
+                renormalize=renormalize,
+                num_expert_group=num_expert_group,
+                topk_group=topk_group,
+                scoring_func=scoring_func,
+                e_score_correction_bias=e_score_correction_bias)
+            if indices_type is not None:
+                topk_ids = topk_ids.to(dtype=indices_type)
+        else:
+            topk_weights, topk_ids = fused_grouped_topk(
+                hidden_states=hidden_states,
+                gating_output=router_logits,
+                e_score_correction_bias=e_score_correction_bias,
+                topk=top_k,
+                renormalize=renormalize,
+                num_expert_group=num_expert_group,
+                topk_group=topk_group,
+            )
+            if indices_type is not None:
+                topk_ids = topk_ids.to(dtype=indices_type)
+    elif custom_routing_function is None:
+        topk_weights, topk_ids = fused_topk(
+            hidden_states=hidden_states,
+            gating_output=router_logits,
+            topk=top_k,
+            renormalize=renormalize,
+            indices_type=indices_type,
+        )
+    else:
+        topk_weights, topk_ids = custom_routing_function(
+            hidden_states=hidden_states,
+            gating_output=router_logits,
+            topk=top_k,
+            renormalize=renormalize,
+        )
+        if indices_type is not None:
+            topk_ids = topk_ids.to(dtype=indices_type)
+
+    return topk_weights, topk_ids
+
+import vllm.model_executor.layers.fused_moe.layer
+vllm.model_executor.layers.fused_moe.layer.FusedMoE.select_experts = select_experts
diff --git a/vllm_musa/vllm_musa/patch/moe/patch_utils.py b/vllm_musa/vllm_musa/patch/moe/patch_utils.py
new file mode 100644
index 000000000..22c232c07
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/moe/patch_utils.py
@@ -0,0 +1,32 @@
+from typing import Optional
+
+import torch
+from vllm_musa import _musa_custom_ops as ops
+from vllm.model_executor.layers.fused_moe.utils import  cdiv
+from vllm_musa.patch.quantization.patch_fp8_utils import sglang_per_token_group_quant_fp8
+
+
+def _fp8_quantize(
+    A: torch.Tensor,
+    A_scale: Optional[torch.Tensor],
+    per_act_token: bool,
+    block_shape: Optional[list[int]] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Perform fp8 quantization on the inputs.  If a block_shape
+    is provided, the output will be blocked.
+    """
+    if block_shape is None:
+        A, A_scale = ops.scaled_fp8_quant(
+            A, A_scale, use_per_token_if_dynamic=per_act_token)
+    else:
+        assert len(block_shape) == 2
+        _, block_k = block_shape[0], block_shape[1]
+        A, A_scale = sglang_per_token_group_quant_fp8(A, block_k)
+        assert cdiv(A.size(-1), block_k) == A_scale.size(-1)
+
+    return A, A_scale
+
+
+import vllm.model_executor.layers.fused_moe.utils
+vllm.model_executor.layers.fused_moe.utils._fp8_quantize = _fp8_quantize
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_attention_layer.py b/vllm_musa/vllm_musa/patch/patch_attention_layer.py
new file mode 100644
index 000000000..ae577701e
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_attention_layer.py
@@ -0,0 +1,47 @@
+import torch
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.platforms import current_platform
+from vllm_musa.patch.patch_vllm_utils import direct_register_custom_op
+
+def unified_attention_with_output(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    forward_context: ForwardContext = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
+    self = forward_context.attn_layers[layer_name]
+    kv_cache = self.kv_cache[forward_context.virtual_engine]
+    self.impl.forward(self,
+                      query,
+                      key,
+                      value,
+                      kv_cache,
+                      attn_metadata,
+                      output=output)
+
+
+def unified_attention_with_output_fake(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    return
+
+
+direct_register_custom_op(
+    op_name="unified_attention_with_output",
+    op_func=unified_attention_with_output,
+    mutates_args=["output"],
+    fake_impl=unified_attention_with_output_fake,
+    dispatch_key=current_platform.dispatch_key,
+)
+
+
+import vllm.attention.layer
+vllm.attention.layer.unified_attention_with_output = unified_attention_with_output
+vllm.attention.layer.unified_attention_with_output_fake = unified_attention_with_output_fake
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_linear.py b/vllm_musa/vllm_musa/patch/patch_linear.py
new file mode 100644
index 000000000..d721dee45
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_linear.py
@@ -0,0 +1,328 @@
+import itertools
+from typing import Optional, Callable
+
+import torch
+from torch.nn.parameter import Parameter
+
+from vllm.logger import init_logger
+from vllm.distributed import (get_tensor_model_parallel_rank,
+                              get_tensor_model_parallel_world_size)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig)
+
+from vllm.model_executor.parameter import (BasevLLMParameter,
+                                           BlockQuantScaleParameter,
+                                           PackedColumnParameter,
+                                           PackedvLLMParameter,
+                                           PerTensorScaleParameter,
+                                           RowvLLMParameter)
+
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
+                                                adjust_scalar_to_fused_array,
+                                                adjust_marlin_shard,
+                                                adjust_bitsandbytes_4bit_shard)
+
+from vllm_musa.platforms.musa import on_ph1
+from vllm_musa import _musa_custom_ops as ops
+
+
+logger = init_logger(__name__)
+
+class MergedColumnParallelLinear(ColumnParallelLinear):
+    """Packed linear layers with column parallelism.
+
+    Similar to ColumnParallelLinear, but the weight matrix is concatenated
+    along the output dimension. When the weight matrix is loaded, the
+    different partitions are sharded separately.
+
+    Args:
+        input_size: input dimension of the linear layer.
+        output_sizes: list of output dimensions of the linear layer.
+        bias: If true, add bias.
+        gather_output: If true, call all-gather on output and make the output
+                       available to all GPUs, otherwise, every GPU will have
+                       its own output.
+        skip_bias_add: This was added to enable performance optimizations where
+                       bias can be fused with other element-wise operations. we
+                       skip adding bias but instead return it.
+        params_dtype: Data type for the parameters.
+        quant_config: Quantization configure.
+        prefix: The name of the layer in the state dict, including all parents
+                        (e.g. model.layers.0.qkv_proj)
+    """
+
+    def __init__(self,
+                 input_size: int,
+                 output_sizes: list[int],
+                 bias: bool = True,
+                 gather_output: bool = False,
+                 skip_bias_add: bool = False,
+                 params_dtype: Optional[torch.dtype] = None,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = "",
+                 save_transpose_weight: bool = False):
+        self.output_sizes = output_sizes
+        self.save_transpose_weight = save_transpose_weight
+        tp_size = get_tensor_model_parallel_world_size()
+        assert all(output_size % tp_size == 0 for output_size in output_sizes)
+        super().__init__(input_size=input_size,
+                         output_size=sum(output_sizes),
+                         bias=bias,
+                         gather_output=gather_output,
+                         skip_bias_add=skip_bias_add,
+                         params_dtype=params_dtype,
+                         quant_config=quant_config,
+                         prefix=prefix)
+
+    def weight_loader(self,
+                      param: Parameter,
+                      loaded_weight: torch.Tensor,
+                      loaded_shard_id: Optional[int] = None):
+
+        # Special case for GGUF
+        # initialize GGUF param after we know the quantize type
+        is_gguf_weight = getattr(param, "is_gguf_weight", False)
+        is_gguf_weight_type = getattr(param, "is_gguf_weight_type", False)
+        if is_gguf_weight_type:
+            if loaded_shard_id is not None:
+                param.data[loaded_shard_id].copy_(loaded_weight)
+                param.shard_weight_type[loaded_shard_id] = loaded_weight.item()
+            else:
+                param.shard_weight_type = {
+                    i: loaded_weight.item()
+                    for i, _ in enumerate(self.output_sizes)
+                }
+            return
+
+        if is_gguf_weight:
+            tp_size = get_tensor_model_parallel_world_size()
+            tp_rank = get_tensor_model_parallel_rank()
+
+            output_dim = getattr(param, "output_dim", None)
+            shard_size = loaded_weight.size(output_dim) // tp_size
+            start_idx = tp_rank * shard_size
+
+            if loaded_shard_id is not None:
+                loaded_weight = loaded_weight.narrow(output_dim, start_idx,
+                                                     shard_size)
+                param.shard_id.append(loaded_shard_id)
+                param.shard_id_map[loaded_shard_id] = len(param.data_container)
+                param.data_container.append(loaded_weight)
+                if len(param.data_container) == 2:
+                    self.qweight = param.materialize_nested()
+                return
+
+        param_data = param.data
+        output_dim = getattr(param, "output_dim", None)
+        # Special case for AQLM codebooks.
+        is_metadata = getattr(param, "is_metadata", False)
+        # Special case for per-tensor scale to load scalar into fused array.
+        needs_scalar_to_array = getattr(param, "needs_scalar_to_array", False)
+
+        if loaded_shard_id is None:
+            # Loaded weight is already fused on disk (mlp).
+            # (e.g., Phi-3's gate_up_proj).
+            if output_dim is None:
+                if needs_scalar_to_array:
+                    param_data, loaded_weight = adjust_scalar_to_fused_array(
+                        param_data, loaded_weight, 0)
+
+                assert param_data.shape == loaded_weight.shape
+                param_data.copy_(loaded_weight)
+                return
+            current_shard_offset = 0
+            use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit",
+                                            False)
+            shard_offsets: list[tuple[int, int, int]] = []
+            for i, output_size in enumerate(self.output_sizes):
+                shard_offsets.append((i, current_shard_offset, output_size))
+                current_shard_offset += output_size
+            packed_dim = getattr(param, "packed_dim", None)
+            for shard_id, shard_offset, shard_size in shard_offsets:
+                # Special case for Quantization.
+                # If quantized, we need to adjust the offset and size to account
+                # for the packing.
+                if packed_dim == output_dim:
+                    shard_size = shard_size // param.pack_factor
+                    shard_offset = shard_offset // param.pack_factor
+                    # Special case for Marlin.
+                    shard_size, shard_offset = adjust_marlin_shard(
+                        param, shard_size, shard_offset)
+
+                if use_bitsandbytes_4bit:
+                    index = list(itertools.accumulate([0] + self.output_sizes))
+                    orig_offsets = {
+                        str(i): (index[i], size)
+                        for i, size in enumerate(self.output_sizes)
+                    }
+                    orig_offsets["total"] = (self.output_size, 0)
+                    shard_size, shard_offset = adjust_bitsandbytes_4bit_shard(
+                        param, orig_offsets, str(shard_id))
+
+                loaded_weight_shard = loaded_weight.narrow(
+                    output_dim, shard_offset, shard_size)
+                self.weight_loader(param, loaded_weight_shard, shard_id)
+            return
+
+        assert loaded_shard_id < len(self.output_sizes)
+        tp_rank = get_tensor_model_parallel_rank()
+        tp_size = get_tensor_model_parallel_world_size()
+        if output_dim is not None:
+            shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
+            shard_size = self.output_sizes[loaded_shard_id] // tp_size
+            # Special case for quantization.
+            # If quantized, we need to adjust the offset and size to account
+            # for the packing.
+            packed_dim = getattr(param, "packed_dim", None)
+            if packed_dim == output_dim:
+                shard_size = shard_size // param.pack_factor
+                shard_offset = shard_offset // param.pack_factor
+                # Special case for Marlin.
+                shard_size, shard_offset = adjust_marlin_shard(
+                    param, shard_size, shard_offset)
+
+            use_bitsandbytes_4bit = getattr(param, "use_bitsandbytes_4bit",
+                                            False)
+            is_sharded_weight = getattr(param, "is_sharded_weight", False)
+            # bitsandbytes loads the weights of the specific portion
+            # no need to narrow
+            is_sharded_weight = is_sharded_weight or use_bitsandbytes_4bit
+
+            if use_bitsandbytes_4bit:
+                shard_size = loaded_weight.shape[output_dim]
+                shard_offset = loaded_weight.shape[output_dim] * \
+                    loaded_shard_id
+
+            param_data = param_data.narrow(output_dim, shard_offset,
+                                           shard_size)
+            start_idx = tp_rank * shard_size
+            if not is_sharded_weight:
+                loaded_weight = loaded_weight.narrow(output_dim, start_idx,
+                                                     shard_size)
+        # Special case for AQLM codebooks.
+        elif is_metadata:
+            # metadata indicates fixed size concatenated along dim 0
+            shard_size = loaded_weight.shape[0]
+            shard_offset = loaded_shard_id * shard_size
+            param_data = param_data.narrow(0, shard_offset, shard_size)
+
+        # Special case for per-tensor scales in fused case.
+        elif needs_scalar_to_array:
+            param_data, loaded_weight = adjust_scalar_to_fused_array(
+                param_data, loaded_weight, loaded_shard_id)
+
+        else:
+            ignore_warning = getattr(param, "ignore_warning", False)
+            if not ignore_warning:
+                logger.warning(
+                    "Loading a weight without `output_dim` attribute in "
+                    "MergedColumnParallelLinear, assume the weight is "
+                    "the same for all partitions.")
+
+        assert param_data.shape == loaded_weight.shape
+        param_data.copy_(loaded_weight)
+        # add fused gemm swiglu weight
+        if self.save_transpose_weight:
+            self.fused_swiglu_weight = self.weight.T.contiguous()
+
+    def _load_fused_module_from_checkpoint(self, param: BasevLLMParameter,
+                                           loaded_weight: torch.Tensor):
+        """
+        Handle special case for models where MLP layers are already
+        fused on disk. In this case, we have no shard id. This function
+        determmines the shard id by splitting these layers and then calls
+        the weight loader using the shard id.
+
+        An example of a model with these fused layers:
+        https://huggingface.co/microsoft/Phi-3-mini-4k-instruct
+        """
+
+        current_shard_offset = 0
+        shard_offsets: list[tuple[int, int, int]] = []
+        for i, output_size in enumerate(self.output_sizes):
+            shard_offsets.append((i, current_shard_offset, output_size))
+            current_shard_offset += output_size
+
+        for shard_id, shard_offset, shard_size in shard_offsets:
+            # Special case for Quantization.
+            # If quantized, we need to adjust the offset and size to account
+            # for the packing.
+            if isinstance(param, (PackedColumnParameter, PackedvLLMParameter
+                                  )) and param.packed_dim == param.output_dim:
+                shard_size, shard_offset = \
+                    param.adjust_shard_indexes_for_packing(
+                    shard_size=shard_size, shard_offset=shard_offset)
+
+            loaded_weight_shard = loaded_weight.narrow(param.output_dim,
+                                                       shard_offset,
+                                                       shard_size)
+            self.weight_loader_v2(param, loaded_weight_shard, shard_id)
+
+        # add fused gemm swiglu weight, TODO: 此处后续可更换为self.weight,进一步节省显存
+        if self.save_transpose_weight:
+            self.fused_swiglu_weight = self.weight.T.contiguous()
+
+    def weight_loader_v2(self,
+                         param: BasevLLMParameter,
+                         loaded_weight: torch.Tensor,
+                         loaded_shard_id: Optional[int] = None):
+        if loaded_shard_id is None:
+            if isinstance(param, PerTensorScaleParameter):
+                param.load_merged_column_weight(loaded_weight=loaded_weight,
+                                                shard_id=0)
+                return
+            elif type(param) in (RowvLLMParameter, BasevLLMParameter):
+                param.load_merged_column_weight(loaded_weight=loaded_weight)
+                return
+            # TODO: @dsikka - move to parameter.py
+            self._load_fused_module_from_checkpoint(param, loaded_weight)
+            return
+
+        assert loaded_shard_id < len(self.output_sizes)
+
+        tp_size = get_tensor_model_parallel_world_size()
+
+        if isinstance(param, BlockQuantScaleParameter):
+            from vllm.model_executor.layers.quantization.fp8 import (
+                Fp8LinearMethod, Fp8MoEMethod)
+            assert self.quant_method is not None
+            assert isinstance(self.quant_method,
+                              (Fp8LinearMethod, Fp8MoEMethod))
+            weight_block_size = self.quant_method.quant_config.weight_block_size
+            assert weight_block_size is not None
+            block_n, _ = weight_block_size[0], weight_block_size[1]
+            shard_offset = (
+                (sum(self.output_sizes[:loaded_shard_id]) + block_n - 1) //
+                block_n) // tp_size
+            shard_size = ((self.output_sizes[loaded_shard_id] + block_n - 1) //
+                          block_n // tp_size)
+        else:
+            shard_offset = sum(self.output_sizes[:loaded_shard_id]) // tp_size
+            shard_size = self.output_sizes[loaded_shard_id] // tp_size
+
+        param.load_merged_column_weight(loaded_weight=loaded_weight,
+                                        shard_id=loaded_shard_id,
+                                        shard_offset=shard_offset,
+                                        shard_size=shard_size)
+        if self.save_transpose_weight:
+            if hasattr(self, 'qweight'):
+                self.qweight_t = self.qweight.T.contiguous()
+
+def musa_unquantized_gemm(layer: torch.nn.Module,
+                          x: torch.Tensor,
+                          weight: torch.Tensor,
+                          bias: Optional[torch.Tensor] = None):
+    b = x.shape[0]
+    x_dim, weight_dim = x.dim(), weight.dim()
+    # gemv is faster than mudnn on PH1.
+    if b < 3 and x_dim == weight_dim == 2 and bias is None and on_ph1():
+        out = ops.fused_gemv(x, weight)
+    else:
+        out = torch.nn.functional.linear(x, weight, bias)
+    return out
+
+import vllm.model_executor.layers.utils
+vllm.model_executor.layers.utils.default_unquantized_gemm = musa_unquantized_gemm
+
+import vllm.model_executor.layers.linear
+vllm.model_executor.layers.linear.MergedColumnParallelLinear = MergedColumnParallelLinear
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_metrics.py b/vllm_musa/vllm_musa/patch/patch_metrics.py
new file mode 100644
index 000000000..1d45834ba
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_metrics.py
@@ -0,0 +1,72 @@
+from typing import Callable, Optional, Union
+
+import torch
+import torch_musa
+from vllm.spec_decode.metrics import (AsyncMetricsCollector,
+                                      SpecDecodeWorkerMetrics)
+
+Timer = Callable[[], float]
+
+# TODO: revert this patch when the cuda hard code is removed in vllm
+# init_tensors: Modified the hard-coded cuda judgment logic to npu;
+# maybe_collect_rejsample_metrics: Removed the check for current_platform.is_cuda_alike()
+
+
+def init_tensors(self,
+                 rank: int,
+                 device_type: Union[torch.device, str] = 'musa') -> None:
+    self._rank = rank
+    if isinstance(device_type, torch.device):
+        device_type = device_type.type
+    if device_type == 'musa':
+        self._copy_stream = torch.musa.Stream()
+
+
+def maybe_collect_rejsample_metrics(
+        self, k: int) -> Optional[SpecDecodeWorkerMetrics]:
+
+    # If a copy was initiated in the previous call, collect and return.
+    if self._in_flight_copy is not None:
+        ready_event = self._in_flight_copy
+        self._in_flight_copy = None
+        return self._collect_rejsample_metrics(k, ready_event)
+
+    # Otherwise, check if we should start a new copy.
+    if self._should_collect_rejsample_metrics(self._timer()):
+        assert self._in_flight_copy is None
+        self._in_flight_copy = self._copy_rejsample_metrics_async()
+
+    return None
+
+
+
+def _copy_rejsample_metrics_async(self) -> torch.musa.Event:
+    """
+    TODO: torch.cuda.xxx --> torch.npu.xxx
+    Copy rejection/typical-acceptance sampling metrics
+    (number of accepted tokens, etc) to CPU asynchronously.
+    Returns a NPU event recording when the copy is complete.
+    """
+    assert self._copy_stream is not None
+    self._copy_stream.wait_stream(torch.musa.current_stream())
+
+    with torch.musa.stream(self._copy_stream):
+        self._aggregate_num_accepted_tokens.copy_(
+            self.spec_decode_sampler.num_accepted_tokens, non_blocking=True)
+        self._aggregate_num_emitted_tokens.copy_(
+            self.spec_decode_sampler.num_emitted_tokens, non_blocking=True)
+        # Number of draft tokens is calculated on CPU, so no copy is
+        # required.
+        self._aggregate_num_draft_tokens = (
+            self.spec_decode_sampler.num_draft_tokens)
+
+    aggregate_metrics_ready = torch.musa.Event()
+    aggregate_metrics_ready.record(self._copy_stream)
+
+    return aggregate_metrics_ready
+
+
+import vllm.spec_decode.metrics
+vllm.spec_decode.metrics.AsyncMetricsCollector.init_tensors = init_tensors
+vllm.spec_decode.metrics.AsyncMetricsCollector.maybe_collect_rejsample_metrics = maybe_collect_rejsample_metrics
+vllm.spec_decode.metrics.AsyncMetricsCollector._copy_rejsample_metrics_async = _copy_rejsample_metrics_async
diff --git a/vllm_musa/vllm_musa/patch/patch_mla_common.py b/vllm_musa/vllm_musa/patch/patch_mla_common.py
new file mode 100644
index 000000000..4ef8cf1be
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_mla_common.py
@@ -0,0 +1,566 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from abc import abstractmethod
+from typing import Any, Dict, Generic, List, Optional, Tuple
+import math
+
+import torch
+from vllm_musa import _musa_custom_ops as ops
+from vllm import envs
+from vllm.attention.backends.abstract import (AttentionLayer,
+                                              MLAAttentionImpl, T)
+from vllm.distributed import (get_tensor_model_parallel_world_size,
+                              tensor_model_parallel_all_reduce)
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
+                                               LinearBase, RowParallelLinear,
+                                               UnquantizedLinearMethod)
+from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
+from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8
+from vllm.model_executor.layers.quantization.utils.quant_utils import scaled_quantize
+from vllm.model_executor.layers.rotary_embedding import (
+    DeepseekScalingRotaryEmbedding, RotaryEmbedding)
+
+from vllm_musa._musa_custom_ops import flash_attn_fwd
+
+from vllm_musa.platforms.musa import on_qy2
+
+class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
+    """
+    Common class for implementing repeated parts
+
+    Main reference: DeepseekV2 paper, and FlashInfer Implementation
+    (https://arxiv.org/abs/2405.04434 and https://github.com/flashinfer-ai/flashinfer/pull/551).
+
+    Deepseek's MLA attention works the following way:
+    * Use a single latent vector to represent the entire KV cache.
+    * The attention "simulates" a multi-head attention, while the compute is
+      similar to multi-query attention.
+    * The dataflow is as follows,
+
+        * B: batch/sequence length
+        * H: hidden size
+        * N: number of attention heads
+        * Lq: latent dimension for Q
+        * Lkv: latent dimension for K/V
+        * P: nope dimension, P+R is the actual head_dim in common attention.
+        * R: rope dimension, this slide of the head_dim goes through rope.
+        * V: V head dim.
+        * kv_c: latent/compressed KV
+        * q_c: latent/compressed Q
+
+        #
+        # Outside the MLA attention backend
+        #
+
+        1. The hidden states (B, H) are projected down into cq (B, Lq) and
+           kv_c_k_pe (B, Lkv+R).
+        2. The kv_c_k_pe is split into kv_c (B, Lkv) and k_pe (B, R). cq
+           and kv_c are normalized.
+
+        #
+        # Inside the MLA attention backend
+        #
+
+        * if prefill:
+
+        3. The q_c is then projected up into the multi-head version.
+           * q_c goes from (B, Lq) to (B, N, (P+R)), which is split into q_nope
+             (B, N, P) and q_pe (B, N, R).
+        4. q_pe, k_pe are then passed through rotary embeddings.
+        5. kv_c and k_pe are concatenated and inserted into the cache
+        6. The kv_c is then projected up into the multi-head version.
+           * kv_c goes from (B, Lkv) to (B, N, (P+V)) which has the nope
+             dimensions for K and V, which is split into k_nope (B, N, P)
+             and v (B, N, V).
+        7. q (B, N, (P+R)) and k (B, N, (P+R)) matrices are assembled from
+           q_nope, q_pe, k_nope, k_pe.
+        8. Attention is computued with q, k, v.
+        9. The attention computation returns (B, N, V), which is projected back
+           to (B, H) using out projection.
+
+        * if decode:
+
+        3. Here's the change, we do not perform up the full up projection for
+           q_c, and there is no up projection at all for kv_c. This is
+           achieved by the technique of "weight absorption". The paper says
+           "Fortunately, due to the associative law of matrix multiplication,
+           we can absorb WUK into WUQ, and WUV into WO"
+           * The q up projection turns (B, Lq) into (B, N, (P+R)), we split it
+             into W_UQ (Lq, N, P) and W_QR (Lq, N, R).
+           * The kv_c up projection turns (B, Lkv) into (B, N, (P+V)), we split
+             it into W_UK (Lkv, N, P) and W_UV (Lkv, N, V).
+           * The out projection shape W_O (N*V, H) turns (B, N, V) into (B, H).
+           * We can precompute the product of W_UQ and W_UK into
+             W_UQ_UK (Lq, N, Lkv), which is possible due to QK^T operation in
+             attention.
+           * We can precompute the product of W_UV and W_O into
+             W_UV_O (N, Lkv, H), which is possible due to V@O as the
+             "epilogue" of attention
+        4. We still need to compute q_pe (B, N, R) by applying W_QR to q_latent.
+        5. q_pe, k_pe are then passed through rotary embeddings.
+        6. kv_c and k_pe are concatenated and inserted into the cache
+        7. By applying W_UQ_UK to q_latent, we have the new q_nope of shape
+           (B, N, Lkv).
+        8. q (B, N, (Lkv+R)), k (B, (Lkv+R)) are assembled from q_nope, q_pe,
+           kv_a, k_pe. v (B, Lkv) is exactly the same vector as kv_a.
+        9. The attention is computed with q, k, v. Note that we just performed
+           a MQA attention with (LKv+R) as our head dim.
+        10. The KV cache is updated using the new entries k (B, N, (Lkv+R)),
+           which included the v and rope values.
+        11. The attention computation returns (B, N, Lkv), which is projected
+           back to (B, H) using W_UV_O.
+
+    From @tsu-bin's calculation, we only want to use the absorption technique
+    for decode. The prefill algorithm should still use the up-projected MHA
+    for less flops and memory usage.
+
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: Optional[List[float]],
+        sliding_window: Optional[int],
+        kv_cache_dtype: str,
+        blocksparse_params: Optional[Dict[str, Any]],
+        logits_soft_cap: Optional[float],
+        attn_type: str,
+        # MLA Specific Arguments
+        q_lora_rank: Optional[int],
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        rotary_emb: RotaryEmbedding,
+        # q_proj should be q_b_proj if q_lora_rank is not None, but from an
+        # attention backend perspective we rely on the layer to pass in the
+        # correct matrix
+        q_proj: ColumnParallelLinear,
+        kv_b_proj: ColumnParallelLinear,
+        o_proj: RowParallelLinear,
+    ) -> None:
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        self.kv_cache_dtype = kv_cache_dtype
+
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_head_dim = qk_head_dim
+        self.v_head_dim = v_head_dim
+
+        self.rotary_emb = rotary_emb
+        self.use_yarn_rope = isinstance(rotary_emb,
+                                        DeepseekScalingRotaryEmbedding)
+        self.q_proj = q_proj
+        self.kv_b_proj = kv_b_proj
+        self.o_proj = o_proj
+
+        if on_qy2():
+            self._forward_prefill_flash = self._forward_prefill_flash_mp22
+        else:
+            self._forward_prefill_flash = self._forward_prefill_flash_general
+
+    def _v_up_proj_and_o_proj(self, x):
+
+        output_parallel = torch.matmul(x.flatten(start_dim=1),
+                                        self.W_UV_O)
+        if self.tp_size > 1:
+            output = tensor_model_parallel_all_reduce(output_parallel)
+        else:
+            output = output_parallel
+        return output
+    def _q_proj_and_k_up_proj(self, x):
+            return torch.matmul(x, self.W_Q_UK)\
+                .view(-1, self.num_heads, self.kv_lora_rank)
+    def process_weights_after_loading(self, act_dtype: torch.dtype):
+        def get_scale_group_shapes_for_fp8(layer: LinearBase) -> \
+            Tuple[Tuple[int, int], Tuple[int, int]]:
+            if isinstance(layer.quant_method, Fp8LinearMethod):
+                if layer.quant_method.block_quant is not None:
+                    weight_block_size = \
+                        layer.quant_method.quant_config.weight_block_size
+                    # per-token-group (1, X), block-quantized (X, Y)
+                    return (1, weight_block_size[-1]), weight_block_size
+                else:
+                    return (-1, -1), (-1, -1)  # per-tensor, per-tensor
+            # elif isinstance(layer.quant_method, CompressedTensorsLinearMethod)\
+            #     and isinstance(layer.scheme, CompressedTensorsW8A8Fp8):
+                # this is hacky but we always assume the for
+                # CompressedTensorsW8A8Fp8 the input is dynamic per-token
+                # we ignore if it is static-per-tensor since we are going to
+                # requantize after later anyways
+                strategy = layer.scheme.strategy
+                # if strategy == QuantizationStrategy.TENSOR:
+                #     return (1, -1), (-1, -1)  # per-token, per-tensor
+                # elif strategy == QuantizationStrategy.CHANNEL:
+                #     return (1, -1), (-1, 1)  # per-token, per-channel
+                # else:
+                #     raise NotImplementedError(
+                #         f"QuantizationStrategy.{strategy} is not supported for "
+                #         "fp8 MLA, please run with VLLM_MLA_DISABLE=1")
+                raise NotImplementedError(
+                    f"QuantizationStrategy.{strategy} is not supported for "
+                    "fp8 MLA, please run with VLLM_MLA_DISABLE=1")
+            else:
+                raise NotImplementedError(
+                    "Can't determine scale group shapes for "
+                    f"{layer.quant_method}, please run with VLLM_MLA_DISABLE=1"
+                )
+
+        def get_scales(layer: LinearBase) -> torch.Tensor:
+            if hasattr(layer, "weight_scale_inv"):
+                return layer.weight_scale_inv
+            return layer.weight_scale
+
+        def get_layer_weight(layer):
+            WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
+            for attr in WEIGHT_NAMES:
+                if hasattr(layer, attr):
+                    return getattr(layer, attr)
+            raise AttributeError(
+                f"Layer '{layer}' has no recognized weight attribute:"
+                f" {WEIGHT_NAMES}.")
+
+        def get_and_maybe_dequant_weights(layer: LinearBase):
+            # if is_layer_fp8(layer):
+            #     if isinstance(layer.quant_method, \
+            #         CompressedTensorsLinearMethod) and \
+            #         isinstance(layer.scheme, CompressedTensorsW8A8Fp8):
+            #         # NOTE(lucas): note sure why but `CompressedTensorsW8A8Fp8`
+            #         # seems to store weights as (input, output) instead of
+            #         # (output, input) so we need to transpose
+            #         weight = layer.weight.T  # standardize to (output, input)
+            #     else:
+            #         weight = layer.weight
+            #     _, weight_scale_group_shape = \
+            #         get_scale_group_shapes_for_fp8(layer)
+            #     scales = get_scales(layer)
+
+            #     return scaled_dequantize(weight, scales,
+            #                              weight_scale_group_shape)
+            if not isinstance(layer.quant_method, UnquantizedLinearMethod):
+                # NOTE: This should only be used offline, since it's O(N^3)
+                eye = torch.eye(layer.input_size_per_partition,
+                                dtype=act_dtype,
+                                device=get_layer_weight(layer).device)
+                dequant_weights = layer.quant_method.apply(layer,
+                                                           eye,
+                                                           bias=None)
+                del eye
+                # standardize to (output, input)
+                return dequant_weights.T
+            return layer.weight
+
+        # if not (quantization_scheme_supported(self.kv_b_proj) and\
+        #     quantization_scheme_supported(self.q_proj) and\
+        #         quantization_scheme_supported(self.o_proj)):
+        #     raise NotImplementedError(
+        #         "Only FP8 and UnquantizedLinearMethod are supported for MLA"
+        #         ", please run with VLLM_MLA_DISABLE=1")
+
+        weight_dtype = get_layer_weight(self.kv_b_proj).dtype
+        assert get_layer_weight(self.o_proj).dtype == weight_dtype
+        assert get_layer_weight(self.q_proj).dtype == weight_dtype
+
+        kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
+        assert kv_b_proj_weight.shape == (
+            self.kv_lora_rank,
+            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), (
+                f"{kv_b_proj_weight.shape=}, "
+                f"{self.kv_lora_rank=}, "
+                f"{self.num_heads=}, "
+                f"{self.qk_nope_head_dim=}, "
+                f"{self.v_head_dim=}")
+        kv_b_proj_weight = kv_b_proj_weight.view(
+            self.kv_lora_rank,
+            self.num_heads,
+            self.qk_nope_head_dim + self.v_head_dim,
+        )
+
+        W_UK, W_UV = kv_b_proj_weight.split(
+            [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        q_proj_weight = get_and_maybe_dequant_weights(self.q_proj).T\
+                .view(-1, self.num_heads, self.qk_head_dim)
+
+        # can be W_Q or W_UQ depending q_lora_rank, the former if
+        # q_lora_rank is None, the latter otherwise. From the Attention backend
+        # perspective though we call these both W_Q and rely on the layer
+        # to pass in the correct matrix
+        W_Q = q_proj_weight[..., :self.qk_nope_head_dim]
+        self.W_QR = q_proj_weight[..., self.qk_nope_head_dim:]\
+            .flatten(start_dim=1).contiguous()
+
+        # W_QR is small so for simplicity we dont bother requantizing it
+        self.W_QR = self.W_QR.to(act_dtype)
+
+        if True:
+            requantization_enabled = True
+            if is_fp8(weight_dtype) and requantization_enabled:
+                # This assumes it wise to requantize using the same group shapes
+                # (i.e. strategy, per-tensor, per-channel, block etc.) that the
+                # weights were originally quantized
+                requant_input_group_shape, requant_weight_group_shape = \
+                    get_scale_group_shapes_for_fp8(self.q_proj)
+                assert (requant_input_group_shape, requant_weight_group_shape)\
+                    == get_scale_group_shapes_for_fp8(self.kv_b_proj)
+                assert (requant_input_group_shape, requant_weight_group_shape)\
+                    == get_scale_group_shapes_for_fp8(self.o_proj)
+                self.reqaunt_input_group_shape = requant_input_group_shape
+                self.reqaunt_weight_group_shape = requant_weight_group_shape
+
+            #
+            # Perform matrix-absorption following
+            #     https://github.com/flashinfer-ai/flashinfer/pull/551
+            # for decode, as a result we end up with absorbed weights for decode
+            # and another copy of raw weights for prefill.
+            #
+            self.W_UK, self.W_UV = kv_b_proj_weight.split(
+                [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+            # We absorb `W_UK` into `W_Q` resulting in either W_Q_UK or W_UQ_UK
+            # depending q_lora_rank, the former if q_lora_rank is None, the
+            # latter otherwise
+            # basically if q_lora_rank is none we are absorbing into q_proj
+            # instead of UQ
+            W_Q_UK = torch.einsum("qnd,lnd -> qnl", W_Q, W_UK)\
+                .flatten(start_dim=1).contiguous()
+
+            self.W_Q_UK = W_Q_UK.to(act_dtype)
+
+            W_O = get_and_maybe_dequant_weights(self.o_proj)\
+                .view(-1, self.num_heads, self.v_head_dim)
+            W_UV_O = torch.einsum("lnd,hnd -> nlh", W_UV, W_O)\
+                .flatten(start_dim=0, end_dim=1).contiguous()
+
+            self.W_UV_O = W_UV_O.to(act_dtype)
+
+            self.tp_size = get_tensor_model_parallel_world_size()
+        else:
+            if is_fp8(weight_dtype):
+                raise NotImplementedError(
+                    "Currently fp8 requires matrix absorption")
+
+            self.W_UV = W_UV
+            self.W_UK = W_UK
+            self.W_Q = W_Q.flatten(start_dim=1)
+
+    @abstractmethod
+    def _forward_prefill(
+        self,
+        q: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _forward_decode(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: T,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def apply_pure_rope(
+        self,
+        input_positions: torch.Tensor,
+        q_pe: torch.Tensor,
+        k_pe: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        seq_len = input_positions.size(0)
+        ori_q_pe_shape, ori_k_pe_shape = q_pe.shape, k_pe.shape
+
+        q_pe, k_pe = self.rotary_emb(
+            input_positions,
+            q_pe.reshape(seq_len, -1),
+            k_pe.reshape(seq_len, -1),
+        )
+        q_pe, k_pe = q_pe.view(ori_q_pe_shape), k_pe.view(ori_k_pe_shape)
+
+        return q_pe, k_pe
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        hidden_states_or_q_c: torch.Tensor,  # query in unified attn
+        k_c_normed: torch.Tensor,  # key in unified attn
+        k_pe: torch.Tensor,  # value in unified attn
+        kv_cache: torch.Tensor,
+        attn_metadata: T,
+        output: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if output is not None:
+            raise NotImplementedError(
+                "output is not yet supported for MLAImplBase")
+
+        is_decode = attn_metadata.decode_metadata is not None
+        is_prefill = attn_metadata.prefill_metadata is not None
+
+        if (is_decode and is_prefill):
+            raise NotImplementedError(
+                "chunked prefill is not supported for MLAImplBase")
+
+        # Restore head dim (for rotary embedding)
+        k_pe = k_pe.unsqueeze(1)
+        assert hasattr(attn_metadata, "input_positions")
+
+
+        num_prefill_tokens: int = attn_metadata.num_prefill_tokens
+
+        decode_hs_or_q_c = hidden_states_or_q_c[num_prefill_tokens:]
+        decode_k_pe = k_pe[num_prefill_tokens:]
+        decode_input_positions = \
+            attn_metadata.input_positions[num_prefill_tokens:]
+
+        prefill_hs_or_q_c = hidden_states_or_q_c[:num_prefill_tokens]
+        prefill_k_pe = k_pe[:num_prefill_tokens]
+        prefill_input_positions = \
+            attn_metadata.input_positions[:num_prefill_tokens]
+        prefill_k_c_normed = k_c_normed[:num_prefill_tokens]
+
+
+        rope_fn = (self.rotary_emb
+                   if self.use_yarn_rope else self.apply_pure_rope)
+
+        if is_decode:
+            decode_ql_nope = self._q_proj_and_k_up_proj(decode_hs_or_q_c)
+            decode_q_pe = torch.matmul(decode_hs_or_q_c, self.W_QR)\
+                .view(-1, self.num_heads, self.qk_rope_head_dim)
+            decode_q_pe, decode_k_pe = rope_fn(decode_input_positions, decode_q_pe, decode_k_pe)
+        else:
+            assert is_prefill
+            prefill_q = self.q_proj(prefill_hs_or_q_c)[0]\
+                .view(-1, self.num_heads, self.qk_head_dim)
+            prefill_q_pe = prefill_q[..., self.qk_nope_head_dim:]
+
+            # TODO(lucas): there must be a nicer way to write this line
+            prefill_q_pe[...], prefill_k_pe[...] = \
+                rope_fn(
+                    prefill_input_positions,
+                    prefill_q_pe,
+                    prefill_k_pe)
+
+        # write the latent and rope to kv cache
+        if kv_cache.numel() > 0:
+            ops.concat_and_cache_mla(
+                k_c_normed,
+                k_pe.squeeze(1),
+                kv_cache,
+                attn_metadata.slot_mapping.flatten(),
+                kv_cache_dtype=self.kv_cache_dtype,
+                scale=layer._k_scale,
+            )
+
+        output = torch.empty(attn_metadata.num_prefill_tokens +
+                            attn_metadata.num_decode_tokens,
+                            self.o_proj.output_size,
+                            device=hidden_states_or_q_c.device,
+                            dtype=hidden_states_or_q_c.dtype)
+
+        if attn_metadata.prefill_metadata is not None:
+            output[:num_prefill_tokens] = self._forward_prefill(
+                prefill_q, prefill_k_c_normed, prefill_k_pe,
+                attn_metadata)
+
+        if attn_metadata.decode_metadata is not None:
+            output[num_prefill_tokens:] = self._forward_decode(
+                decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)
+        return output
+
+
+    # Optional common flash-attn based prefill
+    def _forward_prefill_flash_mp22(
+        self,
+        q: torch.Tensor,
+        k_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        seq_start_loc: torch.Tensor,
+        max_prefill_seq_len: int,
+        attn_metadata
+    ) -> torch.Tensor:
+
+        kv_nope = self.kv_b_proj(k_c_normed)[0]\
+            .view(-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = kv_nope\
+            .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
+
+        # For MLA the v head dim is smaller than qk head dim so we pad out
+        # v with 0s to match the qk head dim
+        cu_seqlen =  attn_metadata.prefill_metadata.query_start_loc
+        v = v.contiguous()
+        o = torch.empty((q.shape[0], q.shape[1], v.shape[-1]), device=q.device, dtype=q.dtype)
+        attn_output = flash_attn_fwd(
+            q,
+            k,
+            v,
+            o,
+            cu_seqlen,
+            cu_seqlen,
+            max_prefill_seq_len,
+            max_prefill_seq_len,
+            causal=True,
+            varlen=True,
+            return_softmax_lse=False
+        )
+
+        attn_output = attn_output\
+                .reshape(-1, self.num_heads * v.shape[-1])
+
+        return self.o_proj(attn_output)[0]
+
+    def _forward_prefill_flash_general(
+        self,
+        q: torch.Tensor,
+        k_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        seq_start_loc: torch.Tensor,
+        max_prefill_seq_len: int,
+        attn_metadata
+    ) -> torch.Tensor:
+        kv_nope = self.kv_b_proj(k_c_normed)[0]\
+            .view(-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = kv_nope\
+            .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
+
+        # For MLA the v head dim is smaller than qk head dim so we pad out
+        # v with 0s to match the qk head dim
+        v_padded = torch.nn.functional.pad(v, [0, q.shape[-1] - v.shape[-1]],
+                                           value=0)
+
+        sum_seq, q_h, qk_d = q.shape
+        _, kv_h, v_d = v.shape
+        bs = seq_start_loc.shape[0] - 1
+
+        q_pad = torch.empty(bs, q_h, max_prefill_seq_len, qk_d, device=q.device, dtype=q.dtype)
+        k_pad = torch.empty(bs, kv_h, max_prefill_seq_len, qk_d, device=q.device, dtype=q.dtype)
+        v_pad = torch.empty(bs, kv_h, max_prefill_seq_len, qk_d, device=q.device, dtype=q.dtype)
+        output = torch.empty(sum_seq, q_h, qk_d, device=q.device, dtype=q.dtype)
+
+        ops.varlen_fa_seqlen_pad(q, k, v_padded, q_pad, k_pad, v_pad, seq_start_loc, sum_seq, max_prefill_seq_len, qk_d, q_h, kv_h, bs)
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            q_pad,
+            k_pad,
+            v_pad,
+            is_causal=True,
+        )
+        ops.varlen_fa_seqlen_unpad(attn_output, output, seq_start_loc, sum_seq, max_prefill_seq_len, qk_d, q_h, bs)
+
+        return self.o_proj(output[..., :v_d].reshape(-1, q_h*v_d))[0]
+
+import vllm.attention.backends.mla.common
+vllm.attention.backends.mla.common.MLACommonImpl = MLACommonImpl
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_musa_custom_ops.py b/vllm_musa/vllm_musa/patch/patch_musa_custom_ops.py
new file mode 100644
index 000000000..454c582ba
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_musa_custom_ops.py
@@ -0,0 +1,85 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import torch
+
+from vllm import _custom_ops
+
+from vllm_musa import _musa_custom_ops as musa_ops
+
+_custom_ops.reshape_and_cache = musa_ops.reshape_and_cache
+_custom_ops.copy_blocks = musa_ops.copy_blocks
+_custom_ops.swap_blocks = musa_ops.swap_blocks
+_custom_ops.gather_cache = musa_ops.gather_cache
+_custom_ops.paged_attention_v1 = musa_ops.paged_attention_v1
+_custom_ops.paged_attention_v2 = musa_ops.paged_attention_v2
+_custom_ops.copy_blocks_mla = musa_ops.copy_blocks_mla
+_custom_ops.concat_and_cache_mla = musa_ops.concat_and_cache_mla
+_custom_ops.topk_softmax = musa_ops.topk_softmax
+_custom_ops.moe_sum = musa_ops.moe_sum
+_custom_ops.moe_align_block_size = musa_ops.moe_align_block_size
+_custom_ops.sgl_moe_align_block_size = musa_ops.sgl_moe_align_block_size
+_custom_ops.silu_and_mul = musa_ops.silu_and_mul
+_custom_ops.fused_grouped_topk = musa_ops.fused_grouped_topk
+_custom_ops.weak_ref_tensor = musa_ops.weak_ref_tensor
+_custom_ops.mudnn_w8a8_scaled_mm = musa_ops.mudnn_w8a8_scaled_mm
+_custom_ops.varlen_fa_seqlen_pad = musa_ops.varlen_fa_seqlen_pad
+_custom_ops.varlen_fa_seqlen_unpad = musa_ops.varlen_fa_seqlen_unpad
+_custom_ops.varlen_fa_lse_seqlen_unpad = musa_ops.varlen_fa_lse_seqlen_unpad
+_custom_ops.rms_norm = musa_ops.rms_norm
+_custom_ops.fused_add_rms_norm = musa_ops.fused_add_rms_norm
+_custom_ops.sgl_per_token_group_quant_fp8 = musa_ops.sgl_per_token_group_quant_fp8
+_custom_ops.mla_mudnn_bmm = musa_ops.mla_mudnn_bmm
+
+torch.ops._C.reshape_and_cache = musa_ops.reshape_and_cache
+torch.ops._C.copy_blocks = musa_ops.copy_blocks
+torch.ops._C.swap_blocks = musa_ops.swap_blocks
+torch.ops._C.paged_attention_v1 = musa_ops.paged_attention_v1
+torch.ops._C.paged_attention_v2 = musa_ops.paged_attention_v2
+torch.ops._C.merge_attn_states = musa_ops.merge_attn_states
+torch.ops._C.copy_blocks_mla = musa_ops.copy_blocks_mla
+torch.ops._C.concat_and_cache_mla = musa_ops.concat_and_cache_mla
+torch.ops._C.topk_softmax = musa_ops.topk_softmax
+torch.ops._C.moe_sum = musa_ops.moe_sum
+torch.ops._C.moe_align_block_size = musa_ops.moe_align_block_size
+torch.ops._C.sgl_moe_align_block_size = musa_ops.sgl_moe_align_block_size
+torch.ops._C.silu_and_mul = musa_ops.silu_and_mul
+torch.ops._C.fused_grouped_topk = musa_ops.fused_grouped_topk
+torch.ops._C.weak_ref_tensor = musa_ops.weak_ref_tensor
+torch.ops._C.static_scaled_fp8_quant = musa_ops.static_scaled_fp8_quant
+torch.ops._C.dynamic_scaled_fp8_quant = musa_ops.dynamic_scaled_fp8_quant
+torch.ops._C.dynamic_per_token_scaled_fp8_quant = musa_ops.dynamic_per_token_scaled_fp8_quant
+torch.ops._C_cache_ops.reshape_and_cache_flash = musa_ops.reshape_and_cache_flash
+
+torch.ops._C_cache_ops.reshape_and_cache = musa_ops.reshape_and_cache
+torch.ops._C_cache_ops.copy_blocks = musa_ops.copy_blocks
+torch.ops._C_cache_ops.swap_blocks = musa_ops.swap_blocks
+torch.ops._C_cache_ops.gather_cache = musa_ops.gather_cache
+torch.ops._C_cache_ops.copy_blocks_mla = musa_ops.copy_blocks_mla
+torch.ops._C_cache_ops.concat_and_cache_mla = musa_ops.concat_and_cache_mla
+torch.ops._moe_C.topk_softmax = musa_ops.topk_softmax
+torch.ops._moe_C.moe_sum = musa_ops.moe_sum
+torch.ops._moe_C.moe_align_block_size = musa_ops.moe_align_block_size
+torch.ops._moe_C.sgl_moe_align_block_size = musa_ops.sgl_moe_align_block_size
+
+# custom ar(all-reduce)
+_custom_ops.init_custom_ar = musa_ops.init_custom_ar
+_custom_ops.all_reduce = musa_ops.all_reduce
+_custom_ops.dispose = musa_ops.dispose
+_custom_ops.meta_size = musa_ops.meta_size
+_custom_ops.register_buffer = musa_ops.register_buffer
+_custom_ops.get_graph_buffer_ipc_meta = musa_ops.get_graph_buffer_ipc_meta
diff --git a/vllm_musa/vllm_musa/patch/patch_parallel_state.py b/vllm_musa/vllm_musa/patch/patch_parallel_state.py
new file mode 100644
index 000000000..0e863006f
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_parallel_state.py
@@ -0,0 +1,38 @@
+import torch
+import torch_musa
+
+from contextlib import contextmanager, nullcontext
+from typing import Optional
+
+from vllm.distributed.parallel_state import GraphCaptureContext
+
+@contextmanager
+def graph_capture(
+        self, graph_capture_context: Optional[GraphCaptureContext] = None):
+    if graph_capture_context is None:
+        stream = torch.musa.Stream()
+        graph_capture_context = GraphCaptureContext(stream)
+    else:
+        stream = graph_capture_context.stream
+
+    # only musa uses this function,
+    # so we don't abstract it into the base class
+    maybe_ca_context = nullcontext()
+    from vllm_musa.distributed.device_communicators.musa_communicator import MusaCommunicator
+    if self.device_communicator is not None:
+        assert isinstance(self.device_communicator, MusaCommunicator)
+        ca_comm = self.device_communicator.ca_comm
+        if ca_comm is not None:
+            maybe_ca_context = ca_comm.capture()  # type: ignore
+
+    # ensure all initialization operations complete before attempting to
+    # capture the graph on another stream
+    curr_stream = torch.musa.current_stream()
+    if curr_stream != stream:
+        stream.wait_stream(curr_stream)
+
+    with torch.musa.stream(stream), maybe_ca_context:
+        yield graph_capture_context
+
+import vllm.distributed.parallel_state
+vllm.distributed.parallel_state.GroupCoordinator.graph_capture = graph_capture
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_ray.py b/vllm_musa/vllm_musa/patch/patch_ray.py
new file mode 100644
index 000000000..860b875e6
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_ray.py
@@ -0,0 +1,86 @@
+import torch_musa  # noqa: F401
+import vllm
+import vllm.envs as envs
+from vllm.executor.ray_utils import RayWorkerWrapper
+from vllm.logger import init_logger
+from vllm.executor.ray_distributed_executor import RayDistributedExecutor
+logger = init_logger(__name__)
+
+if RayWorkerWrapper is not None:
+
+    class MTGPURayWorkerWrapper(RayWorkerWrapper):
+        """Importing torch_musa in other Ray processes through an empty class and
+        a monkey patch.
+
+        When Ray performs a remote call, it serializes the Task or Actor and passes
+        it to the Worker process, where it is deserialized and executed.
+
+        If no patch is applied, the default code of the RayWorkerWrapper provided
+        by vLLM is used, which does not import torch_musa, causing an error in the
+        Worker process.
+        See https://github.com/vllm-project/vllm-ascend/pull/92.
+        """
+
+        pass
+
+    # vllm.executor.ray_utils.RayWorkerWrapper = MTGPURayWorkerWrapper
+
+    def _compiled_ray_dag(self, enable_asyncio: bool):
+        assert self.parallel_config.use_ray
+        self._check_ray_adag_installation()
+        from ray.dag import InputNode, MultiOutputNode
+        from ray.experimental.channel.torch_tensor_type import TorchTensorType
+
+        logger.info("VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL = %s",
+                    envs.VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL)
+        logger.info("VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM = %s",
+                    envs.VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM)
+        with InputNode() as input_data:
+            # Example DAG: PP=2, TP=4
+            # (ExecuteModelReq, None) -> 0 -> (ExecuteModelReq, IntermediateOutput) -> 4 -> SamplerOutput   # noqa: E501
+            #                         -> 1 -> (ExecuteModelReq, IntermediateOutput) -> 5 -> SamplerOutput   # noqa: E501
+            #                         -> 2 -> (ExecuteModelReq, IntermediateOutput) -> 6 -> SamplerOutput   # noqa: E501
+            #                         -> 3 -> (ExecuteModelReq, IntermediateOutput) -> 7 -> SamplerOutput   # noqa: E501
+
+            # All workers in the first TP group will take in the
+            # ExecuteModelRequest as input.
+            outputs = [input_data for _ in self.pp_tp_workers[0]]
+            for pp_rank, tp_group in enumerate(self.pp_tp_workers):
+                # Each PP worker takes in the output of the previous PP worker,
+                # and the TP group executes in SPMD fashion.
+                if self.use_v1:
+                    outputs = [
+                        worker.execute_model.
+                        bind(  # type: ignore[attr-defined]
+                            outputs[i]) for i, worker in enumerate(tp_group)
+                    ]
+                else:
+                    outputs = [
+                        worker.execute_model_spmd.
+                        bind(  # type: ignore[attr-defined]
+                            outputs[i]) for i, worker in enumerate(tp_group)
+                    ]
+
+                last_pp_rank = len(self.pp_tp_workers) - 1
+                if pp_rank < last_pp_rank:
+                    # Specify how intermediate tensors should be passed
+                    # between pp stages, no need to specify for the last
+                    # pp stage.
+                    transport = "mccl" \
+                        if envs.VLLM_USE_RAY_COMPILED_DAG_NCCL_CHANNEL \
+                        else "auto"
+                    logger.info(f"Using transport {transport} for pp stage {pp_rank}.")
+                    outputs = [
+                        output.with_type_hint(
+                            TorchTensorType(transport=transport))
+                        for output in outputs
+                    ]
+
+            forward_dag = MultiOutputNode(outputs)
+
+        return forward_dag.experimental_compile(
+            enable_asyncio=enable_asyncio,
+            _overlap_gpu_communication=envs.
+            VLLM_USE_RAY_COMPILED_DAG_OVERLAP_COMM)
+
+    # RayDistributedExecutor._compiled_ray_dag = _compiled_ray_dag
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_rotary_embedding.py b/vllm_musa/vllm_musa/patch/patch_rotary_embedding.py
new file mode 100644
index 000000000..35eca71d2
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_rotary_embedding.py
@@ -0,0 +1,38 @@
+import torch
+from vllm.model_executor.layers.rotary_embedding import (DeepseekScalingRotaryEmbedding,
+                                                         _yarn_linear_ramp_mask,
+                                                         _yarn_find_correction_range)
+
+
+def _compute_inv_freq(self, scaling_factor: float) -> torch.Tensor:
+    pos_freqs = self.base**(torch.arange(
+        0, self.rotary_dim, 2, dtype=torch.float, device="musa") /
+                            self.rotary_dim)
+    inv_freq_extrapolation = 1.0 / pos_freqs
+    inv_freq_interpolation = 1.0 / (scaling_factor * pos_freqs)
+
+    low, high = _yarn_find_correction_range(self.beta_fast, self.beta_slow,
+                                            self.rotary_dim, self.base,
+                                            self.max_position_embeddings)
+    # Get n-d rotational scaling corrected for extrapolation
+    inv_freq_mask = (1 - _yarn_linear_ramp_mask(
+        low, high, self.rotary_dim // 2,
+        dtype=torch.float)) * self.extrapolation_factor
+    inv_freq = inv_freq_interpolation * (
+        1 - inv_freq_mask) + inv_freq_extrapolation * inv_freq_mask
+    return inv_freq
+
+def _compute_cos_sin_cache(self) -> torch.Tensor:
+    inv_freq = self._compute_inv_freq(self.scaling_factor)
+    t = torch.arange(self.max_position_embeddings * self.scaling_factor,
+                        device="musa",
+                        dtype=torch.float32)
+    freqs = torch.einsum("i,j -> ij", t, inv_freq)
+    cos = (freqs.cos() * self.mscale)
+    sin = (freqs.sin() * self.mscale)
+    cache = torch.cat((cos, sin), dim=-1)
+    return cache
+
+
+DeepseekScalingRotaryEmbedding._compute_inv_freq = _compute_inv_freq
+DeepseekScalingRotaryEmbedding._compute_cos_sin_cache = _compute_cos_sin_cache
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_sampler.py b/vllm_musa/vllm_musa/patch/patch_sampler.py
new file mode 100644
index 000000000..6dbb7c522
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_sampler.py
@@ -0,0 +1,32 @@
+import torch
+
+from typing import List, Optional
+from vllm.model_executor.sampling_metadata import SequenceGroupToSample
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+def _multinomial(
+    probs: torch.Tensor,
+    num_samples: int,
+    seq_groups: Optional[list[SequenceGroupToSample]] = None,
+) -> torch.Tensor:
+    if num_samples > 1:
+        probs = probs.repeat_interleave(num_samples, dim=0)
+    q = torch.empty_like(probs)
+    if seq_groups is None:
+        q.exponential_()
+    else:
+        sample_idx = 0
+        for seq_group in seq_groups:
+            seq_ids = seq_group.seq_ids
+            stride = len(seq_ids) * num_samples
+            assert seq_group.generator is not None
+            q[sample_idx:sample_idx +
+              stride].exponential_(generator=seq_group.generator)
+            sample_idx += stride
+    q.clamp_min_(torch.finfo(q.dtype).tiny)
+    return probs.div_(q).argmax(dim=1).view(-1, num_samples)
+
+import vllm.model_executor.layers.sampler
+vllm.model_executor.layers.sampler._multinomial = _multinomial
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_torch_custom_ops.py b/vllm_musa/vllm_musa/patch/patch_torch_custom_ops.py
new file mode 100644
index 000000000..bac979aef
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_torch_custom_ops.py
@@ -0,0 +1,154 @@
+import torch
+
+from typing import Optional, Sequence
+
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm_musa.patch.patch_vllm_utils import direct_register_custom_op
+from vllm.platforms import current_platform
+
+
+def unified_attention_with_output(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    forward_context: ForwardContext = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
+    self = forward_context.attn_layers[layer_name]
+    kv_cache = self.kv_cache[forward_context.virtual_engine]
+    self.impl.forward(self,
+                      query,
+                      key,
+                      value,
+                      kv_cache,
+                      attn_metadata,
+                      output=output)
+
+def unified_attention_with_output_fake(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    return
+
+
+direct_register_custom_op(
+    op_name="unified_attention_with_output",
+    op_func=unified_attention_with_output,
+    mutates_args=["output"],
+    fake_impl=unified_attention_with_output_fake,
+    dispatch_key=current_platform.dispatch_key,
+)
+
+
+import vllm.attention.layer
+vllm.attention.layer.unified_attention_with_output = unified_attention_with_output
+vllm.attention.layer.unified_attention_with_output_fake = unified_attention_with_output_fake
+
+
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts_impl
+
+def inplace_fused_experts(hidden_states: torch.Tensor,
+                          w1: torch.Tensor,
+                          w2: torch.Tensor,
+                          topk_weights: torch.Tensor,
+                          topk_ids: torch.Tensor,
+                          use_fp8_w8a8: bool = False,
+                          use_int8_w8a16: bool = False,
+                          use_int4_w4a16: bool = False,
+                          w1_scale: Optional[torch.Tensor] = None,
+                          w2_scale: Optional[torch.Tensor] = None,
+                          w1_zp: Optional[torch.Tensor] = None,
+                          w2_zp: Optional[torch.Tensor] = None,
+                          a1_scale: Optional[torch.Tensor] = None,
+                          a2_scale: Optional[torch.Tensor] = None,
+                          block_shape: Optional[Sequence[int]] = None) -> torch.Tensor:
+    fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids, True,
+                       use_fp8_w8a8, use_int8_w8a16, use_int4_w4a16, w1_scale,
+                       w2_scale, w1_zp, w2_zp, a1_scale, a2_scale, block_shape)
+
+
+def inplace_fused_experts_fake(
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        use_fp8_w8a8: bool = False,
+        use_int8_w8a16: bool = False,
+        use_int4_w4a16: bool = False,
+        w1_scale: Optional[torch.Tensor] = None,
+        w2_scale: Optional[torch.Tensor] = None,
+        w1_zp: Optional[torch.Tensor] = None,
+        w2_zp: Optional[torch.Tensor] = None,
+        a1_scale: Optional[torch.Tensor] = None,
+        a2_scale: Optional[torch.Tensor] = None,
+        block_shape: Optional[Sequence[int]] = None) -> torch.Tensor:
+    pass
+
+
+direct_register_custom_op(
+    op_name="inplace_fused_experts",
+    op_func=inplace_fused_experts,
+    mutates_args=["hidden_states"],
+    fake_impl=inplace_fused_experts_fake,
+)
+
+def outplace_fused_experts(
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        use_fp8_w8a8: bool = False,
+        use_int8_w8a16: bool = False,
+        use_int4_w4a16: bool = False,
+        w1_scale: Optional[torch.Tensor] = None,
+        w2_scale: Optional[torch.Tensor] = None,
+        w1_zp: Optional[torch.Tensor] = None,
+        w2_zp: Optional[torch.Tensor] = None,
+        a1_scale: Optional[torch.Tensor] = None,
+        a2_scale: Optional[torch.Tensor] = None,
+        block_shape: Optional[Sequence[int]] = None) -> torch.Tensor:
+    return fused_experts_impl(hidden_states, w1, w2, topk_weights, topk_ids,
+                              False, use_fp8_w8a8, use_int8_w8a16,
+                              use_int4_w4a16, w1_scale, w2_scale, w1_zp, w2_zp,
+                              a1_scale, a2_scale, block_shape)
+
+
+def outplace_fused_experts_fake(
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        use_fp8_w8a8: bool = False,
+        use_int8_w8a16: bool = False,
+        use_int4_w4a16: bool = False,
+        w1_scale: Optional[torch.Tensor] = None,
+        w2_scale: Optional[torch.Tensor] = None,
+        w1_zp: Optional[torch.Tensor] = None,
+        w2_zp: Optional[torch.Tensor] = None,
+        a1_scale: Optional[torch.Tensor] = None,
+        a2_scale: Optional[torch.Tensor] = None,
+        block_shape: Optional[Sequence[int]] = None) -> torch.Tensor:
+    return torch.empty_like(hidden_states)
+
+
+direct_register_custom_op(
+    op_name="outplace_fused_experts",
+    op_func=outplace_fused_experts,
+    mutates_args=[],
+    fake_impl=outplace_fused_experts_fake,
+)
+
+
+import vllm.model_executor.layers.fused_moe.fused_moe
+vllm.model_executor.layers.fused_moe.fused_moe.inplace_fused_experts = inplace_fused_experts
+vllm.model_executor.layers.fused_moe.fused_moe.inplace_fused_experts_fake = inplace_fused_experts_fake
+vllm.model_executor.layers.fused_moe.fused_moe.outplace_fused_experts = outplace_fused_experts
+vllm.model_executor.layers.fused_moe.fused_moe.outplace_fused_experts_fake = outplace_fused_experts_fake
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_triton_decode_attention.py b/vllm_musa/vllm_musa/patch/patch_triton_decode_attention.py
new file mode 100644
index 000000000..630b239aa
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_triton_decode_attention.py
@@ -0,0 +1,62 @@
+import torch, torch_musa
+import triton
+from vllm_musa._musa_custom_ops import decode_mla
+from vllm.attention.ops.triton_decode_attention import _fwd_kernel_stage2
+
+
+def _decode_softmax_reducev_fwd(
+    logits,
+    q,
+    o,
+    v_buffer,
+    b_seq_len,
+    num_kv_splits,
+):
+    batch, head_num = q.shape[0], q.shape[1]
+    Lv = v_buffer.shape[-1]
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    NUM_KV_SPLITS = num_kv_splits
+
+    extra_kargs = {}
+
+    grid = (batch, head_num)
+    _fwd_kernel_stage2[grid](
+        logits,
+        o,
+        b_seq_len,
+        logits.stride(0),
+        logits.stride(1),
+        logits.stride(2),
+        o.stride(0),
+        o.stride(1),
+        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        BLOCK_DV=BLOCK_DV,
+        Lv=Lv,
+        num_warps=4,
+        num_stages=2,
+        **extra_kargs,
+    )
+
+
+def decode_attention_fwd_grouped(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    req_to_token,
+    b_seq_len,
+    attn_logits,
+    num_kv_splits,
+    sm_scale,
+    page_size,
+    logit_cap=0.0,
+):
+
+    decode_mla(q, k_buffer, v_buffer, sm_scale, num_kv_splits, req_to_token, b_seq_len, attn_logits)
+    _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, b_seq_len,
+                                num_kv_splits)
+
+
+import vllm.attention.ops.triton_decode_attention
+vllm.attention.ops.triton_decode_attention.decode_attention_fwd_grouped = decode_attention_fwd_grouped
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_triton_mla.py b/vllm_musa/vllm_musa/patch/patch_triton_mla.py
new file mode 100644
index 000000000..ad5719fab
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_triton_mla.py
@@ -0,0 +1,753 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from collections import defaultdict
+from contextlib import contextmanager
+from dataclasses import dataclass
+from itertools import accumulate
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
+
+from vllm.multimodal import MultiModalPlaceholderMap
+
+try:
+    from flashinfer import BatchDecodeMlaWithPagedKVCacheWrapper
+    FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
+except ImportError:
+    BatchDecodeMlaWithPagedKVCacheWrapper = None
+    FLASHINFER_WORKSPACE_BUFFER_SIZE = 0
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.attention.backends.abstract import (AttentionBackend,
+                                              AttentionMetadata,
+                                              AttentionMetadataBuilder,
+                                              AttentionState, AttentionType)
+from vllm.attention.backends.mla.common import MLACommonMetadata
+from vllm_musa.patch.patch_mla_common import MLACommonImpl
+
+from vllm.attention.backends.utils import (PAD_SLOT_ID, compute_slot_mapping,
+                                           compute_slot_mapping_start_idx,
+                                           is_block_tables_empty)
+from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
+from vllm.utils import async_tensor_h2d, make_tensor_with_pad
+
+if TYPE_CHECKING:
+    from vllm.worker.model_runner import (ModelInputForGPUBuilder,
+                                          ModelInputForGPUWithSamplingMetadata)
+
+
+class TritonMLABackend(AttentionBackend):
+
+    @staticmethod
+    def get_name() -> str:
+        return "TRITON_MLA"
+
+    @staticmethod
+    def get_impl_cls() -> Type["TritonMLAImpl"]:
+        return TritonMLAImpl
+
+    @staticmethod
+    def get_metadata_cls() -> Type["AttentionMetadata"]:
+        return TritonMLAMetadata
+
+    @staticmethod
+    def get_builder_cls() -> Type["TritonMLAMetadataBuilder"]:
+        return TritonMLAMetadataBuilder
+
+    @staticmethod
+    def get_state_cls() -> Type["TritonMLAState"]:
+        return TritonMLAState
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,  # assumed to be 1 for MLA
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        return (num_blocks, block_size, head_size)
+
+    @staticmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: torch.Tensor,
+    ) -> None:
+        ops.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
+
+    @staticmethod
+    def copy_blocks(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: torch.Tensor,
+    ) -> None:
+        ops.copy_blocks_mla(kv_caches, src_to_dists)
+
+    @staticmethod
+    def get_supported_head_sizes() -> List[int]:
+        return [576]
+
+
+class TritonMLAState(AttentionState):
+
+    def __init__(self, runner):
+        self.runner = runner
+        self._is_graph_capturing = False
+
+    @contextmanager
+    def graph_capture(self, max_batch_size: int):
+        self._is_graph_capturing = True
+
+        self._graph_slot_mapping = torch.full((max_batch_size, ),
+                                              PAD_SLOT_ID,
+                                              dtype=torch.long,
+                                              device=self.runner.device)
+        self._graph_seq_lens = torch.ones(max_batch_size,
+                                          dtype=torch.int32,
+                                          device=self.runner.device)
+        self._graph_block_tables = torch.from_numpy(
+            self.runner.graph_block_tables).to(device=self.runner.device)
+
+        self._positions = torch.zeros((max_batch_size, ),
+                                      dtype=torch.long,
+                                      device=self.runner.device)
+
+        yield
+
+        self._is_graph_capturing = False
+        del self._graph_slot_mapping
+        del self._graph_seq_lens
+        del self._graph_block_tables
+        del self._positions
+
+    def graph_clone(self, batch_size: int):
+        assert self._is_graph_capturing
+        return self.__class__(self.runner)
+
+    def graph_capture_get_metadata_for_batch(
+            self, batch_size: int, is_encoder_decoder_model: bool = False):
+        assert self._is_graph_capturing
+
+        attn_metadata = self.runner.attn_backend.make_metadata(
+            num_prefills=0,
+            num_prefill_tokens=0,
+            num_decode_tokens=batch_size,
+            slot_mapping=self._graph_slot_mapping[:batch_size],
+            multi_modal_placeholder_index_maps=None,
+            enable_kv_scales_calculation=True,
+            seq_lens=None,
+            seq_lens_tensor=self._graph_seq_lens[:batch_size],
+            max_query_len=1,
+            max_decode_query_len=1,
+            max_prefill_seq_len=0,
+            max_decode_seq_len=self.runner.max_seq_len_to_capture,
+            query_start_loc=None,
+            seq_start_loc=None,
+            context_lens_tensor=None,
+            block_tables=self._graph_block_tables[:batch_size],
+            use_cuda_graph=True,
+            input_positions=self._positions[:batch_size],
+            head_dim=self.runner.model_config.get_head_size())
+
+        if is_encoder_decoder_model:
+            raise NotImplementedError(
+                "TritonMLAState does not support encoder/decoder yet")
+
+        return attn_metadata
+
+    def get_graph_input_buffers(self,
+                                attn_metadata,
+                                is_encoder_decoder_model: bool = False):
+        input_buffers = {
+            "slot_mapping": attn_metadata.slot_mapping,
+            "seq_lens_tensor": attn_metadata.decode_metadata.seq_lens_tensor,
+            "block_tables": attn_metadata.decode_metadata.block_tables,
+            "input_positions": attn_metadata.decode_metadata.input_positions,
+        }
+        if is_encoder_decoder_model:
+            raise NotImplementedError(
+                "TritonMLAState does not support encoder/decoder yet")
+
+        return input_buffers
+
+    def prepare_graph_input_buffers(self,
+                                    input_buffers,
+                                    attn_metadata,
+                                    is_encoder_decoder_model: bool = False):
+        input_positions = attn_metadata.input_positions
+        num_positions = input_positions.shape[0]
+        input_buffers["seq_lens_tensor"].copy_(
+            attn_metadata.decode_metadata.seq_lens_tensor, non_blocking=True)
+        input_buffers["block_tables"].copy_(
+            attn_metadata.decode_metadata.block_tables, non_blocking=True)
+        # CUDA graph buffer is padded so only perform a partial copy based on
+        # num_positions
+        input_buffers["input_positions"][:num_positions].copy_(
+            input_positions, non_blocking=True)
+        if is_encoder_decoder_model:
+            raise NotImplementedError(
+                "TritonMLAState does not support encoder/decoder yet")
+
+    def begin_forward(self, model_input):
+        return
+
+
+@dataclass
+class TritonMLAMetadata(MLACommonMetadata):
+    """Metadata for TritonMLAMetadata.
+
+    NOTE: Any python object stored here is not updated when it is
+    cuda-graph replayed. If you have values that need to be changed
+    dynamically, it should be stored in tensor. The tensor has to be
+    updated from `CUDAGraphRunner.forward` API.
+    """
+    # (batch_size,). The sequence length per sequence. Sequence length means
+    # the computed tokens + new tokens None if it is a decoding.
+    seq_lens: Optional[List[int]]
+    # seq_lens stored as a tensor.
+    seq_lens_tensor: Optional[torch.Tensor]
+
+    # NOTE(sang): Definition of context_len, query_len, and seq_len.
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seq_len ---------------------|
+    #                                   |-- query_len ---|
+
+    # Maximum sequence length among prefill batch. 0 if there are decoding
+    # requests only.
+    max_prefill_seq_len: int
+    # Maximum sequence length among decode batch. 0 if there are prefill
+    # requests only.
+    max_decode_seq_len: int
+    # (batch_size,) A tensor of context lengths (tokens that are computed
+    # so far).
+    context_lens_tensor: Optional[torch.Tensor]
+
+    # (batch_size, max_blocks_per_seq).
+    # Block addresses per sequence. (Seq id -> list of physical block)
+    # E.g., [0, 1, 2] means tokens are stored in 0th, 1st, and 2nd blocks
+    # in the kv cache. Each block can contain up to block_size tokens.
+    # 2nd dimensions are padded up to max_blocks_per_seq if it is cuda-graph
+    # captured.
+    block_tables: Optional[torch.Tensor]
+
+    # Whether or not if cuda graph is enabled.
+    # Cuda-graph is currently enabled for decoding only.
+    # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
+
+    use_cuda_graph: bool
+
+    # Maximum query length in the batch.
+    max_query_len: Optional[int] = None
+
+    # Max number of query tokens among request in the batch.
+    max_decode_query_len: Optional[int] = None
+
+    # (batch_size + 1,). The cumulative subquery lengths of the sequences in
+    # the batch, used to index into subquery. E.g., if the subquery length
+    # is [4, 6], it is [0, 4, 10].
+    query_start_loc: Optional[torch.Tensor] = None
+    # (batch_size + 1,). The cumulative sequence lengths of the sequences in
+    # the batch, used to index into sequence. E.g., if the sequence length is
+    # [4, 6], it is [0, 4, 10].
+    seq_start_loc: Optional[torch.Tensor] = None
+
+    _cached_prefill_metadata: Optional["TritonMLAMetadata"] = None
+    _cached_decode_metadata: Optional["TritonMLAMetadata"] = None
+
+    num_prefill_tokens: int
+
+    num_kv_splits: int = 4  # TODO(lucas) add heuristic
+    attn_logits: Optional[torch.Tensor] = None
+    req_idx: Optional[torch.Tensor] = None
+
+    # The dimension of the attention heads
+    head_dim: Optional[int] = None
+
+    def __post_init__(self):
+        supported_head_sizes = TritonMLABackend.get_supported_head_sizes()
+        if self.head_dim is not None and self.head_dim \
+                not in supported_head_sizes:
+            raise ValueError(
+                f"Only {supported_head_sizes} are supported for head_dim,",
+                f"received {self.head_dim}.")
+
+    @property
+    def prefill_metadata(self) -> Optional["TritonMLAMetadata"]:
+        if self.num_prefills == 0:
+            return None
+
+        if self._cached_prefill_metadata is not None:
+            return self._cached_prefill_metadata
+
+        assert self.seq_lens is not None
+        assert self.seq_lens_tensor is not None
+
+        # Compute some attn_metadata fields which default to None
+        query_start_loc = (None if self.query_start_loc is None else
+                           self.query_start_loc[:self.num_prefills + 1])
+        slot_mapping = (None if self.slot_mapping is None else
+                        self.slot_mapping[:self.num_prefill_tokens])
+        seq_lens = (None if self.seq_lens is None else
+                    self.seq_lens[:self.num_prefills])
+        seq_lens_tensor = (None if self.seq_lens_tensor is None else
+                           self.seq_lens_tensor[:self.num_prefills])
+        seq_start_loc = (None if self.seq_start_loc is None else
+                         self.seq_start_loc[:self.num_prefills + 1])
+        context_lens_tensor = (None if self.context_lens_tensor is None else
+                               self.context_lens_tensor[:self.num_prefills])
+        block_tables = (None if self.block_tables is None else
+                        self.block_tables[:self.num_prefills])
+        input_positions = (None if self.input_positions is None else
+                           self.input_positions[:self.num_prefill_tokens])
+
+        self._cached_prefill_metadata = TritonMLAMetadata(
+            num_prefills=self.num_prefills,
+            num_prefill_tokens=self.num_prefill_tokens,
+            num_decode_tokens=0,
+            slot_mapping=slot_mapping,
+            multi_modal_placeholder_index_maps=self.
+            multi_modal_placeholder_index_maps,
+            enable_kv_scales_calculation=self.enable_kv_scales_calculation,
+            input_positions=input_positions,
+            seq_lens=seq_lens,
+            seq_lens_tensor=seq_lens_tensor,
+            max_query_len=self.max_query_len,
+            max_prefill_seq_len=self.max_prefill_seq_len,
+            max_decode_query_len=0,
+            max_decode_seq_len=0,
+            query_start_loc=query_start_loc,
+            seq_start_loc=seq_start_loc,
+            context_lens_tensor=context_lens_tensor,
+            block_tables=block_tables,
+            use_cuda_graph=False,
+            head_dim=self.head_dim)
+        return self._cached_prefill_metadata
+
+    @property
+    def decode_metadata(self) -> Optional["TritonMLAMetadata"]:
+        if self.num_decode_tokens == 0:
+            return None
+
+        if self._cached_decode_metadata is not None:
+            return self._cached_decode_metadata
+        assert self.seq_lens_tensor is not None
+
+        # Compute some attn_metadata fields which default to None
+        slot_mapping = (None if self.slot_mapping is None else
+                        self.slot_mapping[self.num_prefill_tokens:])
+        seq_lens_tensor = (None if self.seq_lens_tensor is None else
+                           self.seq_lens_tensor[self.num_prefills:])
+        block_tables = (None if self.block_tables is None else
+                        self.block_tables[self.num_prefills:])
+        input_positions = (None if self.input_positions is None else
+                           self.input_positions[self.num_prefill_tokens:])
+
+        self._cached_decode_metadata = TritonMLAMetadata(
+            num_prefills=0,
+            num_prefill_tokens=0,
+            num_decode_tokens=self.num_decode_tokens,
+            slot_mapping=slot_mapping,
+            multi_modal_placeholder_index_maps=None,
+            enable_kv_scales_calculation=True,
+            seq_lens=None,
+            seq_lens_tensor=seq_lens_tensor,
+            max_decode_query_len=self.max_decode_query_len,
+            max_query_len=self.max_query_len,
+            max_prefill_seq_len=0,
+            max_decode_seq_len=self.max_decode_seq_len,
+            # Batch may be composed of prefill|decodes, adjust query start
+            # indices to refer to the start of decodes. E.g.
+            # in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].
+            query_start_loc=(self.query_start_loc[self.num_prefills:] -
+                             self.query_start_loc[self.num_prefills])
+            if self.query_start_loc is not None else None,
+            seq_start_loc=self.seq_start_loc[self.num_prefills:]
+            if self.seq_start_loc is not None else None,
+            context_lens_tensor=None,
+            block_tables=block_tables,
+            use_cuda_graph=self.use_cuda_graph,
+            input_positions=input_positions,
+            head_dim=self.head_dim)
+        return self._cached_decode_metadata
+
+    def advance_step(self,
+                     model_input: "ModelInputForGPUWithSamplingMetadata",
+                     sampled_token_ids: Optional[torch.Tensor],
+                     block_size: int,
+                     num_seqs: int,
+                     num_queries: int,
+                     turn_prefills_into_decodes: bool = False):
+        """
+        Update metadata in-place to advance one decode step.
+        """
+        # When using cudagraph, the num_seqs is padded to the next captured
+        # batch sized, but num_queries tracks the actual number of requests in
+        # the batch. For --enforce-eager mode, num_seqs == num_queries
+        if num_seqs != num_queries:
+            assert num_seqs > num_queries
+            assert self.use_cuda_graph
+
+        if turn_prefills_into_decodes:
+            # When Mutli-Step is enabled with Chunked-Prefill, prefills and
+            # decodes are scheduled together. In the first step, all the
+            # prefills turn into decodes. This update reflects that
+            # conversion.
+            assert self.num_decode_tokens + self.num_prefills == num_seqs
+            self.num_decode_tokens += self.num_prefills
+            self.num_prefills = 0
+            self.num_prefill_tokens = 0
+            self.max_prefill_seq_len = 0
+            self.max_query_len = 1
+
+            self.slot_mapping = self.slot_mapping[:num_seqs]
+        else:
+            assert self.seq_lens is not None
+            assert self.max_decode_seq_len == max(self.seq_lens)
+
+        assert self.num_prefills == 0
+        assert self.num_prefill_tokens == 0
+        assert self.num_decode_tokens == num_seqs
+        assert self.slot_mapping.shape == (num_seqs, )
+
+        assert self.seq_lens is not None
+        assert len(self.seq_lens) == num_seqs
+        assert self.seq_lens_tensor is not None
+        assert self.seq_lens_tensor.shape == (num_seqs, )
+        assert self.max_query_len == 1
+        assert self.max_prefill_seq_len == 0
+
+        assert self.query_start_loc is not None
+        assert self.query_start_loc.shape == (num_queries + 1, )
+        assert self.seq_start_loc is not None
+        assert self.seq_start_loc.shape == (num_seqs + 1, )
+
+        assert self.context_lens_tensor is not None
+        assert self.context_lens_tensor.shape == (num_queries, )
+
+        assert self.block_tables is not None
+        assert self.block_tables.shape[0] == num_seqs
+
+        # Update query lengths. Note that we update only queries and not seqs,
+        # since tensors may be padded due to captured cuda graph batch size
+        for i in range(num_queries):
+            self.seq_lens[i] += 1
+        self.max_decode_seq_len = max(self.seq_lens)
+
+        ops.advance_step_flashattn(num_seqs=num_seqs,
+                                   num_queries=num_queries,
+                                   block_size=block_size,
+                                   input_tokens=model_input.input_tokens,
+                                   sampled_token_ids=sampled_token_ids,
+                                   input_positions=model_input.input_positions,
+                                   seq_lens=self.seq_lens_tensor,
+                                   slot_mapping=self.slot_mapping,
+                                   block_tables=self.block_tables)
+
+
+class TritonMLAMetadataBuilder(AttentionMetadataBuilder[TritonMLAMetadata]):
+
+    def __init__(self, input_builder: "ModelInputForGPUBuilder"):
+        self.input_builder = input_builder
+        self.runner = input_builder.runner
+        self.sliding_window = input_builder.sliding_window
+        self.block_size = input_builder.block_size
+
+    def prepare(self):
+        self.slot_mapping: List[int] = []
+        self.prefill_seq_lens: List[int] = []
+        self.context_lens: List[int] = []
+        self.block_tables: List[List[int]] = []
+        self.curr_seq_lens: List[int] = []
+        self.input_positions: List[int] = []
+        self.multimodal_placeholder_maps: Dict[
+            str,
+            MultiModalPlaceholderMap] = defaultdict(MultiModalPlaceholderMap)
+        self.num_prefills = 0
+        self.num_prefill_tokens = 0
+        self.num_decode_tokens = 0
+        self.has_prefix_cache_hit = False
+
+    def _add_seq_group(
+            self, inter_data: "ModelInputForGPUBuilder.InterDataForSeqGroup",
+            chunked_prefill_enabled: bool, prefix_cache_hit: bool):
+        """Add a sequence group to the metadata. Specifically update/append
+        1. context length.
+        2. block table.
+        3. slot mapping.
+        """
+        is_prompt = inter_data.is_prompt
+        block_tables = inter_data.block_tables
+
+        for (seq_id, token_len, seq_len, curr_seq_len, query_len, context_len,
+             curr_sliding_window_block, input_positions) in zip(
+                 inter_data.seq_ids, [len(t) for t in inter_data.input_tokens],
+                 inter_data.orig_seq_lens, inter_data.seq_lens,
+                 inter_data.query_lens, inter_data.context_lens,
+                 inter_data.curr_sliding_window_blocks,
+                 inter_data.input_positions):
+            self.input_positions.extend(input_positions)
+            self.context_lens.append(context_len)
+            if is_prompt:
+                mm_maps = inter_data.multi_modal_placeholder_maps
+                if mm_maps:
+                    for modality, placeholders in mm_maps.items():
+                        self.multimodal_placeholder_maps[modality].extend(
+                            placeholders)
+
+                self.num_prefills += 1
+                self.num_prefill_tokens += token_len
+                self.prefill_seq_lens.append(seq_len)
+            else:
+                self.num_decode_tokens += query_len
+                self.curr_seq_lens.append(curr_seq_len)
+
+            # Compute block table.
+            # TODO(sang): Combine chunked prefill and prefix caching by
+            # only allowing multiple of block_size chunk size.
+            # NOTE: This only works for oooooooxxx style attention.
+            block_table = []
+            if prefix_cache_hit:
+                # NOTE(woosuk): For flash-attn, the block table should
+                # include the entries for the incoming prefill tokens.
+                block_table = block_tables[seq_id]
+            elif ((chunked_prefill_enabled or not is_prompt)
+                  and block_tables is not None):
+                if curr_sliding_window_block == 0:
+                    block_table = block_tables[seq_id]
+                else:
+                    block_table = block_tables[seq_id][
+                        -curr_sliding_window_block:]
+            self.block_tables.append(block_table)
+
+            # Compute slot mapping.
+            is_profile_run = is_block_tables_empty(block_tables)
+            start_idx = compute_slot_mapping_start_idx(is_prompt, query_len,
+                                                       context_len,
+                                                       self.sliding_window)
+            compute_slot_mapping(is_profile_run, self.slot_mapping, seq_id,
+                                 seq_len, context_len, start_idx,
+                                 self.block_size, inter_data.block_tables)
+
+    def _get_graph_runner_block_tables(
+            self, num_seqs: int,
+            block_tables: List[List[int]]) -> torch.Tensor:
+        # The shape of graph_block_tables is
+        # [max batch size, max context len // block size].
+        max_batch_size, max_blocks = self.runner.graph_block_tables.shape
+        assert max_batch_size >= num_seqs
+
+        graph_block_tables = self.runner.graph_block_tables[:num_seqs]
+        for i, block_table in enumerate(block_tables):
+            if block_table:
+                num_blocks = len(block_table)
+                if num_blocks <= max_blocks:
+                    graph_block_tables[i, :num_blocks] = block_table
+                else:
+                    # It may be possible to have more blocks allocated due
+                    # to lookahead slots of multi-step, however, they are
+                    # not used anyway, so can be safely ignored.
+                    graph_block_tables[
+                        i, :max_blocks] = block_table[:max_blocks]
+
+        return torch.from_numpy(graph_block_tables).to(
+            device=self.runner.device, non_blocking=True)
+
+    def build(self, seq_lens: List[int], query_lens: List[int],
+              cuda_graph_pad_size: int, batch_size: int):
+        """Build attention metadata with on-device tensors.
+
+        Args:
+            seq_lens: The maybe padded sequence lengths of the input sequences.
+            query_lens: The query lengths of the input sequences.
+            cuda_graph_pad_size: The padding size for cuda graph.
+                                 -1 if cuda graph is not used.
+            batch_size: The maybe padded batch size.
+        """
+        prefix_cache_hit = any([
+            inter_data.prefix_cache_hit
+            for inter_data in self.input_builder.inter_data_list
+        ])
+        for inter_data in self.input_builder.inter_data_list:
+            self._add_seq_group(inter_data,
+                                self.input_builder.chunked_prefill_enabled,
+                                prefix_cache_hit)
+
+        device = self.runner.device
+        use_captured_graph = cuda_graph_pad_size != -1
+
+        max_query_len = max(query_lens)
+        decode_query_lens = query_lens[self.num_prefills:]
+        if len(decode_query_lens) > 0:
+            max_decode_query_len = max(decode_query_lens)
+        else:
+            max_decode_query_len = 1
+        max_prefill_seq_len = max(self.prefill_seq_lens, default=0)
+        max_decode_seq_len = max(self.curr_seq_lens, default=0)
+        num_decode_tokens = self.num_decode_tokens
+        query_start_loc = list(accumulate(query_lens, initial=0))
+        seq_start_loc = list(accumulate(seq_lens, initial=0))
+
+        num_seqs = len(seq_lens)
+        if use_captured_graph:
+            self.slot_mapping.extend([PAD_SLOT_ID] * cuda_graph_pad_size)
+            self.block_tables.extend([] * cuda_graph_pad_size)
+            num_decode_tokens = batch_size - self.num_prefill_tokens
+            block_tables = self._get_graph_runner_block_tables(
+                num_seqs, self.block_tables)
+        else:
+            block_tables = make_tensor_with_pad(
+                self.block_tables,
+                pad=0,
+                dtype=torch.int,
+                device=device,
+            )
+        assert max_query_len > 0, ("query_lens: {}".format(query_lens))
+
+        assert device is not None
+        context_lens_tensor = async_tensor_h2d(self.context_lens, torch.int,
+                                               device, self.runner.pin_memory)
+        seq_lens_tensor = async_tensor_h2d(seq_lens, torch.int, device,
+                                           self.runner.pin_memory)
+        input_positions = async_tensor_h2d(self.input_positions, torch.long,
+                                           device, self.runner.pin_memory)
+        slot_mapping_tensor = async_tensor_h2d(self.slot_mapping, torch.long,
+                                               device, self.runner.pin_memory)
+        query_start_loc_tensor = async_tensor_h2d(query_start_loc, torch.int32,
+                                                  device,
+                                                  self.runner.pin_memory)
+        seq_start_loc_tensor = async_tensor_h2d(seq_start_loc, torch.int32,
+                                                device, self.runner.pin_memory)
+        placeholder_index_maps = {
+            modality: placeholder_map.index_map()
+            for modality, placeholder_map in
+            self.multimodal_placeholder_maps.items()
+        }
+
+        return TritonMLAMetadata(
+            num_prefills=self.num_prefills,
+            slot_mapping=slot_mapping_tensor,
+            num_prefill_tokens=self.num_prefill_tokens,
+            num_decode_tokens=num_decode_tokens,
+            seq_lens=seq_lens,
+            multi_modal_placeholder_index_maps=placeholder_index_maps,
+            enable_kv_scales_calculation=True,
+            input_positions=input_positions,
+            seq_lens_tensor=seq_lens_tensor,
+            max_query_len=max_query_len,
+            max_decode_query_len=max_decode_query_len,
+            max_prefill_seq_len=max_prefill_seq_len,
+            max_decode_seq_len=max_decode_seq_len,
+            query_start_loc=query_start_loc_tensor,
+            seq_start_loc=seq_start_loc_tensor,
+            context_lens_tensor=context_lens_tensor,
+            block_tables=block_tables,
+            use_cuda_graph=use_captured_graph,
+            num_kv_splits=16,  # TODO:启发式的值，显存足够可以继续增加此值，mla速度会有提升
+            head_dim=self.runner.model_config.get_head_size(),
+        )
+
+
+class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
+
+    def __init__(
+            self,
+            num_heads: int,
+            head_size: int,
+            scale: float,
+            num_kv_heads: int,
+            alibi_slopes: Optional[List[float]],
+            sliding_window: Optional[int],
+            kv_cache_dtype: str,
+            blocksparse_params: Optional[Dict[str, Any]],
+            logits_soft_cap: Optional[float],
+            attn_type: str,
+            # MLA Specific Arguments
+            **kwargs) -> None:
+        super().__init__(num_heads, head_size, scale, num_kv_heads,
+                         alibi_slopes, sliding_window, kv_cache_dtype,
+                         blocksparse_params, logits_soft_cap, attn_type,
+                         **kwargs)
+
+        unsupported_features = [
+            alibi_slopes, sliding_window, blocksparse_params, logits_soft_cap
+        ]
+        if any(unsupported_features):
+            raise NotImplementedError(
+                "TritonMLAImpl does not support one of the following: "
+                "alibi_slopes, sliding_window, blocksparse_params, "
+                "logits_soft_cap")
+
+        if attn_type != AttentionType.DECODER:
+            raise NotImplementedError("Encoder self-attention and "
+                                      "encoder/decoder cross-attention "
+                                      "are not implemented for "
+                                      "TritonMLAImpl")
+
+    def _forward_prefill(
+        self,
+        q: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        attn_metadata: MLACommonMetadata,
+    ) -> torch.Tensor:
+        assert isinstance(attn_metadata, MLACommonMetadata)
+        return self._forward_prefill_flash(q, kv_c_normed, k_pe,
+                                           attn_metadata.seq_start_loc,
+                                           attn_metadata.max_prefill_seq_len, attn_metadata)
+
+    def _forward_decode(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: TritonMLAMetadata,
+    ) -> torch.Tensor:
+        assert kv_c_and_k_pe_cache.numel() > 0
+        if self.kv_cache_dtype.startswith("fp8"):
+            raise NotImplementedError("FP8 Triton MLA not yet supported")
+
+        decode_meta = attn_metadata.decode_metadata
+        assert decode_meta is not None
+        B = q_nope.shape[0]
+
+        q = torch.cat([q_nope, q_pe], dim=-1)
+        o = torch.zeros(B,
+                        self.num_heads,
+                        self.kv_lora_rank,
+                        dtype=q.dtype,
+                        device=q.device)
+
+        num_kv_splits = 4
+
+        # TODO(lucas) Allocate ahead of time
+        attn_logits = torch.zeros(
+            (
+                B,
+                self.num_heads,
+                num_kv_splits,
+                # NOTE(lucas) idk why the +1 is here but sglang has it so we
+                # just mirror that
+                self.kv_lora_rank + 1,
+            ),
+            dtype=torch.float32,
+            device=q.device,
+        )
+
+        # Add a head dim of 1
+        kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)
+        kv_c_cache = kv_c_and_k_pe_cache[..., :self.kv_lora_rank]
+        PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
+
+        # Run MQA
+        decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
+                             decode_meta.block_tables,
+                             decode_meta.seq_lens_tensor, attn_logits,
+                             num_kv_splits, self.scale,
+                             PAGE_SIZE)
+
+        return self._v_up_proj_and_o_proj(o)
+
+import vllm.attention.backends.triton_mla
+vllm.attention.backends.triton_mla.TritonMLAImpl = TritonMLAImpl
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_vllm_utils.py b/vllm_musa/vllm_musa/patch/patch_vllm_utils.py
new file mode 100644
index 000000000..d1c14030f
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_vllm_utils.py
@@ -0,0 +1,65 @@
+import torch
+
+from typing import Callable, List, Optional, Tuple
+
+from torch.library import Library
+from vllm.utils import vllm_lib
+
+
+
+def direct_register_custom_op(
+    op_name: str,
+    op_func: Callable,
+    mutates_args: List[str],
+    fake_impl: Optional[Callable] = None,
+    target_lib: Optional[Library] = None,
+    dispatch_key: str = "MUSA",
+):
+    """
+    `torch.library.custom_op` can have significant overhead because it
+    needs to consider complicated dispatching logic. This function
+    directly registers a custom op and dispatches it to the CUDA backend.
+    See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
+    for more details.
+
+    By default, the custom op is registered to the vLLM library. If you
+    want to register it to a different library, you can pass the library
+    object to the `target_lib` argument.
+
+    IMPORTANT: the lifetime of the operator is tied to the lifetime of the
+    library object. If you want to bind the operator to a different library,
+    make sure the library object is alive when the operator is used.
+    """
+
+    import torch.library
+    if hasattr(torch.library, "infer_schema"):
+        schema_str = torch.library.infer_schema(op_func,
+                                                mutates_args=mutates_args)
+    else:
+        # for pytorch 2.4
+        import torch._custom_op.impl
+        schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args=mutates_args)
+    my_lib = target_lib or vllm_lib
+    my_lib.define(op_name + schema_str)
+    my_lib.impl(op_name, op_func, dispatch_key=dispatch_key)
+
+
+def noop(
+    op_name: str,
+    op_func: Callable,
+    mutates_args: List[str],
+    fake_impl: Optional[Callable] = None,
+    target_lib: Optional[Library] = None,
+    dispatch_key: str = "MUSA",
+    tags: Tuple[torch.Tag, ...] = (),
+):
+    pass
+
+
+def cuda_is_initialized():
+    return True
+
+
+import vllm.utils
+vllm.utils.direct_register_custom_op = noop
+vllm.utils.cuda_is_initialized = cuda_is_initialized
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/patch_vocab_parallel_embedding.py b/vllm_musa/vllm_musa/patch/patch_vocab_parallel_embedding.py
new file mode 100644
index 000000000..9a8189788
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/patch_vocab_parallel_embedding.py
@@ -0,0 +1,101 @@
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+from torch.nn.parameter import Parameter
+
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase, method_has_implemented_embedding)
+import vllm.model_executor.layers.vocab_parallel_embedding
+from vllm.model_executor.utils import set_weight_attrs
+
+DEFAULT_VOCAB_PADDING_SIZE = 64
+
+class UnquantizedEmbeddingMethod(QuantizeMethodBase):
+    """Unquantized method for embeddings."""
+
+    def create_weights(self, layer: torch.nn.Module,
+                       input_size_per_partition: int,
+                       output_partition_sizes: List[int], input_size: int,
+                       output_size: int, params_dtype: torch.dtype,
+                       **extra_weight_attrs):
+        """Create weights for embedding layer."""
+        weight = Parameter(torch.empty(sum(output_partition_sizes),
+                                       input_size_per_partition,
+                                       dtype=params_dtype),
+                           requires_grad=False)
+        set_weight_attrs(weight, {"input_dim": 1, "output_dim": 0})
+        layer.register_parameter("weight", weight)
+        set_weight_attrs(weight, extra_weight_attrs)
+
+    def apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+        return F.linear(x, layer.weight, bias)
+
+    def embedding(self, layer: torch.nn.Module,
+                  input_: torch.Tensor) -> torch.Tensor:
+        return F.embedding(input_, layer.weight)
+
+
+def pad_vocab_size(vocab_size: int,
+                   pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int:
+    """Pad the vocab size to the given value."""
+    return ((vocab_size + pad_to - 1) // pad_to) * pad_to
+
+class VocabEmbedding(torch.nn.Module):
+    def __init__(self,
+                 num_embeddings: int,
+                 embedding_dim: int,
+                 params_dtype: Optional[torch.dtype] = None,
+                 org_num_embeddings: Optional[int] = None,
+                 padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
+                 quant_config: Optional[QuantizationConfig] = None,
+                 prefix: str = ""):
+        super().__init__()
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        self.org_vocab_size = org_num_embeddings or num_embeddings
+        num_added_embeddings = num_embeddings - self.org_vocab_size
+        self.org_vocab_size_padded = pad_vocab_size(self.org_vocab_size, padding_size)
+        self.num_embeddings_padded = pad_vocab_size(self.org_vocab_size_padded + num_added_embeddings, padding_size)
+
+        self.num_embeddings_per_partition = self.num_embeddings_padded
+
+
+        quant_method = None
+        if quant_config is not None:
+            quant_method = quant_config.get_quant_method(self, prefix=prefix)
+        if quant_method is None:
+            quant_method = UnquantizedEmbeddingMethod()
+        is_embedding_layer = type(self) is VocabEmbedding
+        linear_method_implements_embedding = method_has_implemented_embedding(type(quant_method))
+        if is_embedding_layer and not linear_method_implements_embedding:
+            raise NotImplementedError(
+                f"The class {type(quant_method).__name__} must implement the 'embedding' method, see UnquantizedEmbeddingMethod.")
+        self.quant_method: QuantizeMethodBase = quant_method
+
+        if params_dtype is None:
+            params_dtype = torch.get_default_dtype()
+
+
+        self.quant_method.create_weights(self,
+                                          self.embedding_dim,
+                                          [self.num_embeddings_per_partition],
+                                          self.embedding_dim,
+                                          self.num_embeddings_padded,
+                                          params_dtype=params_dtype,
+                                          weight_loader=self.weight_loader)
+
+    def weight_loader(self, param: torch.nn.Parameter, loaded_weight: torch.Tensor):
+
+        param.data.copy_(loaded_weight)
+
+    def forward(self, input_):
+
+        output = self.quant_method.embedding(self, input_)
+        return output
+
+
+vllm.model_executor.layers.vocab_parallel_embedding.VocabParallelEmbedding = VocabEmbedding
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/quantization/__init__.py b/vllm_musa/vllm_musa/patch/quantization/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/patch/quantization/patch_fp8.py b/vllm_musa/vllm_musa/patch/quantization/patch_fp8.py
new file mode 100644
index 000000000..f39924432
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/quantization/patch_fp8.py
@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+import functools
+from vllm.logger import init_logger
+import vllm.model_executor.layers.quantization.utils.fp8_utils
+from vllm.platforms import current_platform
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    Fp8LinearOp, cutlass_block_fp8_supported, cutlass_fp8_supported,
+    all_close_1d, normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
+from vllm.model_executor.layers.quantization.fp8 import (Fp8Config,
+                                                         _is_col_major)
+from vllm_musa.ops.deep_gemm import (get_col_major_tma_aligned_tensor)
+
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+                                                         apply_w8a8_block_fp8_linear)
+import vllm._custom_ops as ops
+from typing import Optional
+
+
+logger = init_logger(__name__)
+
+def Fp8LinearMethod__init__(self, quant_config: Fp8Config):
+        self.quant_config = quant_config
+        self.cutlass_block_fp8_supported = False
+        self.out_dtype = torch.get_default_dtype()
+
+        # For GPUs that lack FP8 hardware support, we can leverage the Marlin
+        # kernel for fast weight-only FP8 quantization
+        self.use_marlin = False
+
+        self.block_quant = self.quant_config.weight_block_size is not None
+        if self.block_quant:
+            # Marlin doesn't support block-wise fp8
+            self.use_marlin = False
+
+        self.fp8_linear = Fp8LinearOp(
+            # Default to using per_token quantization if cutlass is supported
+            use_per_token_if_dynamic=True)
+
+
+def Fp8LinearMethod_apply(self,
+              layer: torch.nn.Module,
+              x: torch.Tensor,
+              bias: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        if self.use_marlin:
+            raise "Marlin is not implemented yet for musa"
+
+        if self.block_quant:
+            assert self.quant_config.weight_block_size is not None
+            return apply_w8a8_block_fp8_linear(
+                input=x,
+                weight=layer.weight,
+                block_size=self.quant_config.weight_block_size,
+                weight_scale=layer.weight_scale_inv,
+                input_scale=layer.input_scale,
+                bias=bias,
+                cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
+            )
+
+        return self.fp8_linear.apply(input=x,
+                                     weight=layer.weight,
+                                     weight_scale=layer.weight_scale,
+                                     out_dtype=self.out_dtype,
+                                     input_scale=layer.input_scale,
+                                     bias=bias)
+
+
+
+def Fp8MoEMethod__init__(self, quant_config: Fp8Config):
+    from vllm.model_executor.layers.fused_moe import fused_experts
+    self.quant_config = quant_config
+    self.block_quant = self.quant_config.weight_block_size is not None
+    self.topk_indices_dtype = None
+    self.rocm_aiter_moe_enabled = False
+    self.use_marlin = False
+
+    # Check for DeepGemm support.
+    self.allow_deep_gemm = False
+    self.fused_experts = functools.partial(  # type: ignore
+            fused_experts,
+            use_fp8_w8a8=True,
+            block_shape=self.quant_config.weight_block_size,
+            allow_deep_gemm=self.allow_deep_gemm)
+
+
+def Fp8MoEMethod_process_weights_after_loading(self, layer: Module) -> None:
+        # TODO (rob): refactor block quant into separate class.
+        if self.block_quant:
+            assert self.quant_config.activation_scheme == "dynamic"
+            w13_weight = layer.w13_weight.data
+            w13_weight_scale_inv = layer.w13_weight_scale_inv.data
+            w2_weight = layer.w2_weight
+            w2_weight_scale_inv = layer.w2_weight_scale_inv
+
+            # torch.compile() cannot use Parameter subclasses.
+            layer.w13_weight = Parameter(w13_weight, requires_grad=False)
+            layer.w13_weight_scale_inv = Parameter(w13_weight_scale_inv,
+                                                   requires_grad=False)
+            layer.w2_weight = Parameter(w2_weight, requires_grad=False)
+            layer.w2_weight_scale_inv = Parameter(w2_weight_scale_inv,
+                                                  requires_grad=False)
+
+            # DeepGemm scales need to be transposed and aligned.  We try to do
+            # it ahead of time for performance reasons.
+            if self.allow_deep_gemm:
+                # Lazy import to avoid CUDA initialization problems.
+                if _is_col_major(layer.w13_weight_scale_inv):
+                    layer.w13_weight_scale_inv = \
+                        get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv).contiguous()
+                if _is_col_major(layer.w2_weight_scale_inv):
+                    layer.w2_weight_scale_inv = \
+                        get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous()
+
+            return
+
+        # If checkpoint is fp16, quantize in place.
+        if not self.quant_config.is_checkpoint_fp8_serialized:
+            fp8_dtype = current_platform.fp8_dtype()
+            w13_weight = torch.empty_like(layer.w13_weight.data,
+                                          dtype=fp8_dtype)
+            w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype)
+
+            # Re-initialize w13_scale because we directly quantize
+            # merged w13 weights and generate a single scaling factor.
+            layer.w13_weight_scale = torch.nn.Parameter(torch.ones(
+                layer.local_num_experts,
+                dtype=torch.float32,
+                device=w13_weight.device),
+                    requires_grad=False)
+            for expert in range(layer.local_num_experts):
+                w13_weight[expert, :, :], layer.w13_weight_scale[
+                    expert] = ops.scaled_fp8_quant(
+                        layer.w13_weight.data[expert, :, :])
+                w2_weight[expert, :, :], layer.w2_weight_scale[
+                    expert] = ops.scaled_fp8_quant(
+                        layer.w2_weight.data[expert, :, :])
+            layer.w13_weight = torch.nn.Parameter(w13_weight,
+                                                  requires_grad=False)
+            layer.w2_weight = torch.nn.Parameter(w2_weight,
+                                                 requires_grad=False)
+            return
+
+        # If checkpoint is fp8, we need to handle that the
+        # MoE kernels require single activation scale and single weight
+        # scale for w13 per expert.
+        else:
+            # Fp8 moe kernels require a single activation scale.
+            # We take the max of all the scales in case they differ.
+            if self.quant_config.activation_scheme == "static":
+                if (layer.w13_input_scale is None
+                        or layer.w2_input_scale is None):
+                    raise ValueError(
+                        "QuantConfig has static quantization, but found "
+                        "activation scales are None.")
+                if (not all_close_1d(layer.w13_input_scale)
+                        or not all_close_1d(layer.w2_input_scale)):
+                    logger.warning_once(
+                        "Found input_scales that are not equal for "
+                        "fp8 MoE layer. Using the maximum across experts "
+                        "for each layer.")
+                layer.w13_input_scale = torch.nn.Parameter(
+                    layer.w13_input_scale.max(), requires_grad=False)
+                layer.w2_input_scale = torch.nn.Parameter(
+                    layer.w2_input_scale.max(), requires_grad=False)
+            if current_platform.is_fp8_fnuz():
+                # Normalize the weights and scales
+                w13_weight, w13_weight_scale, w13_input_scale = \
+                    normalize_e4m3fn_to_e4m3fnuz(
+                        layer.w13_weight, layer.w13_weight_scale,
+                        layer.w13_input_scale)
+                w2_weight, w2_weight_scale, w2_input_scale = \
+                    normalize_e4m3fn_to_e4m3fnuz(
+                        layer.w2_weight, layer.w2_weight_scale,
+                        layer.w2_input_scale)
+                # Reset the parameter
+                layer.w13_weight = torch.nn.Parameter(w13_weight,
+                                                      requires_grad=False)
+                layer.w13_weight_scale = torch.nn.Parameter(
+                    w13_weight_scale, requires_grad=False)
+                if w13_input_scale is not None:
+                    layer.w13_input_scale = torch.nn.Parameter(
+                        w13_input_scale, requires_grad=False)
+                layer.w2_weight = torch.nn.Parameter(w2_weight,
+                                                     requires_grad=False)
+                layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale,
+                                                           requires_grad=False)
+                if w2_input_scale is not None:
+                    layer.w2_input_scale = torch.nn.Parameter(
+                        w2_input_scale, requires_grad=False)
+
+            # Fp8 moe kernel needs single weight scale for w13 per expert.
+            # We take the max then dequant and requant each expert.
+            assert layer.w13_weight_scale is not None
+            shard_size = layer.intermediate_size_per_partition
+            max_w13_scales = layer.w13_weight_scale.max(dim=1).values
+            for expert_id in range(layer.local_num_experts):
+                start = 0
+                for shard_id in range(2):
+                    dq_weight = per_tensor_dequantize(
+                        layer.w13_weight[expert_id][start:start +
+                                                    shard_size, :],
+                        layer.w13_weight_scale[expert_id][shard_id])
+                    layer.w13_weight[expert_id][
+                        start:start + shard_size, :], _ = ops.scaled_fp8_quant(
+                            dq_weight, max_w13_scales[expert_id])
+                    start += shard_size
+
+            layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
+                                                        requires_grad=False)
+            return
+
+
+import vllm.model_executor.layers.quantization.fp8
+vllm.model_executor.layers.quantization.fp8.Fp8LinearMethod.__init__ = Fp8LinearMethod__init__
+vllm.model_executor.layers.quantization.fp8.Fp8LinearMethod.apply = Fp8LinearMethod_apply
+vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod.__init__ = Fp8MoEMethod__init__
+vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod.process_weights_after_loading = Fp8MoEMethod_process_weights_after_loading
diff --git a/vllm_musa/vllm_musa/patch/quantization/patch_fp8_utils.py b/vllm_musa/vllm_musa/patch/quantization/patch_fp8_utils.py
new file mode 100644
index 000000000..10c757f1e
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/quantization/patch_fp8_utils.py
@@ -0,0 +1,127 @@
+from typing import List, Optional
+
+import torch
+import torch.nn.functional as F
+import vllm_musa._musa_custom_ops as ops
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+fp8_dtype = torch.float8_e4m3fn
+fp8_max = torch.finfo(fp8_dtype).max
+fp8_min = -fp8_max
+fused_gemv_support_scale_block = [64, 128]
+
+def ceil_div(x: int, y: int) -> int:
+    return (x + y - 1) // y
+
+def align(x: int, y: int) -> int:
+    return ceil_div(x, y) * y
+
+def sglang_per_token_group_quant_fp8(
+    x: torch.Tensor,
+    group_size: int,
+    eps: float = 1e-10,
+    column_major_scales: bool = False,
+    scale_tma_aligned: bool = False,
+    scale_ue8m0: bool = False,
+):
+    assert (
+        x.shape[-1] % group_size == 0
+    ), "the last dimension of `x` cannot be divisible by `group_size`"
+    assert x.is_contiguous(), "`x` is not contiguous"
+
+    x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype)
+    if scale_ue8m0:
+        assert column_major_scales and scale_tma_aligned
+        x_q_mn, x_q_k = x.shape
+        x_s_mn, x_s_k = x_q_mn, x_q_k // 128
+        aligned_mn = align(x_s_mn, 4)
+        aligned_k = align(x_s_k, 4)
+        # TODO(FIXME): Fix cuda kernel and recover here to empty.
+        x_s = torch.zeros(
+            (aligned_k // 4, aligned_mn),
+            device=x.device,
+            dtype=torch.int,
+        ).transpose(0, 1)[:x_s_mn, :]
+    elif column_major_scales:
+        if scale_tma_aligned:
+            # TODO extract "align" function
+            # aligned to 4 * sizeof(float)
+            aligned_size = (x.shape[-2] + 3) // 4 * 4
+            x_s = torch.empty(
+                x.shape[:-2] + (x.shape[-1] // group_size, aligned_size),
+                device=x.device,
+                dtype=torch.float32,
+            ).permute(-1, -2)[: x.shape[-2], :]
+        else:
+            x_s = torch.empty(
+                (x.shape[-1] // group_size,) + x.shape[:-1],
+                device=x.device,
+                dtype=torch.float32,
+            ).permute(-1, -2)
+    else:
+        x_s = torch.empty(
+            x.shape[:-1] + (x.shape[-1] // group_size,),
+            device=x.device,
+            dtype=torch.float32,
+        )
+    if x.shape[0] > 0:
+        ops.sgl_per_token_group_quant_fp8(
+            x, x_q, x_s, group_size, eps, fp8_min, fp8_max, scale_ue8m0
+        )
+
+    return x_q, x_s
+
+
+def apply_w8a8_block_fp8_linear(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    block_size: List[int],
+    weight_scale: torch.Tensor,
+    input_scale: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+    cutlass_block_fp8_supported: bool = False,
+    use_swigelu: bool = False,
+    use_rms_norm: bool = False,
+    gamma: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    assert input_scale is None
+    input_2d = input.view(-1, input.shape[-1])
+    output_shape = [*input_2d.shape[:-1], weight.shape[0] if not use_swigelu else weight.shape[0] // 2]
+    w_block = ceil_div(weight.shape[-1], weight_scale.shape[-1])
+    fused_gemv_block_fp8_supported = w_block in fused_gemv_support_scale_block
+
+    if input_2d.shape[0] < 4 and fused_gemv_block_fp8_supported:
+        output = ops.fused_gemv(input_2d,
+                                weight,
+                                None,
+                                weight_scale,
+                                use_swigelu=use_swigelu,
+                                use_rms_norm=use_rms_norm,
+                                gamma=gamma,
+                                eps=eps)
+    else:
+        q_input, input_scale = sglang_per_token_group_quant_fp8(input_2d, block_size[0])
+        output = ops.mudnn_w8a8_scaled_mm(q_input,
+                                          weight,
+                                          out_dtype=input.dtype,
+                                          scale_a=input_scale,
+                                          scale_b=weight_scale)
+        if use_swigelu:
+            output = F.swish_glu(output)
+        elif use_rms_norm:
+            output = torch.nn.functional.rms_norm(output, (weight.shape[0], ), gamma, eps)
+
+    if bias is not None:
+        if not use_swigelu:
+            output = output + bias
+        else:
+            logger.warning('When fused swigelu activation, w8a8 block matmul can not add bias.')
+
+    return output.to(dtype=input.dtype).view(*output_shape)
+
+
+import vllm.model_executor.layers.quantization.utils.fp8_utils
+vllm.model_executor.layers.quantization.utils.fp8_utils.apply_w8a8_block_fp8_linear = apply_w8a8_block_fp8_linear
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/quantization/patch_quant_utils.py b/vllm_musa/vllm_musa/patch/quantization/patch_quant_utils.py
new file mode 100644
index 000000000..2822f24e8
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/quantization/patch_quant_utils.py
@@ -0,0 +1,47 @@
+from typing import List, Mapping, Optional, Tuple
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.quant_utils import _normalize_quant_group_shape
+
+def scaled_quantize(
+    x: torch.Tensor,
+    group_shape: Tuple[int, int],
+    quant_dtype: torch.dtype,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    group_shape = _normalize_quant_group_shape(x, group_shape)
+    assert quant_dtype.is_floating_point, \
+        "currently `scaled_quantize` only supports floating point dtypes " \
+        "but could be extended to support other dtypes"
+
+    finfo = torch.finfo(quant_dtype)
+
+    # Reshape (M, N) into (BLK_M, BLOCK_SIZE_M, BLK_N, BLOCK_SIZE_N)
+    assert x.ndim == 2
+    assert x.shape[0] % group_shape[0] == 0 and x.shape[1] % group_shape[1] == 0
+    blk_m, blk_n = x.shape[0] // group_shape[0], x.shape[1] // group_shape[1]
+    x_blkd = x.reshape(blk_m, group_shape[0], blk_n, group_shape[1])
+
+    # Permute to (BLK_M, BLK_N, BLOCK_SIZE_M, BLOCK_SIZE_N)
+    x_blkd_permd = x_blkd.permute(0, 2, 1, 3)
+    # Flatten to (BLK_M, BLK_N, BLOCK_SIZE_M * BLOCK_SIZE_N)
+    x_blkd_permd = x_blkd_permd.flatten(start_dim=2)
+    x_blkd_permd = x_blkd_permd.to(torch.float32)
+
+    # Compute scales
+    min_val, max_val = x_blkd_permd.aminmax(dim=-1)
+    amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
+    scale = finfo.max / amax
+
+    # Apply scale and convert form:
+    # (BLK_M, BLK_N, BLOCK_SIZE_M * BLOCK_SIZE_N) to (M, N)
+    x_scl_sat = (x_blkd_permd * scale.unsqueeze(-1))\
+        .clamp(min=finfo.min, max=finfo.max)\
+        .reshape(blk_m, blk_n, group_shape[0], group_shape[1])\
+        .permute(0, 2, 1, 3)\
+        .reshape(x.shape)
+
+    return x_scl_sat.to(quant_dtype).contiguous(), scale.float().reciprocal()
+
+import vllm.model_executor.layers.quantization.utils.quant_utils
+vllm.model_executor.layers.quantization.utils.quant_utils.scaled_quantize = scaled_quantize
diff --git a/vllm_musa/vllm_musa/patch/v1/__init__.py b/vllm_musa/vllm_musa/patch/v1/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_attention_layer.py b/vllm_musa/vllm_musa/patch/v1/patch_attention_layer.py
new file mode 100644
index 000000000..59dc42dc6
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_attention_layer.py
@@ -0,0 +1,52 @@
+import torch
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.platforms import current_platform
+from vllm_musa.patch.patch_vllm_utils import direct_register_custom_op
+from vllm.attention.layer import wait_for_kv_layer_from_connector, maybe_save_kv_layer_to_connector
+
+def unified_attention_with_output(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    wait_for_kv_layer_from_connector(layer_name)
+    forward_context: ForwardContext = get_forward_context()
+    attn_metadata = forward_context.attn_metadata
+    if isinstance(attn_metadata, dict):
+        attn_metadata = attn_metadata[layer_name]
+    self = forward_context.no_compile_layers[layer_name]
+    kv_cache = self.kv_cache[forward_context.virtual_engine]
+    self.impl.forward(self,
+                      query,
+                      key,
+                      value,
+                      kv_cache,
+                      attn_metadata,
+                      output=output)
+    maybe_save_kv_layer_to_connector(layer_name, kv_cache)
+
+
+def unified_attention_with_output_fake(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    return
+
+
+direct_register_custom_op(
+    op_name="unified_attention_with_output",
+    op_func=unified_attention_with_output,
+    mutates_args=["output"],
+    fake_impl=unified_attention_with_output_fake,
+    dispatch_key=current_platform.dispatch_key,
+)
+
+
+import vllm.attention.layer
+vllm.attention.layer.unified_attention_with_output = unified_attention_with_output
+vllm.attention.layer.unified_attention_with_output_fake = unified_attention_with_output_fake
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_gpu_model_runner.py b/vllm_musa/vllm_musa/patch/v1/patch_gpu_model_runner.py
new file mode 100644
index 000000000..891d7dd73
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_gpu_model_runner.py
@@ -0,0 +1,631 @@
+from vllm.utils import cdiv
+from vllm.v1.worker.block_table import BlockTable
+from vllm.utils import DeviceMemoryProfiler, GiB_bytes
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.models.interfaces import has_step_pooler, is_mixture_of_experts
+from vllm.distributed.parallel_state import prepare_communication_buffer_for_model
+from vllm.distributed.eplb.eplb_state import EplbState
+import torch
+import weakref
+import numpy as np
+from typing import Any, Optional, Union
+from tqdm import tqdm
+import time
+
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+from vllm.attention import get_attn_backend
+from vllm.config import (CompilationLevel, VllmConfig)
+from vllm.compilation.counter import compilation_counter
+from vllm.distributed.parallel_state import (
+    get_pp_group, get_tp_group, graph_capture, is_global_first_rank)
+from vllm.v1.attention.backends.mamba_attn import Mamba2AttentionBackend
+from vllm.v1.kv_cache_interface import (AttentionSpec,
+                                        KVCacheConfig,
+                                        MambaSpec)
+from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT,
+                             ModelRunnerOutput)
+from vllm.distributed.kv_transfer import (get_kv_transfer_group,
+                                          has_kv_transfer_group)
+from vllm.forward_context import set_forward_context
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.spec_decode.eagle import EagleProposer
+from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
+from vllm.sequence import IntermediateTensors
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.v1.sample.logits_processor import LogitsProcessorManager
+from vllm.utils import round_up
+
+logger = init_logger(__name__)
+
+
+class CustomGPUModelRunner(GPUModelRunner):
+    def __init__(
+            self,
+            vllm_config: VllmConfig,
+            device: torch.device,
+        ):
+        super().__init__(vllm_config, device)
+        self.use_cuda_graph = True
+        if self.model_config.enforce_eager:
+            self.use_cuda_graph = False
+        if not self.compilation_config.use_cudagraph:
+            self.use_cuda_graph = False
+        if (self.compilation_config.level != CompilationLevel.PIECEWISE
+                and not self.compilation_config.simple_cuda_graph):
+            self.use_cuda_graph = False
+
+        self.uniform_decode_query_len = 1
+        # Currently, we only support simple cuda graph.
+        if self.use_cuda_graph:
+            self.full_cuda_graph = self.use_cuda_graph
+            self.uniform_decode_query_len = 1 if not self.speculative_config else \
+            1 + self.speculative_config.num_speculative_tokens
+
+    def initialize_attn_backend(self, kv_cache_config: KVCacheConfig) -> None:
+        """
+        Initialize the attention backends and attention metadata builders.
+        """
+        assert len(self.attn_backends) == 0 and len(
+            self.attn_metadata_builders
+        ) == 0, "Attention backends are already initialized"
+        for i, kv_cache_group_spec in enumerate(
+                kv_cache_config.kv_cache_groups):
+            kv_cache_spec = kv_cache_group_spec.kv_cache_spec
+            if isinstance(kv_cache_spec, AttentionSpec):
+                attn_backend_i = get_attn_backend(
+                    kv_cache_spec.head_size,
+                    self.dtype,
+                    kv_cache_spec.dtype,
+                    kv_cache_spec.block_size,
+                    self.model_config.is_attention_free,
+                    use_mla=kv_cache_spec.use_mla,
+                )
+                if attn_backend_i is None:
+                    error_msg = (f"Error with get_attn_backend: "
+                                 f"{kv_cache_spec.head_size=}, "
+                                 f"{self.dtype=}, {kv_cache_spec.dtype=}, "
+                                 f"{kv_cache_spec.block_size=}, "
+                                 f"{self.model_config.is_attention_free=}, "
+                                 f"{kv_cache_spec.use_mla=}")
+                    logger.error(error_msg)
+                    raise NotImplementedError(
+                        "Non-Attention backend is not supported by V1 "
+                        "GPUModelRunner.")
+            elif isinstance(kv_cache_spec, MambaSpec):
+                attn_backend_i = Mamba2AttentionBackend
+            else:
+                raise ValueError(
+                    f"Unknown KV cache spec type: {type(kv_cache_spec)}")
+
+            block_table_i = self.input_batch.block_table[i]
+            attn_metadata_builder_i = attn_backend_i.get_builder_cls()(
+                weakref.proxy(self),
+                kv_cache_spec,
+                block_table_i,
+            )
+
+            if (self.full_cuda_graph
+                    and not attn_metadata_builder_i.full_cudagraph_supported):
+                logger.warning(
+                    f"Full CUDAGraph not supported for "
+                    f"{attn_backend_i.__name__}. Turn off CompilationConfig."
+                    f"full_cuda_graph or use a different attention backend.")
+
+            self.attn_backends.append(attn_backend_i)
+            self.attn_metadata_builders.append(attn_metadata_builder_i)
+
+    def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None:
+        """
+        Initialize KV cache based on `kv_cache_config`.
+        Args:
+            kv_cache_config: Configuration for the KV cache, including the KV
+            cache size of each layer
+        """
+        self.kv_cache_config = kv_cache_config
+        self.may_reinitialize_input_batch(kv_cache_config)
+        self.initialize_attn_backend(kv_cache_config)
+        kv_caches = self.initialize_kv_cache_tensors(kv_cache_config)
+
+        if self.speculative_config and self.speculative_config.use_eagle() and get_pp_group().is_last_rank:
+            assert isinstance(self.drafter, EagleProposer)
+            # validate all draft model layers belong to the same kv cache
+            # group
+            self.drafter.validate_same_kv_cache_group(kv_cache_config)
+
+        if has_kv_transfer_group():
+            get_kv_transfer_group().register_kv_caches(kv_caches)
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+        intermediate_tensors: Optional[IntermediateTensors] = None,
+    ) -> Union[ModelRunnerOutput, IntermediateTensors]:
+        self._update_states(scheduler_output)
+        if not scheduler_output.total_num_scheduled_tokens:
+            if not has_kv_transfer_group():
+                # Return empty ModelRunnerOutput if there's no work to do.
+                return EMPTY_MODEL_RUNNER_OUTPUT
+
+            return self.kv_connector_no_forward(scheduler_output)
+
+        # Prepare the decoder inputs.
+        (attn_metadata, attention_cuda_graphs, logits_indices,
+         spec_decode_metadata,
+         num_scheduled_tokens_np) = (self._prepare_inputs(scheduler_output))
+        num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
+
+        # Note(zhenxue.xu): Currently, only simple full graph is supported during the decode stage:
+        # Case 1: Normal decode — when max_query_len == 1
+        # Case 2: Speculative decode — when max_query_len == 1 + num_speculative_tokens
+        max_query_len = np.max(num_scheduled_tokens_np)
+        use_cuda_graph = self.use_cuda_graph and attention_cuda_graphs \
+                and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1] \
+                and max_query_len == self.uniform_decode_query_len
+
+        if use_cuda_graph:
+            # Use CUDA graphs.
+            # Add padding to the batch size.
+            num_input_tokens = self.vllm_config.pad_for_cudagraph(
+                num_scheduled_tokens)
+        else:
+            # Eager mode.
+            # Pad tokens to multiple of tensor_parallel_size when
+            # enabled collective fusion for SP
+            tp_size = self.vllm_config.parallel_config.tensor_parallel_size
+            if self.compilation_config.pass_config. \
+                enable_sequence_parallelism and tp_size > 1:
+                num_input_tokens = round_up(num_scheduled_tokens, tp_size)
+            else:
+                num_input_tokens = num_scheduled_tokens
+
+        # Padding for DP
+        num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
+        num_input_tokens += num_pad
+
+        # _prepare_inputs may reorder the batch, so we must gather multi
+        # modal outputs after that to ensure the correct order
+        if self.is_multimodal_model:
+            # Run the multimodal encoder if any.
+            self._execute_mm_encoder(scheduler_output)
+            mm_embeds = self._gather_mm_embeddings(scheduler_output)
+        else:
+            mm_embeds = []
+
+        if self.is_multimodal_model and get_pp_group().is_first_rank:
+            # NOTE(woosuk): To unify token ids and soft tokens (vision
+            # embeddings), we always use embeddings (rather than token ids)
+            # as input to the multimodal model, even when the input is text.
+            input_ids = self.input_ids[:num_scheduled_tokens]
+            if mm_embeds:
+                inputs_embeds = self.model.get_input_embeddings(
+                    input_ids, mm_embeds)
+            else:
+                inputs_embeds = self.model.get_input_embeddings(input_ids)
+            # TODO(woosuk): Avoid the copy. Optimize.
+            self.inputs_embeds[:num_scheduled_tokens].copy_(inputs_embeds)
+            inputs_embeds = self.inputs_embeds[:num_input_tokens]
+            input_ids = None
+        else:
+            # For text-only models, we use token ids as input.
+            # While it is possible to use embeddings as input just like the
+            # multimodal models, it is not desirable for performance since
+            # then the embedding layer is not included in the CUDA graph.
+            input_ids = self.input_ids[:num_input_tokens]
+            inputs_embeds = None
+        if self.uses_mrope:
+            positions = self.mrope_positions[:, :num_input_tokens]
+        else:
+            positions = self.positions[:num_input_tokens]
+
+        if get_pp_group().is_first_rank:
+            intermediate_tensors = None
+        else:
+            intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                num_input_tokens, intermediate_tensors, True)
+
+        # Run the model.
+        # Use persistent buffers for CUDA graphs.
+        with set_forward_context(
+                attn_metadata,
+                self.vllm_config,
+                num_tokens=num_input_tokens,
+                num_tokens_across_dp=num_tokens_across_dp,
+                skip_cuda_graphs= not use_cuda_graph):
+            self.maybe_setup_kv_connector(scheduler_output)
+
+            model_output = self.model(
+                input_ids=input_ids,
+                positions=positions,
+                intermediate_tensors=intermediate_tensors,
+                inputs_embeds=inputs_embeds,
+            )
+
+            self.maybe_wait_for_kv_save()
+            finished_sending, finished_recving = (
+                self.get_finished_kv_transfers(scheduler_output))
+
+        if self.use_aux_hidden_state_outputs:
+            hidden_states, aux_hidden_states = model_output
+        else:
+            hidden_states = model_output
+            aux_hidden_states = None
+
+        # Broadcast PP output for external_launcher (torchrun)
+        # to make sure we are synced across pp ranks
+        # TODO: Support overlapping mirco-batches
+        # https://github.com/vllm-project/vllm/issues/18019
+        broadcast_pp_output = \
+            self.parallel_config.distributed_executor_backend \
+            == "external_launcher" and len(get_pp_group().ranks) > 0
+        if not get_pp_group().is_last_rank:
+            # For mid-pipeline stages, return the hidden states.
+            if not broadcast_pp_output:
+                return hidden_states
+            assert isinstance(hidden_states, IntermediateTensors)
+            get_pp_group().send_tensor_dict(hidden_states.tensors,
+                                            all_gather_group=get_tp_group())
+            logits = None
+        else:
+            if self.input_batch.pooling_params:
+                return self._pool(hidden_states, num_scheduled_tokens,
+                                  num_scheduled_tokens_np, finished_sending,
+                                  finished_recving)
+
+            sample_hidden_states = hidden_states[logits_indices]
+            logits = self.model.compute_logits(sample_hidden_states, None)
+        if broadcast_pp_output:
+            model_output_broadcast_data = {
+                "logits": logits.contiguous(),
+            } if logits is not None else {}
+            model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
+                model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
+            assert model_output_broadcast_data is not None
+            logits = model_output_broadcast_data["logits"]
+
+        # Apply structured output bitmasks if present
+        if scheduler_output.grammar_bitmask is not None:
+            self.apply_grammar_bitmask(scheduler_output, logits)
+
+        # Sample the next token and get logprobs if needed.
+        sampling_metadata = self.input_batch.sampling_metadata
+        if spec_decode_metadata is None:
+            sampler_output = self.sampler(
+                logits=logits,
+                sampling_metadata=sampling_metadata,
+            )
+        else:
+            # When indexing with a tensor (bonus_logits_indices), PyTorch
+            # creates a new tensor with separate storage from the original
+            # logits tensor. This means any in-place operations on bonus_logits
+            # won't affect the original logits tensor.
+            assert logits is not None
+            bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
+            sampler_output = self.sampler(
+                logits=bonus_logits,
+                sampling_metadata=sampling_metadata,
+            )
+            bonus_token_ids = sampler_output.sampled_token_ids
+
+            # Just like `bonus_logits`, `target_logits` is a new tensor with
+            # separate storage from the original `logits` tensor. Therefore,
+            # it is safe to update `target_logits` in place.
+            target_logits = logits[spec_decode_metadata.target_logits_indices]
+            output_token_ids = self.rejection_sampler(
+                spec_decode_metadata,
+                None,  # draft_probs
+                target_logits,
+                bonus_token_ids,
+                sampling_metadata,
+            )
+            sampler_output.sampled_token_ids = output_token_ids
+
+        num_nans_in_logits = {}
+        if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
+            num_nans_in_logits = self._get_nans_in_logits(logits)
+
+        # TODO(woosuk): The following loop can be slow since it iterates over
+        # the requests one by one. Optimize.
+        discard_sampled_tokens_req_indices = []
+        for i, req_id in enumerate(self.input_batch.req_ids):
+            req_state = self.requests[req_id]
+            seq_len = (req_state.num_computed_tokens +
+                       scheduler_output.num_scheduled_tokens[req_id])
+            if seq_len < req_state.num_tokens:
+                # Ignore the sampled token for partial prefills.
+                # Rewind the generator state as if the token was not sampled.
+                # This relies on cuda-specific torch-internal impl details
+                generator = self.input_batch.generators.get(i)
+                if generator is not None:
+                    generator.set_offset(generator.get_offset() - 4)
+                # Record the index of the request that should not be sampled,
+                # so that we could clear the sampled tokens before returning.
+                discard_sampled_tokens_req_indices.append(i)
+
+        # NOTE: GPU -> CPU Sync happens here.
+        # Move as many CPU operations as possible before this sync point.
+        logprobs_tensors = sampler_output.logprobs_tensors
+        logprobs_lists = logprobs_tensors.tolists() \
+            if logprobs_tensors is not None else None
+
+        # Compute prompt logprobs if needed.
+        prompt_logprobs_dict = self._get_prompt_logprobs_dict(
+            hidden_states[:num_scheduled_tokens],
+            scheduler_output,
+        )
+
+        # Get the valid generated tokens.
+        sampled_token_ids = sampler_output.sampled_token_ids
+        max_gen_len = sampled_token_ids.shape[-1]
+        if max_gen_len == 1:
+            # No spec decode tokens.
+            valid_sampled_token_ids = sampled_token_ids.tolist()
+        else:
+            # Includes spec decode tokens.
+            valid_sampled_token_ids = self.rejection_sampler.parse_output(
+                sampled_token_ids,
+                self.input_batch.vocab_size,
+            )
+        # Mask out the sampled tokens that should not be sampled.
+        for i in discard_sampled_tokens_req_indices:
+            valid_sampled_token_ids[i].clear()
+
+        # Cache the sampled tokens in the model runner, so that the scheduler
+        # doesn't need to send them back.
+        # NOTE(woosuk): As an exception, when using PP, the scheduler sends
+        # the sampled tokens back, because there's no direct communication
+        # between the first-stage worker and the last-stage worker.
+        for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
+            if not sampled_ids:
+                continue
+
+            start_idx = self.input_batch.num_tokens_no_spec[req_idx]
+            end_idx = start_idx + len(sampled_ids)
+            assert end_idx <= self.max_model_len, (
+                "Sampled token IDs exceed the max model length. "
+                f"Total number of tokens: {end_idx} > max_model_len: "
+                f"{self.max_model_len}")
+
+            self.input_batch.token_ids_cpu[req_idx,
+                                           start_idx:end_idx] = sampled_ids
+            self.input_batch.num_tokens_no_spec[req_idx] = end_idx
+            self.input_batch.num_tokens[req_idx] = end_idx
+            req_id = self.input_batch.req_ids[req_idx]
+            req_state = self.requests[req_id]
+            req_state.output_token_ids.extend(sampled_ids)
+
+        if not self.speculative_config:
+            # Speculative decoding is not enabled.
+            spec_token_ids = None
+        else:
+            spec_token_ids = self.propose_draft_token_ids(
+                scheduler_output,
+                valid_sampled_token_ids,
+                sampling_metadata,
+                hidden_states,
+                sample_hidden_states,
+                aux_hidden_states,
+                spec_decode_metadata,
+                attn_metadata,
+            )
+
+        if has_kv_transfer_group():
+            get_kv_transfer_group().clear_connector_metadata()
+
+        self.eplb_step()
+
+        return ModelRunnerOutput(
+            req_ids=self.input_batch.req_ids,
+            req_id_to_index=self.input_batch.req_id_to_index,
+            sampled_token_ids=valid_sampled_token_ids,
+            spec_token_ids=spec_token_ids,
+            logprobs=logprobs_lists,
+            prompt_logprobs_dict=prompt_logprobs_dict,
+            pooler_output=[],
+            finished_sending=finished_sending,
+            finished_recving=finished_recving,
+            num_nans_in_logits=num_nans_in_logits,
+        )
+
+    @torch.inference_mode()
+    def _dummy_run(
+        self,
+        num_tokens: int,
+        capture_attn_cudagraph: bool = False,
+        skip_eplb: bool = False,
+        is_profile: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+
+        # Padding for DP
+        num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
+        num_tokens += num_pad
+
+        max_query_len = self.uniform_decode_query_len
+
+        # Set num_scheduled_tokens based on num_tokens and max_num_seqs
+        # for dummy run with LoRA so that the num_reqs collectively
+        # has num_tokens in total.
+        assert num_tokens <= self.scheduler_config.max_num_batched_tokens
+        max_num_reqs = self.scheduler_config.max_num_seqs
+        num_reqs = min(num_tokens, max_num_reqs)
+
+        if capture_attn_cudagraph:
+            num_reqs = cdiv(num_tokens, max_query_len)
+
+        min_tokens_per_req = num_tokens // num_reqs
+        num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs
+        num_scheduled_tokens_list[-1] += num_tokens % num_reqs
+        assert sum(num_scheduled_tokens_list) == num_tokens
+        assert len(num_scheduled_tokens_list) == num_reqs
+        num_scheduled_tokens = np.array(num_scheduled_tokens_list,
+                                        dtype=np.int32)
+
+        attn_metadata: Optional[dict[str, Any]] = None
+        if capture_attn_cudagraph:
+            attn_metadata = {}
+
+            query_start_loc = self.query_start_loc[:num_reqs + 1]
+            # Make sure max_model_len is used at the graph capture time.
+            self.seq_lens_np[:num_reqs] = self.max_model_len
+            self.seq_lens_np[num_reqs:] = 0
+            self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
+                                           non_blocking=True)
+            seq_lens = self.seq_lens[:num_reqs]
+
+            common_attn_metadata = CommonAttentionMetadata(
+                query_start_loc=query_start_loc,
+                seq_lens=seq_lens,
+                num_reqs=num_reqs,
+                num_actual_tokens=num_tokens,
+                max_query_len=max_query_len,
+            )
+
+            for kv_cache_group_id, kv_cache_group_spec in enumerate(
+                    self.kv_cache_config.kv_cache_groups):
+
+                attn_metadata_i = self.attn_metadata_builders[
+                    kv_cache_group_id].build_for_cudagraph_capture(
+                        common_attn_metadata)
+                for layer_name in kv_cache_group_spec.layer_names:
+                    attn_metadata[layer_name] = attn_metadata_i
+
+        with self.maybe_dummy_run_with_lora(self.lora_config,
+                                            num_scheduled_tokens):
+            model = self.model
+            if self.is_multimodal_model:
+                input_ids = None
+                inputs_embeds = self.inputs_embeds[:num_tokens]
+            else:
+                input_ids = self.input_ids[:num_tokens]
+                inputs_embeds = None
+            if self.uses_mrope:
+                positions = self.mrope_positions[:, :num_tokens]
+            else:
+                positions = self.positions[:num_tokens]
+
+            if get_pp_group().is_first_rank:
+                intermediate_tensors = None
+            else:
+                if self.intermediate_tensors is None:
+                    self.intermediate_tensors = (
+                        self.model.make_empty_intermediate_tensors(
+                            batch_size=self.max_num_tokens,
+                            dtype=self.model_config.dtype,
+                            device=self.device))
+
+                intermediate_tensors = self.sync_and_slice_intermediate_tensors(
+                    num_tokens, None, False)
+
+            with self.maybe_randomize_inputs(input_ids), set_forward_context(
+                    attn_metadata,
+                    self.vllm_config,
+                    num_tokens=num_tokens,
+                    num_tokens_across_dp=num_tokens_across_dp):
+                outputs = model(
+                    input_ids=input_ids,
+                    positions=positions,
+                    intermediate_tensors=intermediate_tensors,
+                    inputs_embeds=inputs_embeds,
+                )
+            if self.use_aux_hidden_state_outputs:
+                hidden_states, _ = outputs
+            else:
+                hidden_states = outputs
+
+            if self.speculative_config and self.speculative_config.use_eagle() and get_pp_group().is_last_rank:
+                assert isinstance(self.drafter, EagleProposer)
+                self.drafter.dummy_run(num_tokens)
+
+        # This is necessary to avoid blocking DP.
+        # For dummy runs, we typically skip EPLB since we don't have any real
+        # requests to process.
+        # However, in DP settings, there may be cases when some DP ranks do
+        # not have any requests to process, so they're executing dummy batches.
+        # In such cases, we still have to trigger EPLB to make sure
+        # ranks execute the rearrangement in synchronization.
+        if not skip_eplb:
+            self.eplb_step(is_dummy=True, is_profile=is_profile)
+
+        logit_indices = np.cumsum(num_scheduled_tokens) - 1
+        return hidden_states, hidden_states[logit_indices]
+
+    @torch.inference_mode()
+    def _dummy_sampler_run(
+        self,
+        hidden_states: torch.Tensor,
+    ) -> torch.Tensor:
+        # The dummy hidden states may contain special values,
+        # like `inf` or `nan`.
+        # To avoid breaking the sampler, we use a random tensor here instead.
+        hidden_states = torch.rand_like(hidden_states)
+
+        logits = self.model.compute_logits(hidden_states, None)
+        num_reqs = logits.size(0)
+
+        dummy_tensors = lambda v: torch.full(
+            (num_reqs, ), v, device=self.device)
+
+        dummy_metadata = SamplingMetadata(
+            temperature=dummy_tensors(0.5),
+            all_greedy=False,
+            all_random=False,
+            top_p=dummy_tensors(0.9),
+            top_k=dummy_tensors(logits.size(1) - 1),
+            generators={},
+            max_num_logprobs=None,
+            no_penalties=True,
+            prompt_token_ids=None,
+            frequency_penalties=dummy_tensors(0.1),
+            presence_penalties=dummy_tensors(0.1),
+            repetition_penalties=dummy_tensors(0.1),
+            output_token_ids=[[] for _ in range(num_reqs)],
+            allowed_token_ids_mask=None,
+            bad_words_token_ids={},
+            logitsprocs=LogitsProcessorManager(),
+        )
+        try:
+            sampler_output = self.sampler(logits=logits,
+                                          sampling_metadata=dummy_metadata)
+        except RuntimeError as e:
+            if 'out of memory' in str(e):
+                raise RuntimeError(
+                    "CUDA out of memory occurred when warming up sampler with "
+                    f"{num_reqs} dummy requests. Please try lowering "
+                    "`max_num_seqs` or `gpu_memory_utilization` when "
+                    "initializing the engine.") from e
+            else:
+                raise e
+        # if self.speculative_config:
+            # draft_token_ids = [[0] for _ in range(num_reqs)]
+            # dummy_spec_decode_metadata = SpecDecodeMetadata.make_dummy(
+            #     draft_token_ids, self.device)
+
+            # num_tokens = sum(len(ids) for ids in draft_token_ids)
+            # draft_probs = torch.randn(
+            #     num_tokens, logits.shape[-1], device=self.device,
+            #     dtype=logits.dtype)
+            # draft_probs = None
+            # target_logits = torch.randn(num_tokens,
+            #                             logits.shape[-1],
+            #                             device=self.device,
+            #                             dtype=logits.dtype)
+            # NOTE(woosuk): Here, we should use int32 because the sampler uses
+            # int32 for bonus_token_ids. If the dtype mismatches, re-compilation
+            # will occur at runtime.
+            # bonus_token_ids = torch.zeros(num_reqs,
+            #                               device=self.device,
+            #                               dtype=torch.int32)
+            # self.rejection_sampler(
+            #     dummy_spec_decode_metadata,
+            #     draft_probs,
+            #     target_logits,
+            #     bonus_token_ids,
+            #     dummy_metadata,
+            # )
+        return sampler_output
+
+import vllm.v1.worker.gpu_model_runner
+vllm.v1.worker.gpu_model_runner.GPUModelRunner = CustomGPUModelRunner
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_merge_attn_states.py b/vllm_musa/vllm_musa/patch/v1/patch_merge_attn_states.py
new file mode 100644
index 000000000..4ca1b49b2
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_merge_attn_states.py
@@ -0,0 +1,46 @@
+# SPDX-License-Identifier: Apache-2.0
+from typing import Optional
+
+import torch
+
+from vllm.platforms import current_platform
+
+
+def merge_attn_states(
+    output: torch.Tensor,
+    prefix_output: torch.Tensor,
+    prefix_lse: torch.Tensor,
+    suffix_output: torch.Tensor,
+    suffix_lse: torch.Tensor,
+    output_lse: Optional[torch.Tensor] = None,
+) -> None:
+
+    # NOTE(DefTruth): Currently, custom merge_attn_states CUDA kernel
+    # is not support for FP8 dtype, fallback to use Triton kernel.
+    def supported_dtypes(o: torch.Tensor) -> bool:
+        return o.dtype in [torch.float32, torch.half, torch.bfloat16]
+
+    # NOTE(DefTruth): Currently, custom merge_attn_states CUDA
+    # kernel load/store 128b(16 bytes) per memory issue within
+    # thread. Namely, the headsize(headdim) must be multiple of
+    # pack_size (float32 -> 4, half/bfloat16 -> 8).
+    def supported_headdim(o: torch.Tensor) -> bool:
+        headdim = o.shape[2]  # [NUM_TOKENS, NUM_HEADS, HEAD_SIZE]
+        if o.dtype == torch.float32:
+            return headdim % 4 == 0
+        return headdim % 8 == 0
+
+    if (current_platform.is_out_of_tree() and supported_dtypes(output)
+            and supported_headdim(output)):
+        # from vllm._custom_ops import merge_attn_states
+        from vllm_musa._musa_custom_ops import merge_attn_states
+        return merge_attn_states(output, prefix_output, prefix_lse,
+                                 suffix_output, suffix_lse, output_lse)
+    else:
+        from vllm.attention.ops.triton_merge_attn_states import (
+            merge_attn_states)
+        return merge_attn_states(output, prefix_output, prefix_lse,
+                                 suffix_output, suffix_lse, output_lse)
+
+import vllm.attention.ops.merge_attn_states
+vllm.attention.ops.merge_attn_states.merge_attn_states = merge_attn_states
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_mla_common.py b/vllm_musa/vllm_musa/patch/v1/patch_mla_common.py
new file mode 100644
index 000000000..7f9698007
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_mla_common.py
@@ -0,0 +1,665 @@
+
+# SPDX-License-Identifier: Apache-2.0
+
+import functools
+from abc import abstractmethod
+from dataclasses import dataclass
+import math
+from typing import Any, Generic, Optional, TypeVar
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.attention.backends.abstract import (AttentionLayer,
+                                              MLAAttentionImpl)
+
+from vllm.v1.attention.backends.utils import CommonAttentionMetadata
+from vllm.attention.ops.merge_attn_states import merge_attn_states
+
+from vllm.model_executor.layers.linear import (ColumnParallelLinear,
+                                               LinearBase,
+                                               UnquantizedLinearMethod)
+from vllm.v1.attention.backends.mla.common import (MLACommonMetadata,
+                                                   MLACommonMetadataBuilder,
+                                                   )
+
+from vllm.v1.kv_cache_interface import AttentionSpec
+from vllm.v1.worker.block_table import BlockTable
+
+from vllm_musa._musa_custom_ops import flash_attn_fwd
+from vllm.attention.ops.triton_flash_attention import triton_attention
+from vllm_musa.platforms.musa import on_qy2
+
+M = TypeVar("M", bound=MLACommonMetadata)
+
+class CustomMLACommonMetadataBuilder(MLACommonMetadataBuilder):
+
+    def __init__(self,
+                 runner: "GPUModelRunner",
+                 kv_cache_spec: AttentionSpec,
+                 block_table: BlockTable,
+                 metadata_cls: Optional[type[M]] = None):
+        super().__init__(runner, kv_cache_spec, block_table, metadata_cls=metadata_cls)
+
+        # Note(zhenxue.xu): Full CUDA Graph for speculative decode
+        self.aot_schedule = True
+        self.page_size = self.kv_cache_spec.block_size
+
+
+    def reorder_batch(self, input_batch: "InputBatch",
+                      scheduler_output: "SchedulerOutput") -> bool:
+        # We now want to reorder the batch so that the "decode" requests are and
+        # the front and the "prefill" requests are at the using the least amount
+        # swaps possible. (NOTE for now we loosely use "decode" to mean requests
+        # where attention is likely memory-bound and "prefill" to mean requests
+        # where attention is likely compute-bound, TODO(lucas): figure out a
+        # better naming here)
+        decodes = []
+        prefills = []
+        num_decode_tokens = 0
+        num_prefill_tokens = 0
+
+        for i, req_id in enumerate(input_batch.req_ids):
+            num_tokens = scheduler_output.num_scheduled_tokens[req_id]
+            # Note(zhenxue.xu): Support spec decode in full cuda graph
+            if num_tokens in [1, 2]:
+                decodes.append(i)
+                num_decode_tokens += num_tokens
+            else:
+                prefills.append(i)
+                num_prefill_tokens += num_tokens
+
+        # We hope that this is fairly minimal since decodes
+        # should be around for a number of iterations so hopefully they are
+        # relatively stationary (and new request are generally appended to the
+        # persistent batch so already should be at the back)
+        # To achieve this we loop over the decodes in descending order and
+        # the prefills in ascending order. We swap decodes from the  "back"
+        # i.e. past where the last decode should be in the reodorered with
+        # prefills from the front of the batch.
+        # `decodes` and `prefills` are already in ascending order just based on
+        # the above loop
+        num_decodes = len(decodes)
+        num_prefills = len(prefills)
+        modified_batch = False
+
+        for i in range(1, min(num_decodes, num_prefills) + 1):
+            # If the decode is at the "back" of the batch, i, we can swap it
+            # with the prefill closest to the front of the batch
+            decode_idx = decodes[num_decodes - i]
+            if decode_idx < num_decodes:
+                break
+
+            input_batch.swap_states(prefills[i - 1], decode_idx)
+            modified_batch = True
+
+        # Save for next `build` call
+        # TODO(lucas): this is a bit of a hack, we should probably have a
+        # better way of doing this
+        self._num_decodes = num_decodes
+        self._num_prefills = num_prefills
+        self._num_decode_tokens = num_decode_tokens
+        self._num_prefill_tokens = num_prefill_tokens
+
+        return modified_batch
+
+    # Note(zhenxue.xu): Build attention metadata for capturing cudagraph
+    # in normal decode and spec decode
+    def build_for_cudagraph_capture(
+            self, common_attn_metadata: CommonAttentionMetadata) -> M:
+        """
+        This method builds the metadata for full cudagraph capture.
+        Currently, only decode is supported for full cudagraphs with MLA.
+        """
+        m = common_attn_metadata
+        assert m.num_reqs * self.runner.uniform_decode_query_len == m.num_actual_tokens, \
+            "MLA only supports decode-only full CUDAGraph capture. " \
+            "Make sure all cudagraph capture sizes <= max_num_seq."
+
+        m.max_query_len = self.runner.uniform_decode_query_len
+        self._num_decodes = m.num_reqs
+        self._num_decode_tokens = m.num_actual_tokens
+        self._num_prefills = 0
+        self._num_prefill_tokens = 0
+
+        return self.build(0, m)
+
+    # Note(zhenxue.xu): Supports normal and speculative decoding with CUDA Graphs
+    def can_run_in_cudagraph(
+            self, common_attn_metadata: CommonAttentionMetadata) -> bool:
+        return common_attn_metadata.max_query_len == self.runner.uniform_decode_query_len
+
+
+class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
+    """
+    NOTE: Please read the comment at the top of the file before trying to
+    understand this class
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: Optional[list[float]],
+        sliding_window: Optional[int],
+        kv_cache_dtype: str,
+        blocksparse_params: Optional[dict[str, Any]],
+        logits_soft_cap: Optional[float],
+        attn_type: str,
+        kv_sharing_target_layer_name: Optional[str],
+        # MLA Specific Arguments
+        q_lora_rank: Optional[int],
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        kv_b_proj: ColumnParallelLinear,
+    ) -> None:
+        if kv_sharing_target_layer_name is not None:
+            raise NotImplementedError("KV sharing is not supported for MLA")
+
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        self.kv_cache_dtype = kv_cache_dtype
+
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_head_dim = qk_head_dim
+        self.v_head_dim = v_head_dim
+        self.kv_b_proj = kv_b_proj
+        self.triton_fa_func = triton_attention
+
+        # We don't need to pad V if we are on a hopper system with FA3
+        self._pad_v = False
+
+    def _forward_prefill_flash_general(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        max_seqlen_q: int,
+        max_seqlen_k: int,
+        dropout_p: float = 0.0,
+        causal: bool = True,
+        softmax_scale: float = None,
+        return_softmax_lse = False
+    ) -> torch.Tensor:
+        sum_seq, q_h, qk_d = q.shape
+        _, v_h, v_d = v.shape
+        bs = cu_seqlens_q.shape[0] - 1
+
+        max_seqlen = max(max_seqlen_q, max_seqlen_k)
+
+        q_pad = torch.empty(bs, q_h, max_seqlen, qk_d, device=q.device, dtype=q.dtype)
+        k_pad = torch.empty(bs, v_h, max_seqlen, qk_d, device=q.device, dtype=q.dtype)
+        v_pad = torch.empty(bs, v_h, max_seqlen, v_d, device=q.device, dtype=q.dtype)
+        output = torch.empty(sum_seq, q_h, v_d, device=q.device, dtype=q.dtype)
+
+        ops.varlen_fa_seqlen_pad(q, k, v, q_pad, k_pad, v_pad, cu_seqlens_q, cu_seqlens_k, sum_seq, max_seqlen, bs)
+        attn_output, attn_lse, _ = torch.ops.aten._scaled_dot_product_attention_flash_musa(
+            q_pad,
+            k_pad,
+            v_pad,
+            # scale=softmax_scale,
+            is_causal=causal
+        )
+        ops.varlen_fa_seqlen_unpad(attn_output, output, cu_seqlens_q, sum_seq, max_seqlen, v_d, q_h, bs)
+        if return_softmax_lse:
+            attn_lse = attn_lse * math.log(2) # workarround mudnn lse accurate issus
+            lse = torch.empty(q_h, sum_seq, device=q.device, dtype=torch.float32)
+            ops.varlen_fa_lse_seqlen_unpad(attn_lse, lse, cu_seqlens_q, sum_seq, max_seqlen, q_h, bs)
+            return output, lse
+        return output
+
+
+    def _flash_attn_varlen_diff_headdims_triton(self, q, k, v, softmax_scale,
+                                         return_softmax_lse, **kwargs):
+        maybe_padded_v = v
+        if self._pad_v:
+            maybe_padded_v = torch.nn.functional.pad(
+                v, [0, q.shape[-1] - v.shape[-1]], value=0)
+
+        attn_out = self.triton_fa_func(
+            q,
+            k,
+            maybe_padded_v,
+            None,  # output
+            kwargs["cu_seqlens_q"],
+            kwargs["cu_seqlens_k"],
+            kwargs["max_seqlen_q"],
+            kwargs["max_seqlen_k"],
+            kwargs["causal"],
+            softmax_scale,
+            None,  # bias
+        )
+
+        # Unpack the output if there is multiple results,
+        # triton always returns (output, softmax_lse),
+        # vllm_flash_attn returns (output, softmax_lse) when
+        #  `return_softmax_lse = True`
+        # flash_attn (RoCM) returns (output, softmax_lse, ...) when
+        #  `return_attn_probs = True`
+        rest = None
+        if isinstance(attn_out, tuple):
+            attn_out, *rest = attn_out
+
+
+        # unpad if necessary
+        if self._pad_v:
+            attn_out = attn_out[..., :v.shape[-1]]
+
+        # Remain consistent with old `flash_attn_varlen_func` where there
+        # is only one output tensor if `return_softmax_lse` is False.
+        if return_softmax_lse:
+            assert rest is not None
+            return attn_out, rest[0]
+        return attn_out
+
+
+    def _flash_attn_varlen_diff_headdims(self,
+                                         q,
+                                         k,
+                                         v,
+                                         cu_seqlens_q,
+                                         cu_seqlens_k,
+                                         max_seqlen_q,
+                                         max_seqlen_k,
+                                         softmax_scale=None,
+                                         causal=True,
+                                         return_softmax_lse=False):
+        assert on_qy2(), "_flash_attn_varlen_diff_headdims is only supported on PY2."
+        maybe_padded_v = v
+        if self._pad_v:
+            maybe_padded_v = torch.nn.functional.pad(
+                v, [0, q.shape[-1] - v.shape[-1]], value=0)
+
+        # vllm original code
+        # attn_out = self.flash_attn_varlen_func(
+        #     q=q,
+        #     k=k,
+        #     v=maybe_padded_v,
+        #     return_softmax_lse=return_softmax_lse,
+        #     softmax_scale=softmax_scale,
+        #     **kwargs,
+        # )
+
+        # vllm_musa code
+        # JIRA: https://jira.mthreads.com/browse/KUAE-595
+
+        o = torch.empty((q.shape[0], q.shape[1], v.shape[-1]), device=q.device, dtype=q.dtype)
+        attn_out = flash_attn_fwd(q=q,
+                                k=k,
+                                v=maybe_padded_v.contiguous(),
+                                o=o,
+                                cu_q_seqlen=cu_seqlens_q,
+                                cu_kv_seqlen=cu_seqlens_k,
+                                max_seqlen_q=max_seqlen_q,
+                                max_seqlen_kv=max_seqlen_k,
+                                causal=True,
+                                varlen=True,
+                                return_softmax_lse=return_softmax_lse)
+
+        if isinstance(attn_out, tuple):
+            attn_out, lse = attn_out[0], attn_out[1]
+
+        # unpad if necessary
+        if self._pad_v:
+            attn_out = attn_out[..., :v.shape[-1]]
+
+        # Remain consistent with old `flash_attn_varlen_func` where there
+        # is only one output tensor if `return_softmax_lse` is False.
+        if return_softmax_lse:
+            return attn_out, lse
+        return attn_out
+
+    def _v_up_proj(self, x):
+        # Convert from (B, N, L) to (N, B, L)
+        x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
+        # Multiply (N, B, L) x (N, L, V) -> (N, B, V)
+        x = torch.bmm(x, self.W_UV)
+        # Convert from (N, B, V) to (B, N * V)
+        return x.transpose(0, 1).reshape(-1, self.num_heads * self.v_head_dim)
+
+    def process_weights_after_loading(self, act_dtype: torch.dtype):
+
+        def get_layer_weight(layer):
+            WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
+            for attr in WEIGHT_NAMES:
+                if hasattr(layer, attr):
+                    return getattr(layer, attr)
+            raise AttributeError(
+                f"Layer '{layer}' has no recognized weight attribute:"
+                f" {WEIGHT_NAMES}.")
+
+        def get_and_maybe_dequant_weights(layer: LinearBase):
+            if not isinstance(layer.quant_method, UnquantizedLinearMethod):
+                # NOTE: This should only be used offline, since it's O(N^3)
+                eye = torch.eye(layer.input_size_per_partition,
+                                dtype=act_dtype,
+                                device=get_layer_weight(layer).device)
+                dequant_weights = layer.quant_method.apply(layer,
+                                                           eye,
+                                                           bias=None)
+                del eye
+                # standardize to (output, input)
+                return dequant_weights.T
+            return layer.weight
+
+        # we currently do not have quantized bmm's which are needed for
+        # `W_UV` and `W_UK_T`, we we just store fp16/bf16 copies and perform
+        # the bmm's in 16-bit, the extra memory overhead of this is fairly low
+        kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
+        assert kv_b_proj_weight.shape == (
+            self.kv_lora_rank,
+            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim)), (
+                f"{kv_b_proj_weight.shape=}, "
+                f"{self.kv_lora_rank=}, "
+                f"{self.num_heads=}, "
+                f"{self.qk_nope_head_dim=}, "
+                f"{self.v_head_dim=}")
+        kv_b_proj_weight = kv_b_proj_weight.view(
+            self.kv_lora_rank,
+            self.num_heads,
+            self.qk_nope_head_dim + self.v_head_dim,
+        )
+
+        W_UK, W_UV = kv_b_proj_weight.split(
+            [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        # Convert from (L, N, V) to (N, L, V)
+        self.W_UV = W_UV.transpose(0, 1).contiguous()
+        # Convert from (L, N, P) to (N, P, L)
+        self.W_UK_T = W_UK.permute(1, 2, 0)
+
+    def _compute_prefill_context(
+        self,
+        q: torch.Tensor,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: MLACommonMetadata,
+    ):
+        assert attn_metadata.prefill is not None
+        prefill_metadata = attn_metadata.prefill
+        assert prefill_metadata.chunked_context is not None
+
+        output = None
+        iters = len(prefill_metadata.chunked_context.seq_tot)
+        workspace = prefill_metadata.chunked_context.workspace
+
+        for i in range(iters):
+            toks = prefill_metadata.chunked_context.seq_tot[i]
+
+            ops.gather_cache(
+                src_cache=kv_c_and_k_pe_cache,
+                dst=workspace,
+                block_table=prefill_metadata.block_table,
+                cu_seq_lens=prefill_metadata.chunked_context.cu_seq_lens[i],
+                batch_size=attn_metadata.num_prefills,
+                seq_starts=prefill_metadata.chunked_context.starts[i],
+            )
+
+            kv_c_normed = workspace[:toks]\
+                [..., :self.kv_lora_rank].contiguous()
+            k_pe = workspace[:toks]\
+                [..., self.kv_lora_rank:].unsqueeze(1)
+
+            kv_nope = self.kv_b_proj(kv_c_normed)[0].view( \
+                -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
+            k_nope, v = kv_nope\
+                .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+            k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))),
+                          dim=-1)
+
+            # attn_output, attn_softmax_lse = \
+            #     self._flash_attn_varlen_diff_headdims_triton(
+            #     q=q,
+            #     k=k,
+            #     v=v,
+            #     cu_seqlens_q=prefill_metadata.query_start_loc,
+            #     cu_seqlens_k=prefill_metadata.chunked_context.cu_seq_lens[i],
+            #     max_seqlen_q=prefill_metadata.max_query_len,
+            #     max_seqlen_k=prefill_metadata.chunked_context.max_seq_lens[i],
+            #     softmax_scale=self.scale,
+            #     causal=False,  # Context is unmasked
+            #     return_softmax_lse=True,
+            # )
+            if on_qy2():
+                attn_output, attn_softmax_lse = \
+                    self._flash_attn_varlen_diff_headdims(
+                    q,
+                    k,
+                    v,
+                    cu_seqlens_q=prefill_metadata.query_start_loc,
+                    cu_seqlens_k=prefill_metadata.chunked_context.cu_seq_lens[i],
+                    max_seqlen_q=prefill_metadata.max_query_len,
+                    max_seqlen_k=prefill_metadata.chunked_context.max_seq_lens[i],
+                    softmax_scale=self.scale,
+                    causal=False,
+                    return_softmax_lse=True)
+            else:
+                attn_output, attn_softmax_lse = \
+                    self._forward_prefill_flash_general(
+                    q,
+                    k,
+                    v,
+                    cu_seqlens_q=prefill_metadata.query_start_loc,
+                    cu_seqlens_k=prefill_metadata.chunked_context.cu_seq_lens[i],
+                    max_seqlen_q=prefill_metadata.max_query_len,
+                    max_seqlen_k=prefill_metadata.chunked_context.max_seq_lens[i],
+                    softmax_scale=self.scale,
+                    causal=False,
+                    return_softmax_lse=True)
+
+            if output is None:
+                output = attn_output
+                output_lse = attn_softmax_lse
+            else:
+                output_tmp = torch.empty_like(output)
+                output_lse_tmp = torch.empty_like(output_lse)
+                merge_attn_states(
+                    output=output_tmp,
+                    output_lse=output_lse_tmp,
+                    prefix_output=output,
+                    prefix_lse=output_lse,
+                    suffix_output=attn_output,
+                    suffix_lse=attn_softmax_lse,
+                )
+                output = output_tmp
+                output_lse = output_lse_tmp
+
+        return output, output_lse
+
+    def _forward_prefill(
+        self,
+        q: torch.Tensor,
+        kv_c_normed: torch.Tensor,
+        k_pe: torch.Tensor,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: MLACommonMetadata,
+    ) -> torch.Tensor:
+        assert attn_metadata.prefill is not None
+
+        has_context = attn_metadata.prefill.chunked_context is not None
+        # Hardcode: set has_context = False, because suffix_lse is None
+        # JIRA: https://jira.mthreads.com/browse/KUAE-595
+        kv_nope = self.kv_b_proj(kv_c_normed)[0].view(\
+            -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
+        k_nope, v = kv_nope\
+            .split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+        k = torch.cat((k_nope, k_pe.expand((*k_nope.shape[:-1], -1))), dim=-1)
+
+        if on_qy2():
+            output = self._flash_attn_varlen_diff_headdims(
+                q,
+                k,
+                v,
+                cu_seqlens_q=attn_metadata.prefill.query_start_loc,
+                cu_seqlens_k=attn_metadata.prefill.query_start_loc,
+                max_seqlen_q=attn_metadata.prefill.max_query_len,
+                max_seqlen_k=attn_metadata.prefill.max_query_len,
+                softmax_scale=self.scale,
+                causal=True,
+                return_softmax_lse=has_context
+            )
+        else:
+            output = self._forward_prefill_flash_general(
+                q,
+                k,
+                v,
+                cu_seqlens_q=attn_metadata.prefill.query_start_loc,
+                cu_seqlens_k=attn_metadata.prefill.query_start_loc,
+                max_seqlen_q=attn_metadata.prefill.max_query_len,
+                max_seqlen_k=attn_metadata.prefill.max_query_len,
+                softmax_scale=self.scale,
+                causal=True,
+                return_softmax_lse=has_context
+            )
+
+        # output = self._flash_attn_varlen_diff_headdims_triton(
+        #     q=q,
+        #     k=k,
+        #     v=v,
+        #     cu_seqlens_q=attn_metadata.prefill.query_start_loc,
+        #     cu_seqlens_k=attn_metadata.prefill.query_start_loc,
+        #     max_seqlen_q=attn_metadata.prefill.max_query_len,
+        #     max_seqlen_k=attn_metadata.prefill.max_query_len,
+        #     softmax_scale=self.scale,
+        #     causal=True,
+        #     return_softmax_lse=has_context,
+        # )
+
+        if has_context:
+            suffix_output, suffix_lse = output
+            context_output, context_lse = self._compute_prefill_context( \
+                q, kv_c_and_k_pe_cache, attn_metadata)
+
+            output = torch.empty_like(suffix_output)
+            merge_attn_states(
+                output=output,
+                prefix_output=context_output,
+                prefix_lse=context_lse,
+                suffix_output=suffix_output,
+                suffix_lse=suffix_lse,
+            )
+
+        # unpad if necessary
+        if self._pad_v:
+            output = output[..., :v.shape[-1]]
+
+        return output.flatten(start_dim=-2)
+
+    @abstractmethod
+    def _forward_decode(
+        self,
+        ql_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: M,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        q: torch.Tensor,
+        k_c_normed: torch.Tensor,  # key in unified attn
+        k_pe: torch.Tensor,  # value in unified attn
+        kv_cache: torch.Tensor,
+        attn_metadata: M,
+        output: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        assert output is not None, "Output tensor must be provided."
+
+        if attn_metadata is None:
+            # The zero fill is required when used with DP + EP
+            # to ensure all ranks within a DP group compute the
+            # same expert outputs.
+            return output.fill_(0)
+
+        num_actual_toks = attn_metadata.num_actual_tokens
+
+        # Inputs and outputs may be padded for CUDA graphs
+        output_padded = output
+        output = output[:num_actual_toks, ...]
+        q = q[:num_actual_toks, ...]
+        k_c_normed = k_c_normed[:num_actual_toks, ...]
+        k_pe = k_pe[:num_actual_toks, ...]
+
+        assert attn_metadata.num_decodes is not None and \
+            attn_metadata.num_prefills is not None and \
+            attn_metadata.num_decode_tokens is not None
+
+        has_decode = attn_metadata.num_decodes > 0
+        has_prefill = attn_metadata.num_prefills > 0
+        num_decode_tokens = attn_metadata.num_decode_tokens
+
+        decode_q = q[:num_decode_tokens]
+
+        prefill_q = q[num_decode_tokens:]
+        prefill_k_pe = k_pe[num_decode_tokens:]
+        prefill_k_c_normed = k_c_normed[num_decode_tokens:]
+
+        # write the latent and rope to kv cache
+        if kv_cache.numel() > 0:
+            ops.concat_and_cache_mla(
+                k_c_normed,
+                k_pe.squeeze(1),
+                kv_cache,
+                attn_metadata.slot_mapping.flatten(),
+                kv_cache_dtype=self.kv_cache_dtype,
+                scale=layer._k_scale,
+            )
+
+        if has_prefill:
+            output[num_decode_tokens:] = self._forward_prefill(
+                prefill_q, prefill_k_c_normed, prefill_k_pe, kv_cache,
+                attn_metadata)
+
+        if has_decode:
+            assert attn_metadata.decode is not None
+            decode_q_nope, decode_q_pe = decode_q.split(
+                [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+            # # Convert from (B, N, P) to (N, B, P)
+            # decode_q_nope = decode_q_nope.transpose(0, 1)
+            # # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
+            # decode_ql_nope = torch.bmm(decode_q_nope, self.W_UK_T)
+            # # Convert from (N, B, L) to (B, N, L)
+            # decode_ql_nope = decode_ql_nope.transpose(0, 1)
+
+            # Note (zhiguo.qin): mla_mudnn_bmm only verify in dsr1 and dsv2-lite
+            #----------------------------------------------------------------
+            # dsv2-lite:
+            # decode_q_nope torch.Size([1, 16, 128])   stride(2048, 128, 1)
+            # W_UK_T        torch.Size([16, 128, 512]) stride(131072, 512, 1)
+            # o             torch.Size([1, 16, 512])   stride(8192, 512, 1)
+            #----------------------------------------------------------------
+            # dsr1
+            # decode_q_nope torch.Size([1, 16, 128])    stride(2048, 128, 1)
+            # W_UK_T        torch.Size([16, 128, 512])  stride(128, 1, 2048)
+            # o             torch.Size([1, 16, 512])    stride(8192, 512, 1)
+            #----------------------------------------------------------------
+            B, N, P = decode_q_nope.shape
+            L = self.W_UK_T.shape[-1]
+            decode_ql_nope = torch.empty((B, N, L), device=decode_q_nope.device, dtype=decode_q_nope.dtype)
+            ops.mla_mudnn_bmm(decode_ql_nope, decode_q_nope, self.W_UK_T)
+
+            output[:num_decode_tokens] = self._forward_decode(
+                decode_ql_nope, decode_q_pe, kv_cache, attn_metadata)
+
+        return output_padded
+
+import vllm.v1.attention.backends.mla.common
+vllm.v1.attention.backends.mla.common.MLACommonMetadataBuilder = CustomMLACommonMetadataBuilder
+vllm.v1.attention.backends.mla.common.MLACommonImpl = MLACommonImpl
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_moe_layer.py b/vllm_musa/vllm_musa/patch/v1/patch_moe_layer.py
new file mode 100644
index 000000000..765fcfd81
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_moe_layer.py
@@ -0,0 +1,31 @@
+import torch
+from vllm.forward_context import ForwardContext, get_forward_context
+from vllm.platforms import current_platform
+from vllm_musa.patch.patch_vllm_utils import direct_register_custom_op
+
+def moe_forward(hidden_states: torch.Tensor, router_logits: torch.Tensor,
+                layer_name: str) -> torch.Tensor:
+    forward_context: ForwardContext = get_forward_context()
+    self = forward_context.no_compile_layers[layer_name]
+    assert self.quant_method is not None
+
+    return self.forward_impl(hidden_states, router_logits)
+
+
+def moe_forward_fake(hidden_states: torch.Tensor, router_logits: torch.Tensor,
+                     layer_name: str) -> torch.Tensor:
+    return torch.empty_like(hidden_states)
+
+
+direct_register_custom_op(
+    op_name="moe_forward",
+    op_func=moe_forward,
+    mutates_args=[],
+    fake_impl=moe_forward_fake,
+    dispatch_key=current_platform.dispatch_key,
+)
+
+
+import vllm.model_executor.layers.fused_moe.layer
+vllm.model_executor.layers.fused_moe.layer.moe_forward = moe_forward
+vllm.model_executor.layers.fused_moe.layer.moe_forward_fake = moe_forward_fake
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_rocm.py b/vllm_musa/vllm_musa/patch/v1/patch_rocm.py
new file mode 100644
index 000000000..f6157b65a
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_rocm.py
@@ -0,0 +1,6 @@
+
+def on_gfx1x() -> bool:
+    return False
+
+import vllm.platforms.rocm
+vllm.platforms.rocm.on_gfx1x = on_gfx1x
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_sampler.py b/vllm_musa/vllm_musa/patch/v1/patch_sampler.py
new file mode 100644
index 000000000..f25478120
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_sampler.py
@@ -0,0 +1,32 @@
+import torch
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+def random_sample(
+    probs: torch.Tensor,
+    generators: dict[int, torch.Generator],
+) -> torch.Tensor:
+    """Randomly sample from the probabilities.
+
+    We use this function instead of torch.multinomial because torch.multinomial
+    causes CPU-GPU synchronization.
+    """
+    q = torch.empty_like(probs)
+    # NOTE(woosuk): To batch-process the requests without their own seeds,
+    # which is the common case, we first assume that every request does
+    # not have its own seed. Then, we overwrite the values for the requests
+    # that have their own seeds.
+    if len(generators) != probs.shape[0]:
+        q.exponential_()
+    if generators:
+        # TODO(woosuk): This can be slow because we handle each request
+        # one by one. Optimize this.
+        for i, generator in generators.items():
+            q[i].exponential_(generator=generator)
+    q.clamp_min_(torch.finfo(q.dtype).tiny)
+    return probs.div_(q).argmax(dim=-1).view(-1)
+
+
+import vllm.v1.sample.ops.topk_topp_sampler
+vllm.v1.sample.ops.topk_topp_sampler.random_sample = random_sample
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_triton_decode_attention.py b/vllm_musa/vllm_musa/patch/v1/patch_triton_decode_attention.py
new file mode 100644
index 000000000..ba3d10ae2
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_triton_decode_attention.py
@@ -0,0 +1,63 @@
+import torch, torch_musa
+import triton
+from vllm_musa._musa_custom_ops import decode_mla
+from vllm.attention.ops.triton_decode_attention import _fwd_kernel_stage2
+
+
+def _decode_softmax_reducev_fwd(
+    logits,
+    q,
+    o,
+    v_buffer,
+    b_seq_len,
+    num_kv_splits,
+):
+    batch, head_num = q.shape[0], q.shape[1]
+    Lv = v_buffer.shape[-1]
+    BLOCK_DV = triton.next_power_of_2(Lv)
+
+    NUM_KV_SPLITS = num_kv_splits
+
+    extra_kargs = {}
+
+    grid = (batch, head_num)
+    _fwd_kernel_stage2[grid](
+        logits,
+        o,
+        b_seq_len,
+        logits.stride(0),
+        logits.stride(1),
+        logits.stride(2),
+        o.stride(0),
+        o.stride(1),
+        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        BLOCK_DV=BLOCK_DV,
+        Lv=Lv,
+        num_warps=4,
+        num_stages=2,
+        **extra_kargs,
+    )
+
+
+def decode_attention_fwd_grouped(
+    q,
+    k_buffer,
+    v_buffer,
+    o,
+    req_to_token,
+    b_seq_len,
+    attn_logits,
+    num_kv_splits,
+    sm_scale,
+    page_size,
+    logit_cap=0.0,
+):
+    decode_mla(q, k_buffer, v_buffer, sm_scale, num_kv_splits, req_to_token, b_seq_len, attn_logits)
+    # look patch_triton_mla.py
+    if o is not None:
+        _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, b_seq_len,
+                                    num_kv_splits)
+
+
+import vllm.attention.ops.triton_decode_attention
+vllm.attention.ops.triton_decode_attention.decode_attention_fwd_grouped = decode_attention_fwd_grouped
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/patch_triton_mla.py b/vllm_musa/vllm_musa/patch/v1/patch_triton_mla.py
new file mode 100644
index 000000000..3bcae7cae
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/patch_triton_mla.py
@@ -0,0 +1,171 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, Optional
+
+import torch
+
+from vllm.attention.backends.abstract import (AttentionType,
+                                              is_quantized_kv_cache)
+from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
+from vllm.logger import init_logger
+from vllm.v1.attention.backends.mla.common import (MLACommonBackend,
+                                                   MLACommonImpl,
+                                                   MLACommonMetadata)
+from vllm.v1.attention.backends.mla.triton_mla import TritonMLAImpl
+from vllm_musa.platforms.musa import on_ph1, on_qy2
+
+logger = init_logger(__name__)
+
+class CustomTritonMLAImpl(TritonMLAImpl):
+
+    def __init__(
+            self,
+            num_heads: int,
+            head_size: int,
+            scale: float,
+            num_kv_heads: int,
+            alibi_slopes: Optional[list[float]],
+            sliding_window: Optional[int],
+            kv_cache_dtype: str,
+            blocksparse_params: Optional[dict[str, Any]],
+            logits_soft_cap: Optional[float],
+            attn_type: str,
+            kv_sharing_target_layer_name: Optional[str],
+            # MLA Specific Arguments
+            **mla_args) -> None:
+        super().__init__(num_heads, head_size, scale, num_kv_heads,
+                         alibi_slopes, sliding_window, kv_cache_dtype,
+                         blocksparse_params, logits_soft_cap, attn_type,
+                         kv_sharing_target_layer_name, **mla_args)
+
+        if on_ph1():
+            self.num_kv_splits = 1
+            self._forward_decode = self.mp31_forward_decode
+        elif on_qy2():
+            self.num_kv_splits = 4
+            self._forward_decode = self.mp22_forward_decode
+
+    def mp22_forward_decode(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: MLACommonMetadata,
+    ) -> torch.Tensor:
+        assert kv_c_and_k_pe_cache.numel() > 0
+        assert attn_metadata.decode is not None
+
+        if self.kv_cache_dtype.startswith("fp8"):
+            raise NotImplementedError("FP8 Triton MLA not yet supported")
+
+        B = q_nope.shape[0]
+
+        q = torch.cat([q_nope, q_pe], dim=-1)
+
+        o = torch.empty(B,
+                        self.num_heads,
+                        self.kv_lora_rank,
+                        dtype=q.dtype,
+                        device=q.device)
+
+
+        # TODO(lucas) Allocate ahead of time
+        attn_logits = torch.zeros(
+            (
+                B,
+                self.num_heads,
+                self.num_kv_splits,
+                # NOTE(lucas) idk why the +1 is here but sglang has it so we
+                # just mirror that
+                self.kv_lora_rank + 1,
+            ),
+            dtype=torch.float32,
+            device=q.device,
+        )
+
+        # Add a head dim of 1
+        kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)
+        kv_c_cache = kv_c_and_k_pe_cache[..., :self.kv_lora_rank]
+        PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
+
+        # Run MQA
+        decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
+                             attn_metadata.decode.block_table,
+                             attn_metadata.decode.seq_lens, attn_logits,
+                             self.num_kv_splits, self.scale, PAGE_SIZE)
+
+        return self._v_up_proj(o)
+
+    # Note (zhiguo.qin):
+    # When num_kv_splits = 1, the KV cache is not split, meaning there is only one segment to attend over.
+    # In this case, the mp31_decode_mla_kernel can skip calling the _decode_softmax_reducev_fwd Triton kernel,
+    # because no cross-split reduction (e.g., merging attention value contributions from multiple KV splits) is needed.
+    # Skipping this step avoids unnecessary kernel launch overhead and memory operations, improving decoding performance.
+    def mp31_forward_decode(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: MLACommonMetadata,
+    ) -> torch.Tensor:
+        assert kv_c_and_k_pe_cache.numel() > 0
+        assert attn_metadata.decode is not None
+
+        if self.kv_cache_dtype.startswith("fp8"):
+            raise NotImplementedError("FP8 Triton MLA not yet supported")
+
+        B = q_nope.shape[0]
+
+        q = torch.cat([q_nope, q_pe], dim=-1)
+
+        if self.num_kv_splits == 1:
+            o = None
+            attn_logits = torch.zeros(
+                (
+                    B,
+                    self.num_heads,
+                    self.num_kv_splits,
+                    # NOTE(lucas) idk why the +1 is here but sglang has it so we
+                    # just mirror that
+                    self.kv_lora_rank,
+                ),
+                dtype=q.dtype,
+                device=q.device,
+            )
+        else:
+            o = torch.empty(B,
+                            self.num_heads,
+                            self.kv_lora_rank,
+                            dtype=q.dtype,
+                            device=q.device)
+            attn_logits = torch.zeros(
+                (
+                    B,
+                    self.num_heads,
+                    self.num_kv_splits,
+                    # NOTE(lucas) idk why the +1 is here but sglang has it so we
+                    # just mirror that
+                    self.kv_lora_rank + 1,
+                ),
+                dtype=torch.float32,
+                device=q.device,
+            )
+
+        # Add a head dim of 1
+        kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)
+        kv_c_cache = kv_c_and_k_pe_cache[..., :self.kv_lora_rank]
+        PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
+
+        # Run MQA
+        decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
+                             attn_metadata.decode.block_table,
+                             attn_metadata.decode.seq_lens, attn_logits,
+                             self.num_kv_splits, self.scale, PAGE_SIZE)
+
+        if self.num_kv_splits == 1:
+            o = attn_logits
+
+        return self._v_up_proj(o)
+
+import vllm.v1.attention.backends.mla.triton_mla
+vllm.v1.attention.backends.mla.triton_mla.TritonMLAImpl = CustomTritonMLAImpl
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/simple_graph/__init__.py b/vllm_musa/vllm_musa/patch/v1/simple_graph/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/patch/v1/simple_graph/patch_decorators.py b/vllm_musa/vllm_musa/patch/v1/simple_graph/patch_decorators.py
new file mode 100644
index 000000000..61fb360b7
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/simple_graph/patch_decorators.py
@@ -0,0 +1,158 @@
+
+import inspect
+from typing import Callable, Optional, TypeVar, Union, overload
+from unittest.mock import patch
+
+import torch
+import torch.nn as nn
+from torch._dynamo.symbolic_convert import InliningInstructionTranslator
+
+from vllm.compilation.counter import compilation_counter
+from vllm.compilation.wrapper import TorchCompileWrapperWithCustomDispatcher
+from vllm.config import CompilationLevel, VllmConfig
+from vllm.logger import init_logger
+from vllm.sequence import IntermediateTensors
+from vllm.utils import supports_dynamo
+from vllm.compilation.monitor import start_monitoring_torch_compile
+from vllm_musa.patch.v1.simple_graph.wrapper import CudaGraphWrapper
+
+logger = init_logger(__name__)
+
+_T = TypeVar("_T", bound=type[nn.Module])
+
+def _support_torch_compile(
+    cls: _T,
+    dynamic_arg_dims: dict[str, Union[int, list[int]]],
+) -> _T:
+    """
+    A decorator to add support for compiling the forward method of a class.
+    """
+    if TorchCompileWrapperWithCustomDispatcher in cls.__bases__:
+        # support decorating multiple times
+        return cls
+    if CudaGraphWrapper in cls.__bases__:
+        # support decorating multiple times
+        return cls
+
+    # take care of method resolution order
+    # make sure super().__init__ is called on the base class
+    #  other than TorchCompileWrapperWithCustomDispatcher
+    cls.__bases__ = cls.__bases__ + (TorchCompileWrapperWithCustomDispatcher,
+                                     CudaGraphWrapper)
+
+    old_init = cls.__init__
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = '', **kwargs):
+        old_init(self, vllm_config=vllm_config, prefix=prefix, **kwargs)
+        self.vllm_config = vllm_config
+        # for CompilationLevel.DYNAMO_AS_IS , the upper level model runner
+        # will handle the compilation, so we don't need to do anything here.
+        self.do_not_compile = \
+            vllm_config.compilation_config.level in [
+            CompilationLevel.NO_COMPILATION, CompilationLevel.DYNAMO_AS_IS
+        ] or not supports_dynamo()
+        if self.do_not_compile:
+            if vllm_config.compilation_config.simple_cuda_graph:
+                CudaGraphWrapper.__init__(self)
+            return
+
+        compilation_counter.num_models_seen += 1
+        TorchCompileWrapperWithCustomDispatcher.__init__(
+            self, compilation_level=vllm_config.compilation_config.level)
+
+    cls.__init__ = __init__
+
+    def __call__(self, *args, **kwargs):
+        # torch.compiler.is_compiling() means we are inside the compilation
+        # e.g. TPU has the compilation logic in model runner, so we don't
+        # need to compile the model inside.
+
+        # Deleted by qn: NOT support in musa platform
+        # if torch.compiler.is_compiling():
+        #     return self.forward(*args, **kwargs)
+
+        if self.do_not_compile:
+            if not self.vllm_config.compilation_config.simple_cuda_graph:
+                return self.forward(*args, **kwargs)
+            return self.forward_graph(*args, **kwargs)
+
+        # the first compilation needs to have dynamic shapes marked
+        if len(self.compiled_codes) < 1:
+            sig = inspect.signature(self.__class__.forward)
+            bound_args = sig.bind(self, *args, **kwargs)
+            bound_args.apply_defaults()
+            for k, dims in dynamic_arg_dims.items():
+                arg = bound_args.arguments.get(k)
+                if arg is not None:
+                    dims = [dims] if isinstance(dims, int) else dims
+                    if isinstance(arg, torch.Tensor):
+                        # In case dims is specified with negative indexing
+                        dims = [
+                            arg.ndim + dim if dim < 0 else dim for dim in dims
+                        ]
+                        torch._dynamo.mark_dynamic(arg, dims)
+                    elif isinstance(arg, IntermediateTensors):
+                        for tensor in arg.tensors.values():
+                            # In case dims is specified with negative indexing
+                            dims = [
+                                tensor.ndim + dim if dim < 0 else dim
+                                for dim in dims
+                            ]
+                            torch._dynamo.mark_dynamic(tensor, dims)
+                    else:
+                        raise ValueError(
+                            "Unsupported dynamic dimensions"
+                            f" {dims} for argument {k} with type {type(arg)}.")
+            # here, it is the starting point of the `torch.compile` process
+            start_monitoring_torch_compile(self.vllm_config)
+            logger.debug("Start compiling function %s",
+                         self.original_code_object)
+
+        # if we don't use custom dispatcher, we can directly call the
+        # compiled function and let torch.compile handle the dispatching,
+        # with the overhead of guard evaluation and recompilation.
+        if len(self.compiled_codes) < 1 or not self.use_custom_dispatcher:
+            # it seems Dynamo reuse the compilation across instances,
+            # while we need to make sure the compiled code is not reused.
+            # we need to control all the compilation of the model.
+            torch._dynamo.eval_frame.remove_from_cache(
+                self.original_code_object)
+
+            # collect all relevant files traced by Dynamo,
+            # so that the compilation cache can trigger re-compilation
+            # properly when any of these files change.
+
+            # 1. the file containing the top-level forward function
+            self.vllm_config.compilation_config.traced_files.add(
+                self.original_code_object.co_filename)
+
+            # 2. every time Dynamo sees a function call, it will inline
+            # the function by calling InliningInstructionTranslator.inline_call
+            # we hijack this function to know all the functions called
+            # during Dynamo tracing, and their corresponding files
+            inline_call = InliningInstructionTranslator.inline_call
+
+            def patched_inline_call(parent, func, args, kwargs):
+                code = func.get_code()
+                self.vllm_config.compilation_config.traced_files.add(
+                    code.co_filename)
+                return inline_call(parent, func, args, kwargs)
+
+            with patch.object(InliningInstructionTranslator, 'inline_call',
+                              patched_inline_call):
+                output = self.compiled_callable(*args, **kwargs)
+            return output
+
+        # usually, capturing the model once is enough, and then we can
+        # dispatch to the compiled code directly, without going through
+        # the Dynamo guard mechanism.
+        with self.dispatch_to_code(0):
+            model_output = self.forward(*args, **kwargs)
+            return model_output
+
+    cls.__call__ = __call__
+    return cls
+
+
+import vllm.compilation.decorators
+vllm.compilation.decorators._support_torch_compile = _support_torch_compile
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/patch/v1/simple_graph/wrapper.py b/vllm_musa/vllm_musa/patch/v1/simple_graph/wrapper.py
new file mode 100644
index 000000000..e5a980996
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/simple_graph/wrapper.py
@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import sys
+from abc import abstractmethod
+from collections import defaultdict
+from contextlib import contextmanager
+from types import CodeType
+from typing import Callable, Optional
+
+import torch
+
+import vllm.envs as envs
+from vllm.config import CompilationLevel, get_current_vllm_config
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+class CudaGraphWrapper:
+
+    def __init__(self):
+        vllm_config = get_current_vllm_config()
+        self.vllm_config = vllm_config
+        self.compilation_config = vllm_config.compilation_config
+
+        # configs
+        self.cudagraph_capture_sizes = set(
+            self.compilation_config.cudagraph_capture_sizes)
+        self.cudagraph_num_of_warmups = (
+            self.compilation_config.cudagraph_num_of_warmups)
+        assert self.compilation_config.simple_cuda_graph
+        assert self.compilation_config.full_cuda_graph
+
+        # states
+        # batch size -> graph
+        self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
+        self.pool = torch.cuda.graph_pool_handle()
+        # batch size -> hidden states
+        self.hidden_states: dict[int, torch.Tensor] = {}
+        # batch size -> number of warmups
+        self.num_warmups: dict[int, int] = defaultdict(int)
+        # Special flag to handle the first memory profiling run.
+        self.first_run_finished = False
+
+    def capture_graph(self, *args, **kwargs) -> None:
+        batch_size = self._get_batch_size(*args, **kwargs)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph, self.pool):
+            hidden_states = self.forward(*args, **kwargs)
+        self.hidden_states[batch_size] = hidden_states
+        self.graphs[batch_size] = graph
+
+    def forward_graph(self, *args, **kwargs) -> torch.Tensor:
+        if not self.first_run_finished:
+            # Memory profiling run.
+            self.first_run_finished = True
+            return self.forward(*args, **kwargs)
+
+        forward_context = get_forward_context()
+        skip_cuda_graphs = forward_context.skip_cuda_graphs
+        if skip_cuda_graphs:
+            return self.forward(*args, **kwargs)
+
+        batch_size = self._get_batch_size(*args, **kwargs)
+        if batch_size not in self.cudagraph_capture_sizes:
+            # Run in eager mode.
+            return self.forward(*args, **kwargs)
+
+        if self.num_warmups[batch_size] < self.cudagraph_num_of_warmups:
+            # Warmup mode. Run in eager mode.
+            self.num_warmups[batch_size] += 1
+            return self.forward(*args, **kwargs)
+
+        if batch_size not in self.graphs:
+            # Capture the graph.
+            self.capture_graph(*args, **kwargs)
+            return self.hidden_states[batch_size]
+
+        # Run the graph and return the hidden states.
+        graph = self.graphs[batch_size]
+        graph.replay()
+        hidden_states = self.hidden_states[batch_size]
+        return hidden_states
+
+    @abstractmethod
+    def forward(self, *args, **kwargs):
+        ...
+
+    def _get_batch_size(self, *args, **kwargs) -> int:
+        # NOTE(woosuk): Ensure that the keyword arguments here match those
+        # in the model's forward method signature.
+        input_ids = kwargs.get("input_ids")
+        if input_ids is not None:
+            return input_ids.shape[0]
+        input_embeds = kwargs.get("inputs_embeds")
+        if input_embeds is not None:
+            return input_embeds.shape[0]
+        intermediate_tensors = kwargs.get("intermediate_tensors")
+        if intermediate_tensors is not None:
+            return intermediate_tensors.shape[0]
+        # NOTE(woosuk): We don't use the `positions` tensor for batch size
+        # because its first dimension may not be the batch dimension for some
+        # models such as Qwen2.5-VL.
+        if len(args) > 0:
+            # For LoRA models, kwargs could be empty.
+            # FIXME(woosuk): This is a hack. We should find a more robust way
+            # to get the batch size.
+            return args[0].shape[0]
+        raise ValueError("No batch size found in arguments")
diff --git a/vllm_musa/vllm_musa/patch/v1/spec_decode/__init__.py b/vllm_musa/vllm_musa/patch/v1/spec_decode/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/patch/v1/spec_decode/patch_eagle.py b/vllm_musa/vllm_musa/patch/v1/spec_decode/patch_eagle.py
new file mode 100644
index 000000000..26f3136d4
--- /dev/null
+++ b/vllm_musa/vllm_musa/patch/v1/spec_decode/patch_eagle.py
@@ -0,0 +1,65 @@
+import torch
+import torch.nn as nn
+
+from vllm.logger import init_logger
+from vllm.attention.layer import Attention
+from vllm.config import (get_layers_from_vllm_config)
+from vllm.distributed.parallel_state import get_pp_group
+from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.models import supports_multimodal
+
+logger = init_logger(__name__)
+
+
+def load_model(self, target_model: nn.Module) -> None:
+        draft_model_config = \
+            self.vllm_config.speculative_config.draft_model_config
+        target_attn_layer_names = set(
+            get_layers_from_vllm_config(self.vllm_config, Attention).keys())
+
+        from vllm.compilation.backends import set_model_tag
+        with set_model_tag("eagle_head"):
+            self.model = get_model(vllm_config=self.vllm_config,
+                                   model_config=draft_model_config)
+
+        draft_attn_layer_names = (
+            get_layers_from_vllm_config(self.vllm_config, Attention).keys() -
+            target_attn_layer_names)
+
+        self.attn_layer_names = list(draft_attn_layer_names)
+
+        if supports_multimodal(target_model):
+            # handle multimodality
+            self.model.config.image_token_index = (
+                target_model.config.image_token_index)
+            target_language_model = target_model.get_language_model()
+        else:
+            target_language_model = target_model
+        # share embed_tokens with the target model if needed
+        # if get_pp_group().world_size == 1 \
+            #and self.model.model.embed_tokens.weight.shape \
+            #    == target_language_model.model.embed_tokens.weight.shape:
+        if get_pp_group().world_size == 1:
+            logger.info(
+                "Assuming the EAGLE head shares the same vocab embedding" \
+                " with the target model."
+            )
+            #del self.model.model.embed_tokens
+            self.model.model.embed_tokens = (
+                target_language_model.model.embed_tokens)
+        else:
+            logger.info(
+                "The EAGLE head's vocab embedding will be loaded separately" \
+                " from the target model."
+            )
+
+        # share lm_head with the target model if needed
+        # some model definition do not define lm_head explicitly
+        # and reuse embed_tokens for lm_head, e.g., CohereForCausalLM
+        if self.vllm_config.speculative_config.method != "eagle3" and \
+                hasattr(target_language_model, "lm_head"):
+            logger.info("Loading EAGLE LM head weights from the target model.")
+            self.model.lm_head = target_language_model.lm_head
+
+import vllm.v1.spec_decode.eagle
+vllm.v1.spec_decode.eagle.EagleProposer.load_model = load_model
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/platforms/__init__.py b/vllm_musa/vllm_musa/platforms/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/platforms/musa.py b/vllm_musa/vllm_musa/platforms/musa.py
new file mode 100644
index 000000000..5c6886460
--- /dev/null
+++ b/vllm_musa/vllm_musa/platforms/musa.py
@@ -0,0 +1,265 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""A MTGPU platform class.
+Adapted from vllm/platforms/cuda.py
+"""
+
+from functools import cache
+import logging
+import os
+from datetime import timedelta
+from typing import TYPE_CHECKING, Optional, Tuple
+
+import torch
+
+try:
+    import torch_musa  # noqa: F401
+except ImportError:
+    logging.warning("Failed to import torch_musa.")
+
+from torch.distributed import PrefixStore, ProcessGroup
+from torch.distributed.distributed_c10d import is_nccl_available
+
+import vllm.envs as envs
+from vllm.logger import init_logger
+
+from vllm.platforms import Platform, PlatformEnum
+from vllm.platforms.interface import DeviceCapability
+
+if TYPE_CHECKING:
+    from vllm.config import ModelConfig, VllmConfig
+    from vllm.utils import FlexibleArgumentParser
+else:
+    ModelConfig = None
+    VllmConfig = None
+    FlexibleArgumentParser = None
+
+logger = init_logger(__name__)
+
+os.environ["RAY_EXPERIMENTAL_NOSET_MTHREADS_VISIBLE_DEVICES"] = "1"
+
+def _device_id_to_physical_device_id(device_id: int) -> int:
+    device_ids = os.getenv("MTHREADS_VISIBLE_DEVICES")
+    if device_ids is None or device_ids == "all":
+        return device_id
+
+    device_ids = device_ids.split(",")
+    if device_ids == [""]:
+        raise RuntimeError("MTHREADS_VISIBLE_DEVICES is set to empty"
+                            "string, which means MUSA MTGPU support is"
+                            "disabled.")
+    physical_device_id = device_ids[device_id]
+    return int(physical_device_id)
+
+@cache
+def on_qy2() -> bool:
+    major, minor = torch.musa.get_device_capability()
+    return DeviceCapability(major, minor).to_int() == 22
+
+@cache
+def on_ph1() -> bool:
+    major, minor = torch.musa.get_device_capability()
+    return DeviceCapability(major, minor).to_int() == 31
+
+class MusaPlatform(Platform):
+
+    _enum = PlatformEnum.OOT
+    device_name: str = "musa"
+    device_type: str = "musa"
+    # simple_compile_backend: str = "mtgpu"
+    ray_device_key: str = "GPU"
+    device_control_env_var: str = "MTHREADS_VISIBLE_DEVICES"
+    dispatch_key: str = "MUSA" # TODO added by qn: PrivateUse1 or MUSA?
+
+    def is_sleep_mode_available(self) -> bool:
+        return False
+
+    @classmethod
+    def pre_register_and_update(cls,
+                                parser: Optional[FlexibleArgumentParser] = None
+                                ) -> None:
+        # Adapt the global patch here.
+        pass
+
+    @classmethod
+    def get_device_capability(cls, device_id: int = 0) -> Optional[DeviceCapability]:
+        major, minor = torch.musa.get_device_capability()
+        return None
+
+    @classmethod
+    def device_id_to_physical_device_id(cls, device_id: int):
+        return _device_id_to_physical_device_id(device_id)
+
+    @classmethod
+    def get_device_name(cls, device_id: int = 0) -> str:
+        physical_device_id = _device_id_to_physical_device_id(device_id)
+        return torch.musa.get_device_name(physical_device_id) # return 'MTT S4000'
+
+    @classmethod
+    def is_async_output_supported(cls, enforce_eager: Optional[bool]) -> bool:
+        return True
+
+    @classmethod
+    def inference_mode(cls):
+        return torch.inference_mode()
+
+    def is_musa(self) -> bool:
+        return True
+
+    @classmethod
+    def set_device(cls, device: torch.device) -> None:
+        """
+        Set the device for the current platform.
+        """
+        torch.musa.set_device(device)
+        # With this trick we can force the device to be set eagerly
+        # see https://github.com/pytorch/pytorch/issues/155668
+        # for why and when it is needed
+        _ = torch.zeros(1, device=device)
+
+    @classmethod
+    def empty_cache(cls):
+        torch.musa.empty_cache()
+
+    @classmethod
+    def synchronize(cls):
+        torch.musa.synchronize()
+
+    @classmethod
+    def mem_get_info(cls) -> Tuple[int, int]:
+        return torch.musa.mem_get_info()
+
+    @classmethod
+    def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
+        # RayWorkerWrapper monkey patch when setup
+        from vllm_musa.patch import patch_ray  # noqa: F401
+
+        parallel_config = vllm_config.parallel_config
+        if parallel_config and parallel_config.worker_cls == "auto":
+            if envs.VLLM_USE_V1:
+                parallel_config.worker_cls = "vllm_musa.worker.musa_worker.MTGPUWorker"
+            else:
+                if vllm_config.speculative_config:
+                    parallel_config.worker_cls = \
+                    "vllm.spec_decode.spec_decode_worker.create_spec_worker"
+                    parallel_config.sd_worker_cls = \
+                        "vllm_musa.v0.musa_worker.MTGPUWorker"
+                else:
+                    parallel_config.worker_cls = "vllm_musa.v0.musa_worker.MTGPUWorker"
+
+        cache_config = vllm_config.cache_config
+        if cache_config and cache_config.block_size is None:
+            # TODO:
+            # Set block_size to 128 will lead unexpected accuracy issue in mla case.
+            # Please set block_size to 128 back once the problem is fixed.
+            cache_config.block_size = 16
+
+        # PR: https://github.com/vllm-project/vllm/pull/18434
+        if (envs.VLLM_ALL2ALL_BACKEND == "deepep_high_throughput"
+                and parallel_config.data_parallel_size > 1
+                and vllm_config.compilation_config.use_cudagraph):
+            logger.info(
+                "Data Parallel: Forcing enforce eager to be True since DP "
+                "with DeepEP high-throughput kernels are not CUDA Graph "
+                "compatible. The DeepEP low-latency kernels are CUDA Graph "
+                "compatible. Set the all_to_all backend to deepep_low_latency "
+                "to use those kernels instead.")
+            vllm_config.compilation_config.use_cudagraph = False
+            vllm_config.model_config.enforce_eager = True
+            # TODO (varun): Turning this ON gives incorrect results for the
+            # Deepseek-V2-lite model.
+            vllm_config.compilation_config.use_inductor = False
+
+    @classmethod
+    def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
+                             kv_cache_dtype, block_size, use_v1, use_mla):
+        if use_v1:
+            if use_mla:
+                return "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"
+            else:
+                if block_size == 32:
+                    return "vllm_musa.attention.flash_attn.FlashAttentionBackend"
+                else:
+                    raise ValueError("On MUSA platform, V1 flash_attn block_size must be 32.")
+        else:
+            if use_mla:
+                return "vllm.attention.backends.triton_mla.TritonMLABackend"
+            else:
+                return "vllm_musa.v0.flash_attn.FlashAttentionBackend"
+                # if block_size == 64:
+                #     return "vllm_musa.v0.flash_attn.FlashAttentionBackend"
+                # else:
+                #     return "vllm_musa.v0.paged_attention.MUSAFlashAttentionBackend"
+
+
+    @classmethod
+    def get_current_memory_usage(cls,
+                                 device: Optional[torch.types.Device] = None
+                                 ) -> float:
+        torch.musa.empty_cache()
+        torch.musa.reset_peak_memory_stats(device)
+        return torch.musa.max_memory_allocated(device)
+
+    @classmethod
+    def get_device_communicator_cls(cls) -> str:
+        return "vllm_musa.distributed.device_communicators.musa_communicator.MusaCommunicator"
+
+    @classmethod
+    def use_custom_allreduce(cls) -> bool:
+        return True
+
+    @classmethod
+    def stateless_init_device_torch_dist_pg(
+        cls,
+        backend: str,
+        prefix_store: PrefixStore,
+        group_rank: int,
+        group_size: int,
+        timeout: timedelta,
+    ) -> ProcessGroup:
+        assert is_nccl_available()
+        raise RuntimeError(f"Unsupported torch distributed backend: {backend}")
+        # PR: https://github.com/vllm-project/vllm/pull/18763
+        # pg: ProcessGroup = ProcessGroup(
+        #     prefix_store,
+        #     group_rank,
+        #     group_size,
+        # )
+        # from torch.distributed.distributed_c10d import ProcessGroupNCCL
+
+        # backend_options = ProcessGroupNCCL.Options()
+        # backend_options._timeout = timeout
+
+        # backend_class = ProcessGroupNCCL(prefix_store, group_rank, group_size,
+        #                                  backend_options)
+        # backend_type = ProcessGroup.BackendType.NCCL
+        # device = torch.device("cuda")
+        # pg._set_default_backend(backend_type)
+        # backend_class._set_sequence_number_for_group()
+
+        # pg._register_backend(device, backend_type, backend_class)
+        # return pg
+
+    @classmethod
+    def is_pin_memory_available(cls):
+        return True
+
+    @classmethod
+    def supports_v1(cls, model_config: ModelConfig) -> bool:
+        """Returns whether the current platform can support v1 for the supplied
+        model configuration.
+        """
+        return True
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/v0/README.md b/vllm_musa/vllm_musa/v0/README.md
new file mode 100644
index 000000000..b14be8d4d
--- /dev/null
+++ b/vllm_musa/vllm_musa/v0/README.md
@@ -0,0 +1,3 @@
+# Notice
+
+V0 engine is deprecated, code will be removed in future.
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/v0/__init__.py b/vllm_musa/vllm_musa/v0/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/v0/flash_attn.py b/vllm_musa/vllm_musa/v0/flash_attn.py
new file mode 100644
index 000000000..247ea3d12
--- /dev/null
+++ b/vllm_musa/vllm_musa/v0/flash_attn.py
@@ -0,0 +1,785 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Attention layer MUSA GPUs."""
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
+
+import torch
+
+from vllm_musa import _musa_custom_ops as ops
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
+                                              AttentionLayer,
+                                              AttentionMetadata, AttentionType)
+from vllm.attention.backends.utils import (CommonAttentionState,
+                                           CommonMetadataBuilder)
+from vllm.attention.ops.paged_attn import PagedAttentionMetadata
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
+
+from vllm_musa.platforms.musa import on_ph1
+
+logger = init_logger(__name__)
+
+class FlashAttentionBackend(AttentionBackend):
+
+    @staticmethod
+    def get_name() -> str:
+        return "MUSA_FLASH"
+
+    @staticmethod
+    def get_impl_cls() -> Type["FlashAttentionImpl"]:
+        return FlashAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> Type["AttentionMetadata"]:
+        return FlashAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> Type["FlashAttentionMetadataBuilder"]:
+        return FlashAttentionMetadataBuilder
+
+    @staticmethod
+    def get_state_cls() -> Type["CommonAttentionState"]:
+        return CommonAttentionState
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        if block_size % 16 != 0:
+            raise ValueError("Block size must be a multiple of 16.")
+        return (2, num_blocks, block_size, num_kv_heads, head_size)
+
+    @staticmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: torch.Tensor,
+    ) -> None:
+        src_key_cache = src_kv_cache[0]
+        dst_key_cache = dst_kv_cache[0]
+        ops.swap_blocks(src_key_cache, dst_key_cache, src_to_dst)
+        src_value_cache = src_kv_cache[1]
+        dst_value_cache = dst_kv_cache[1]
+        ops.swap_blocks(src_value_cache, dst_value_cache, src_to_dst)
+
+    @staticmethod
+    def copy_blocks(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: torch.Tensor,
+    ) -> None:
+        key_caches = [kv_cache[0] for kv_cache in kv_caches]
+        value_caches = [kv_cache[1] for kv_cache in kv_caches]
+
+        ops.copy_blocks(key_caches, value_caches, src_to_dists)
+
+@dataclass
+class FlashAttentionMetadata(AttentionMetadata, PagedAttentionMetadata):
+    """Metadata for FlashAttentionBackend.
+
+    NOTE: Any python object stored here is not updated when it is
+    cuda-graph replayed. If you have values that need to be changed
+    dynamically, it should be stored in tensor. The tensor has to be
+    updated from `CUDAGraphRunner.forward` API.
+    """
+    # (batch_size,). The sequence length per sequence. Sequence length means
+    # the computed tokens + new tokens None if it is a decoding.
+    seq_lens: Optional[List[int]]
+    # seq_lens stored as a tensor.
+    seq_lens_tensor: Optional[torch.Tensor]
+
+    # NOTE(sang): Definition of context_len, query_len, and seq_len.
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seq_len ---------------------|
+    #                                   |-- query_len ---|
+
+    # Maximum sequence length among prefill batch. 0 if there are decoding
+    # requests only.
+    max_prefill_seq_len: int
+    # Maximum sequence length among decode batch. 0 if there are prefill
+    # requests only.
+    max_decode_seq_len: int
+
+    # Whether or not if cuda graph is enabled.
+    # Cuda-graph is currently enabled for decoding only.
+    # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
+    use_cuda_graph: bool
+
+    # Maximum query length in the batch. None for decoding.
+    max_query_len: Optional[int] = None
+    # (batch_size + 1,). The cumulative subquery lengths of the sequences in
+    # the batch, used to index into subquery. E.g., if the subquery length
+    # is [4, 6], it is [0, 4, 10].
+    query_start_loc: Optional[torch.Tensor] = None
+    # (batch_size + 1,). The cumulative sequence lengths of the sequences in
+    # the batch, used to index into sequence. E.g., if the sequence length is
+    # [4, 6], it is [0, 4, 10].
+    seq_start_loc: Optional[torch.Tensor] = None
+    # (batch_size,) A tensor of context lengths (tokens that are computed
+    # so far).
+    context_lens_tensor: Optional[torch.Tensor] = None
+
+    # Max number of query tokens among request in the batch.
+    max_decode_query_len: Optional[int] = None
+
+    _cached_prefill_metadata: Optional["FlashAttentionMetadata"] = None
+    _cached_decode_metadata: Optional["FlashAttentionMetadata"] = None
+
+    # Begin encoder attn & enc/dec cross-attn fields...
+
+    # Encoder sequence lengths representation
+    encoder_seq_lens: Optional[List[int]] = None
+    encoder_seq_lens_tensor: Optional[torch.Tensor] = None
+
+    # Maximum sequence length among encoder sequences
+    max_encoder_seq_len: Optional[int] = None
+
+    # Number of tokens input to encoder
+    num_encoder_tokens: Optional[int] = None
+
+    # Cross-attention memory-mapping data structures: slot mapping
+    # and block tables
+    cross_slot_mapping: Optional[torch.Tensor] = None
+    cross_block_tables: Optional[torch.Tensor] = None
+
+    @property
+    def prefill_metadata(self) -> Optional["FlashAttentionMetadata"]:
+        if self.num_prefills == 0:
+            return None
+
+        if self._cached_prefill_metadata is not None:
+            return self._cached_prefill_metadata
+
+        assert self.seq_lens is not None
+        assert self.seq_lens_tensor is not None
+        assert self.block_tables is not None
+
+        self._cached_prefill_metadata = FlashAttentionMetadata(
+            num_prefills=self.num_prefills,
+            num_prefill_tokens=self.num_prefill_tokens,
+            num_decode_tokens=0,
+            slot_mapping=self.slot_mapping[:self.num_prefill_tokens],
+            multi_modal_placeholder_index_maps=self.
+            multi_modal_placeholder_index_maps,
+            enable_kv_scales_calculation=self.enable_kv_scales_calculation,
+            seq_lens=self.seq_lens[:self.num_prefills],
+            seq_lens_tensor=self.seq_lens_tensor[:self.num_prefills],
+            max_query_len=self.max_query_len,
+            max_prefill_seq_len=self.max_prefill_seq_len,
+            max_decode_seq_len=0,
+            query_start_loc=None if self.query_start_loc is None else
+            self.query_start_loc[:self.num_prefills + 1],
+            seq_start_loc=None if self.seq_start_loc is None else
+            self.seq_start_loc[:self.num_prefills + 1],
+            context_lens_tensor=None if self.context_lens_tensor is None else
+            self.context_lens_tensor[:self.num_prefills],
+            block_tables=self.block_tables[:self.num_prefills],
+            use_cuda_graph=False,
+            # Begin encoder & cross attn fields below...
+            encoder_seq_lens=self.encoder_seq_lens,
+            encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,
+            max_encoder_seq_len=self.max_encoder_seq_len,
+            cross_slot_mapping=self.cross_slot_mapping,
+            cross_block_tables=self.cross_block_tables)
+        return self._cached_prefill_metadata
+
+    @property
+    def decode_metadata(self) -> Optional["FlashAttentionMetadata"]:
+        if self.num_decode_tokens == 0:
+            return None
+
+        if self._cached_decode_metadata is not None:
+            return self._cached_decode_metadata
+        assert self.block_tables is not None
+        assert self.seq_lens_tensor is not None
+
+        self._cached_decode_metadata = FlashAttentionMetadata(
+            num_prefills=0,
+            num_prefill_tokens=0,
+            num_decode_tokens=self.num_decode_tokens,
+            slot_mapping=self.slot_mapping[self.num_prefill_tokens:],
+            multi_modal_placeholder_index_maps=None,
+            enable_kv_scales_calculation=True,
+            seq_lens=None,
+            seq_lens_tensor=self.seq_lens_tensor[self.num_prefills:],
+            max_query_len=None,
+            max_prefill_seq_len=0,
+            max_decode_seq_len=self.max_decode_seq_len,
+            query_start_loc=None,
+            seq_start_loc=None,
+            context_lens_tensor=None,
+            block_tables=self.block_tables[self.num_prefills:],
+            use_cuda_graph=self.use_cuda_graph,
+            # Begin encoder & cross attn fields below...
+            encoder_seq_lens=self.encoder_seq_lens,
+            encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,
+            max_encoder_seq_len=self.max_encoder_seq_len,
+            cross_slot_mapping=self.cross_slot_mapping,
+            cross_block_tables=self.cross_block_tables)
+        # Batch may be composed of prefill|decodes, adjust query start indices
+        # to refer to the start of decodes when the two are split apart.
+        # E.g. in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].
+        if self._cached_decode_metadata.query_start_loc is not None:
+            qs = self._cached_decode_metadata.query_start_loc
+            self._cached_decode_metadata.query_start_loc = qs - qs[0]
+        return self._cached_decode_metadata
+
+    def advance_step(self,
+                     model_input: "ModelInputForGPUWithSamplingMetadata",
+                     sampled_token_ids: Optional[torch.Tensor],
+                     block_size: int,
+                     num_seqs: int,
+                     num_queries: int,
+                     turn_prefills_into_decodes: bool = False):
+        """
+        Update metadata in-place to advance one decode step.
+        """
+
+        assert not turn_prefills_into_decodes, \
+            ("Chunked prefill is not supported with rocm_flash_attn yet."
+             "turn_prefills_into_decodes is a Multi-Step + Chunked-Prefill "
+             "specific parameter.")
+
+        # When using cudagraph, the num_seqs is padded to the next captured
+        # batch sized, but num_queries tracks the actual number of requests in
+        # the batch. For --enforce-eager mode, num_seqs == num_queries
+        if num_seqs != num_queries:
+            assert num_seqs > num_queries
+            assert self.use_cuda_graph
+
+        assert self.num_prefills == 0
+        assert self.num_prefill_tokens == 0
+        assert self.num_decode_tokens == num_seqs
+        assert self.slot_mapping.shape == (num_seqs, )
+
+        assert self.seq_lens is not None
+        assert len(self.seq_lens) == num_seqs
+        assert self.seq_lens_tensor is not None
+        assert self.seq_lens_tensor.shape == (num_seqs, )
+        assert self.max_query_len == 1
+        assert self.max_prefill_seq_len == 0
+        assert self.max_decode_seq_len == max(self.seq_lens)
+
+        assert self.query_start_loc is not None
+        assert self.query_start_loc.shape == (num_queries + 1, )
+        assert self.seq_start_loc is not None
+        assert self.seq_start_loc.shape == (num_seqs + 1, )
+
+        assert self.context_lens_tensor is not None
+        assert self.context_lens_tensor.shape == (num_queries, )
+
+        assert self.block_tables is not None
+        assert self.block_tables.shape[0] == num_seqs
+
+        # Update query lengths. Note that we update only queries and not seqs,
+        # since tensors may be padded due to captured cuda graph batch size
+        for i in range(num_queries):
+            self.seq_lens[i] += 1
+        self.max_decode_seq_len = max(self.seq_lens)
+
+        ops.advance_step_flashattn(num_seqs=num_seqs,
+                                   num_queries=num_queries,
+                                   block_size=block_size,
+                                   input_tokens=model_input.input_tokens,
+                                   sampled_token_ids=sampled_token_ids,
+                                   input_positions=model_input.input_positions,
+                                   seq_lens=self.seq_lens_tensor,
+                                   slot_mapping=self.slot_mapping,
+                                   block_tables=self.block_tables)
+
+
+class FlashAttentionMetadataBuilder(
+        CommonMetadataBuilder[FlashAttentionMetadata]):
+
+    _metadata_cls = FlashAttentionMetadata
+
+
+def _make_alibi_bias(alibi_slopes: torch.Tensor,
+                     dtype: torch.dtype,
+                     seq_lens: Optional[List[int]],
+                     make_attn_mask: bool = True) -> List[torch.Tensor]:
+    attn_biases = []
+    if seq_lens:
+        for seq_len in seq_lens:
+            bias = torch.arange(seq_len, dtype=dtype)
+            # NOTE(zhuohan): HF uses
+            #     `bias = bias[None, :].repeat(seq_len, 1)`
+            # here. We find that both biases give the same results, but
+            # the bias below more accurately follows the original ALiBi
+            # paper.
+            bias = bias[None, :] - bias[:, None]
+
+            num_heads = alibi_slopes.shape[0]
+            bias = bias[None, :].repeat(
+                (num_heads, 1, 1)).to(alibi_slopes.device)
+            bias.mul_(alibi_slopes[:, None, None])
+            if make_attn_mask:
+                inf_mask = torch.empty(
+                    (1, seq_len, seq_len),
+                    dtype=bias.dtype).fill_(-torch.inf).triu_(diagonal=1).to(
+                        alibi_slopes.device)
+                attn_biases.append((bias + inf_mask).to(dtype))
+            else:
+                attn_biases.append(bias.to(dtype))
+
+    return attn_biases
+
+
+def _get_seq_len_block_table_args(
+    attn_metadata: FlashAttentionMetadata,
+    attn_type: str,
+) -> tuple:
+    '''
+    The particular choice of sequence-length
+    attributes which should be extracted from attn_metadata is dependent
+    on the type of attention operation.
+
+    Decoder attn -> select entirely decoder self-attention-related fields
+    Encoder/decoder cross-attn -> select encoder sequence lengths
+    Encoder attn -> select encoder sequence lengths fields
+
+    Arguments:
+
+    * attn_metadata: Attention metadata structure associated with attention op
+    * attn_type: encoder attention, decoder self-attention,
+                encoder/decoder cross-attention
+
+    Returns:
+
+    * Appropriate sequence-lengths tensors for query and key
+    * Appropriate max sequence-length scalar
+    '''
+
+    partial_prefix_sum = 0
+    if attn_type == AttentionType.ENCODER:
+        assert attn_metadata.encoder_seq_lens is not None
+        assert attn_metadata.encoder_seq_lens_tensor is not None
+        query_seq_start_loc = torch.tensor(
+            [0] + [
+                partial_prefix_sum := partial_prefix_sum + i
+                for i in attn_metadata.encoder_seq_lens
+            ],
+            device=attn_metadata.encoder_seq_lens_tensor.device,
+            dtype=attn_metadata.encoder_seq_lens_tensor.dtype)
+        causal_mask = False
+
+        # No block tables associated with encoder attention
+        return (query_seq_start_loc, attn_metadata.max_encoder_seq_len,
+                query_seq_start_loc, attn_metadata.max_encoder_seq_len,
+                attn_metadata.encoder_seq_lens, causal_mask)
+    elif attn_type == AttentionType.DECODER:
+        # Decoder self-attention
+        # Choose max_seq_len based on whether we are in prompt_run
+        assert attn_metadata.seq_lens is not None
+        assert attn_metadata.seq_lens_tensor is not None
+        query_seq_start_loc = torch.tensor(
+            [0] + [
+                partial_prefix_sum := partial_prefix_sum + i
+                for i in attn_metadata.seq_lens
+            ],
+            device=attn_metadata.seq_lens_tensor.device,
+            dtype=attn_metadata.seq_lens_tensor.dtype)
+        max_seq_len = attn_metadata.max_prefill_seq_len
+        causal_mask = True
+
+        return (query_seq_start_loc, max_seq_len, query_seq_start_loc,
+                max_seq_len, attn_metadata.seq_lens, causal_mask)
+    elif attn_type == AttentionType.ENCODER_DECODER:
+        assert attn_metadata.seq_lens is not None
+        assert attn_metadata.encoder_seq_lens_tensor is not None
+        query_start_loc = torch.tensor(
+            [0] + [
+                partial_prefix_sum := partial_prefix_sum + i
+                for i in attn_metadata.seq_lens
+            ],
+            device=attn_metadata.encoder_seq_lens_tensor.device,
+            dtype=attn_metadata.encoder_seq_lens_tensor.dtype)
+
+        partial_prefix_sum = 0
+        assert attn_metadata.encoder_seq_lens is not None
+        assert attn_metadata.seq_lens_tensor is not None
+        key_seq_start_loc = torch.tensor(
+            [0] + [
+                partial_prefix_sum := partial_prefix_sum + i
+                for i in attn_metadata.encoder_seq_lens
+            ],
+            device=attn_metadata.seq_lens_tensor.device,
+            dtype=attn_metadata.seq_lens_tensor.dtype)
+        causal_mask = False
+
+        # Enc/dec cross-attention KVs match encoder sequence length;
+        # cross-attention utilizes special "cross" block tables
+        return (query_start_loc, attn_metadata.max_prefill_seq_len,
+                key_seq_start_loc, attn_metadata.max_encoder_seq_len,
+                attn_metadata.seq_lens, causal_mask)
+    else:
+        raise AttributeError(f"Invalid attention type {str(attn_type)}")
+
+
+class FlashAttentionImpl(AttentionImpl):
+    """
+    If the input tensors contain prompt tokens, the layout is as follows:
+    |<--------------- num_prefill_tokens ----------------->|
+    |<--prefill_0-->|<--prefill_1-->|...|<--prefill_N-1--->|
+
+    Otherwise, the layout is as follows:
+    |<----------------- num_decode_tokens ------------------>|
+    |<--decode_0-->|..........|<--decode_M-1-->|<--padding-->|
+
+    Generation tokens can contain padding when cuda-graph is used.
+    Currently, prompt tokens don't contain any padding.
+
+    The prompts might have different lengths, while the generation tokens
+    always have length 1.
+
+    If chunked prefill is enabled, prefill tokens and decode tokens can be
+    batched together in a flattened 1D query.
+
+    |<----- num_prefill_tokens ---->|<------- num_decode_tokens --------->|
+    |<-prefill_0->|...|<-prefill_N-1->|<--decode_0-->|...|<--decode_M-1-->|
+
+    Currently, cuda graph is disabled for chunked prefill, meaning there's no
+    padding between prefill and decode tokens.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: Optional[List[float]],
+        sliding_window: Optional[int],
+        kv_cache_dtype: str,
+        blocksparse_params: Optional[Dict[str, Any]] = None,
+        logits_soft_cap: Optional[float] = None,
+        attn_type: str = AttentionType.DECODER,
+        kv_sharing_target_layer_name: Optional[int] = None,
+        **extra_impl_args,
+    ) -> None:
+        if blocksparse_params is not None:
+            raise ValueError(
+                "FlashAttention does not support block-sparse attention.")
+
+        if logits_soft_cap is None:
+            # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
+            self.logits_soft_cap = 0.0
+        else:
+            self.logits_soft_cap = logits_soft_cap
+        self.attn_type = attn_type
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        if alibi_slopes is not None:
+            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
+        self.alibi_slopes = alibi_slopes
+        self.sliding_window = ((sliding_window, sliding_window)
+                               if sliding_window is not None else (-1, -1))
+        self.kv_cache_dtype = kv_cache_dtype
+
+        assert self.num_heads % self.num_kv_heads == 0
+        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+
+        self.num_splits = 4
+        if on_ph1():
+            self.num_splits = 1
+
+    def repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor:
+        """torch.repeat_interleave(x, dim=1, repeats=n_rep)"""
+        tokens, n_kv_heads, head_dim = x.shape
+        return (x[:, :,
+                  None, :].expand(tokens, n_kv_heads, n_rep,
+                                  head_dim).reshape(tokens, n_kv_heads * n_rep,
+                                                    head_dim))
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: FlashAttentionMetadata,
+        output: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        # print("---------------FlashAttention---------------")
+        """Forward pass with FlashAttention.
+
+        Args:
+            query: shape = [num_tokens, num_heads, head_size]
+            key: shape = [num_tokens, num_kv_heads, head_size]
+            value: shape = [num_tokens, num_kv_heads, head_size]
+            output: shape = [num_tokens, num_heads, head_size]
+            kv_cache = [2, num_blocks, block_size, num_kv_heads, head_size]
+                NOTE: kv_cache will be an empty tensor with shape [0]
+                for profiling run.
+            attn_metadata: Metadata for attention.
+        NOTE: It in-place updates the output tensor.
+        """
+        # NOTE(woosuk): FlashAttention does not support FP8 KV cache.
+        assert layer._k_scale_float == 1.0 and layer._v_scale_float == 1.0, (
+            "key/v_scale is not supported in FlashAttention.")
+
+        attn_type = self.attn_type
+        if (attn_type == AttentionType.ENCODER
+                and (not attn_metadata.is_all_encoder_attn_metadata_set)):
+            raise AttributeError("Encoder attention requires setting "
+                                 "encoder metadata attributes.")
+        elif (attn_type == AttentionType.ENCODER_DECODER
+              and (not attn_metadata.is_all_cross_attn_metadata_set)):
+            raise AttributeError("Encoder/decoder cross-attention "
+                                 "requires setting cross-attention "
+                                 "metadata attributes.")
+
+        query = query.view(-1, self.num_heads, self.head_size)
+        if key is not None:
+            assert value is not None
+            key = key.view(-1, self.num_kv_heads, self.head_size)
+            value = value.reshape(-1, self.num_kv_heads, self.head_size)
+        else:
+            assert value is None
+
+        if kv_cache.numel() > 0:
+            key_cache = kv_cache[0]
+            value_cache = kv_cache[1]
+            # We skip updating the KV cache under two conditions:
+            #  a. When the Attention Type is ENCODER. In this phase, we compute
+            #     only the encoder attention without updating the cache.
+            #  b. When both Key and Value are None. This occurs during
+            #     cross-attention computation in the decoding phase, where the
+            #     KV cache is already populated with the cross-attention
+            #     tensor. Thus, we skip cache updates during this time.
+            if (attn_type != AttentionType.ENCODER) and (key is not None) and (
+                    value is not None):
+                if attn_type == AttentionType.ENCODER_DECODER:
+                    # Update cross-attention KV cache (prefill-only)
+                    updated_slot_mapping = attn_metadata.cross_slot_mapping
+                else:
+                    # Update self-attention KV cache (prefill/decode)
+                    updated_slot_mapping = attn_metadata.slot_mapping
+
+                # Reshape the input keys and values and store them in the cache.
+                # If kv_cache is not provided, the new key and value tensors are
+                # not cached. This happens during the initial memory
+                # profiling run.
+                torch.ops._C_cache_ops.reshape_and_cache_flash(
+                    key,
+                    value,
+                    kv_cache[0],
+                    kv_cache[1],
+                    updated_slot_mapping.flatten(),  # type: ignore[union-attr]
+                    self.kv_cache_dtype,
+                    layer._k_scale,
+                    layer._v_scale,
+                )
+
+        if self.attn_type != AttentionType.ENCODER:
+            num_prefill_tokens = attn_metadata.num_prefill_tokens
+        else:
+            assert attn_metadata.num_encoder_tokens is not None
+            num_prefill_tokens = attn_metadata.num_encoder_tokens
+
+        output = torch.empty_like(query)
+        # Query for decode. KV is not needed because it is already cached.
+        decode_query = query[num_prefill_tokens:]
+        # QKV for prefill.
+        query = query[:num_prefill_tokens]
+
+        if key is not None and value is not None \
+            and self.attn_type != AttentionType.ENCODER_DECODER:
+            key = key[:num_prefill_tokens]
+            value = value[:num_prefill_tokens]
+
+        if prefill_meta := attn_metadata.prefill_metadata:
+            # Prompt run.
+            # normal attention and DECODER
+            if self.attn_type == AttentionType.DECODER and (
+                    kv_cache.numel() == 0 or prefill_meta.block_tables is None
+                    or prefill_meta.block_tables.numel() == 0):
+                (query_seq_start_loc, query_max_seq_len, key_seq_start_loc,
+                 key_max_seq_len, seq_lens,
+                 causal_mask) = (prefill_meta.seq_start_loc,
+                                 prefill_meta.max_prefill_seq_len,
+                                 prefill_meta.seq_start_loc,
+                                 prefill_meta.max_prefill_seq_len,
+                                 attn_metadata.seq_lens, True)
+            # prefix-enabled attention and ENCODER/ENCODER_DECODER
+            else:
+                (query_seq_start_loc, query_max_seq_len, key_seq_start_loc,
+                 key_max_seq_len, seq_lens,
+                 causal_mask) = _get_seq_len_block_table_args(
+                     prefill_meta, self.attn_type)
+            # Prompt run.
+            if prefill_meta := attn_metadata.prefill_metadata:
+                # if self.num_kv_heads != self.num_heads:
+                #     # Interleave for MQA workaround.
+                #     key = self.repeat_kv(key, self.num_queries_per_kv)
+                #     value = self.repeat_kv(value, self.num_queries_per_kv)
+
+                # output = sdpa_attention_with_torch_seqlen_pad(
+                output = sdpa_attention_with_kernel_seqlen_pad(
+                    query,
+                    key,
+                    value,
+                    query_seq_start_loc,
+                    query_max_seq_len,
+                )
+
+                return output
+
+        if decode_meta := attn_metadata.decode_metadata:
+            from vllm.attention.ops.triton_decode_attention import _decode_grouped_att_m_fwd, _decode_softmax_reducev_fwd
+
+            decode_meta = attn_metadata.decode_metadata
+            assert decode_meta is not None
+
+            B = decode_query.shape[0]
+            PAGE_SIZE = key_cache.size(1)
+
+            if self.num_splits == 1:
+                attn_logits = torch.zeros(
+                    (
+                        B,
+                        self.num_heads,
+                        self.num_splits,
+                        self.head_size,
+                    ),
+                    dtype=query.dtype,
+                    device=query.device,
+                )
+                output = attn_logits
+            else:
+                output = torch.empty(B,
+                                self.num_heads,
+                                self.head_size,
+                                dtype=query.dtype,
+                                device=query.device)
+                attn_logits = torch.zeros(
+                    (
+                        B,
+                        self.num_heads,
+                        self.num_splits,
+                        self.head_size + 1,
+                    ),
+                    dtype=torch.float32,
+                    device=query.device,
+                )
+
+            if on_ph1():
+                ops.mp31_decode_mha(decode_query,
+                                    key_cache,
+                                    value_cache,
+                                    self.scale,
+                                    self.num_splits,
+                                    decode_meta.block_tables,
+                                    decode_meta.seq_lens_tensor,
+                                    attn_logits)
+            else:
+                # Run MQA
+                _decode_grouped_att_m_fwd(
+                    decode_query,
+                    key_cache,
+                    value_cache,
+                    attn_logits,
+                    decode_meta.block_tables,
+                    decode_meta.seq_lens_tensor,
+                    self.num_splits,
+                    self.scale,
+                    PAGE_SIZE,
+                    0,
+                )
+
+
+            if self.num_splits > 1:
+                _decode_softmax_reducev_fwd(attn_logits,
+                                            decode_query,
+                                            output,
+                                            value_cache,
+                                            decode_meta.seq_lens_tensor,
+                                            self.num_splits)
+
+            return output.view(-1, self.num_heads*self.head_size)
+
+def sdpa_attention_with_kernel_seqlen_pad(
+    query: torch.Tensor, #sum_seq, h_q, d_q
+    key: torch.Tensor,
+    value: torch.Tensor,
+    seq_lens: List[int],
+    max_prefill_seq_len: int,
+    is_causal: bool = True,
+) -> torch.Tensor:
+    bs = seq_lens.shape[0] - 1
+    sum_seq, h_q, d_q = query.shape
+    _,h_kv, d_kv = key.shape
+    device, dtype = query.device, query.dtype
+
+    q_pad = torch.empty((bs, h_q, max_prefill_seq_len, d_q), device=device, dtype=dtype)
+    k_pad = torch.empty((bs, h_kv, max_prefill_seq_len, d_q), device=device, dtype=dtype)
+    v_pad = torch.empty((bs, h_kv, max_prefill_seq_len, d_q), device=device, dtype=dtype)
+    output = torch.empty((sum_seq, h_q, d_q), device=device, dtype=dtype)
+
+    ops.varlen_fa_seqlen_pad(query, key, value, q_pad, k_pad, v_pad, seq_lens, seq_lens, sum_seq, max_prefill_seq_len, bs)
+    attn_out, _, _ = torch.ops.aten._scaled_dot_product_attention_flash_musa(
+        q_pad,
+        k_pad,
+        v_pad,
+        dropout_p=0.0,
+        is_causal=is_causal)
+
+    ops.varlen_fa_seqlen_unpad(attn_out, output, seq_lens, sum_seq, max_prefill_seq_len, d_q, h_q, bs)
+    return output.view(-1, h_q * d_q)
+
+
+def pad_seq(x: torch.Tensor, cu_seqlen:torch.Tensor, max_prefill_seq_len: int, num_heads: int) -> torch.Tensor:
+    bs = cu_seqlen.shape[0] - 1
+    padded = torch.zeros(bs, num_heads, max_prefill_seq_len, x.shape[-1], device=x.device, dtype=x.dtype)
+    for i in range(bs):
+        start = cu_seqlen[i].item()
+        end = cu_seqlen[i + 1].item()
+        seq_len = end - start
+        padded[i, :, :seq_len] = x[start:end, :, :].transpose(0, 1)
+    return padded
+
+def restore_tokens(attn_output: torch.Tensor, cu_seqlen: torch.Tensor, max_prefill_seq_len: int, num_heads: int) -> torch.Tensor:
+    bs = cu_seqlen.shape[0] - 1
+    attn_output_tensor = attn_output.transpose(1, 2).view(bs, max_prefill_seq_len, num_heads, -1)
+    restored_tokens = []
+    for i in range(bs):
+        start = cu_seqlen[i].item()
+        end = cu_seqlen[i + 1].item()
+        seq_len = end - start
+        restored_tokens.append(attn_output_tensor[i, :seq_len])
+    return torch.cat(restored_tokens, dim=0)
+
+
+def sdpa_attention_with_torch_seqlen_pad(
+    query: torch.Tensor, #sum_seq, h_q, d_q
+    key: torch.Tensor,
+    value: torch.Tensor,
+    seq_lens: List[int],
+    max_prefill_seq_len: int,
+    is_causal: bool = True,
+) -> torch.Tensor:
+    sum_seq, h_q, d_q = query.shape
+    _, h_kv, d_kv = key.shape
+
+    q_pad = pad_seq(query, seq_lens, max_prefill_seq_len, h_q)
+    k_pad = pad_seq(key, seq_lens, max_prefill_seq_len, h_kv)
+    v_pad = pad_seq(value, seq_lens, max_prefill_seq_len, h_kv)
+
+    output, _, _ = torch.ops.aten._scaled_dot_product_attention_flash_musa(
+        q_pad,
+        k_pad,
+        v_pad,
+        dropout_p=0.0,
+        is_causal=is_causal)
+
+    output = restore_tokens(output, seq_lens, max_prefill_seq_len, h_q)
+    return output.view(-1, h_q * d_q)
diff --git a/vllm_musa/vllm_musa/v0/musa_worker.py b/vllm_musa/vllm_musa/v0/musa_worker.py
new file mode 100644
index 000000000..abb863c3d
--- /dev/null
+++ b/vllm_musa/vllm_musa/v0/musa_worker.py
@@ -0,0 +1,290 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+# Adapted from vllm-project/vllm/vllm/worker/worker.py
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import gc
+from typing import Dict, List, Optional, Tuple, Type
+
+import os
+import time
+from datetime import datetime
+from distutils.version import StrictVersion
+
+import torch
+import torch.distributed
+from vllm import envs
+from vllm.config import VllmConfig
+from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized
+from vllm.logger import init_logger
+from vllm.model_executor import set_random_seed
+from vllm.sequence import SequenceGroupMetadata
+from vllm.utils import (GiB_bytes, MemorySnapshot, memory_profiling)
+from vllm.worker.cache_engine import CacheEngine
+from vllm.worker.enc_dec_model_runner import EncoderDecoderModelRunner
+from vllm.worker.model_runner_base import ModelRunnerBase
+from vllm.worker.model_runner import GPUModelRunnerBase, ModelRunner
+
+from vllm.worker.pooling_model_runner import PoolingModelRunner
+from vllm.worker.worker_base import WorkerBase
+from vllm.worker.worker import Worker
+from vllm.distributed import (ensure_model_parallel_initialized,
+                              init_distributed_environment,
+                              set_custom_all_reduce)
+
+logger = init_logger(__name__)
+
+class MTGPUWorker(Worker):
+    """A worker class that executes (a partition of) the model on a MTGPU.
+    Each worker is associated with a single MTGPU. The worker is responsible for
+    maintaining the KV cache and executing the model on the MTGPU. In case of
+    distributed inference, each worker is assigned a partition of the model.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        local_rank: int,
+        rank: int,
+        distributed_init_method: str,
+        is_driver_worker: bool = False,
+        model_runner_cls: Optional[Type[ModelRunnerBase]] = None,
+    ) -> None:
+        # Register ops and patch when worker init.
+        from vllm_musa import ops  # noqa: F401
+        from vllm_musa import patch  # noqa: F401
+
+        WorkerBase.__init__(self, vllm_config)
+        self.parallel_config.rank = rank
+        self.local_rank = local_rank
+        self.rank = rank
+        self.distributed_init_method = distributed_init_method
+        self.is_driver_worker = is_driver_worker
+        if self.model_config.trust_remote_code:
+            # note: lazy import to avoid importing torch before initializing
+            from vllm.utils import init_cached_hf_modules
+            init_cached_hf_modules()
+
+        # Return hidden states from target model if the draft model is an
+        # mlp_speculator
+        speculative_config = self.speculative_config
+        model_config = self.model_config
+        speculative_args = {} if speculative_config is None \
+            or (speculative_config.draft_model_config.model ==
+                model_config.model) \
+            or (speculative_config.draft_model_config.hf_config.model_type
+                not in ["medusa", "mlp_speculator", "eagle"]) \
+                    else {"return_hidden_states": True}
+
+        ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner
+        if model_config.runner_type == "pooling":
+            ModelRunnerClass = PoolingModelRunner
+        elif self.model_config.is_encoder_decoder:
+            ModelRunnerClass = EncoderDecoderModelRunner
+        self.model_runner: GPUModelRunnerBase = ModelRunnerClass(
+            vllm_config=self.vllm_config,
+            kv_cache_dtype=self.cache_config.cache_dtype,
+            is_driver_worker=is_driver_worker,
+            **speculative_args,
+        )
+        if model_runner_cls is not None:
+            self.model_runner = model_runner_cls(self.model_runner)
+
+        # Uninitialized cache engine. Will be initialized by
+        # initialize_cache.
+        self.cache_engine: List[CacheEngine]
+        # Initialize gpu_cache as pooling models don't initialize kv_caches
+        self.gpu_cache: Optional[List[List[torch.Tensor]]] = None
+        self._seq_group_metadata_cache: Dict[str, SequenceGroupMetadata] = {}
+
+        # Torch profiler. Enabled and configured through env vars:
+        # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
+        if envs.VLLM_TORCH_PROFILER_DIR:
+            self.torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
+            logger.info(
+                "Profiling enabled. Traces will be saved to: %s",
+                self.torch_profiler_trace_dir,
+            )
+
+            def trace_handler(prof):
+                begin = time.monotonic()
+                current_time = datetime.now().strftime("%Y.%m.%d-%H:%M:%S")
+                rank = torch.distributed.get_rank()
+                save_dir = os.path.join(self.torch_profiler_trace_dir, current_time)
+                if not os.path.exists(save_dir):
+                    os.makedirs(save_dir, exist_ok=True)
+                prof.export_chrome_trace(f"{save_dir}/profile_rank{rank}.json")
+                print(
+                    f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
+                )
+
+            if StrictVersion(torch.__version__) < StrictVersion('2.5.0'):
+                activities=[
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.MUSA,
+                    ]
+            else:
+                activities=[
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.PrivateUse1,
+                    ]
+
+            self.profiler = torch.profiler.profile(
+                activities=activities,
+                profile_memory=True,
+                record_shapes=True,
+                with_flops=True,
+                with_modules=True,
+                with_stack=True,
+                on_trace_ready=trace_handler,
+            )
+        else:
+            self.profiler = None
+
+
+    def init_device(self) -> None:
+        if self.device_config.device.type == "musa":
+            # os.environ["MCCL_PROTOS"] = "2"
+            self.device = torch.device(f"musa:{self.local_rank}")
+            torch.musa.set_device(self.device)
+
+            _check_if_gpu_supports_dtype(self.model_config.dtype)
+            gc.collect()
+            torch.musa.empty_cache()
+        else:
+            raise RuntimeError(
+                f"Not support device type: {self.device_config.device}")
+
+        self.baseline_snapshot = MemorySnapshot()
+        # Initialize the distributed environment.
+        init_worker_distributed_environment(self.vllm_config, self.rank,
+                                            self.distributed_init_method,
+                                            self.local_rank)
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+
+    @torch.inference_mode()
+    def determine_num_available_blocks(self) -> Tuple[int, int]:
+        """Profiles the peak memory usage of the model to determine how many
+        KV blocks may be allocated without OOMs.
+
+        The engine will first conduct a profiling of the existing memory usage.
+        Then, it calculate the maximum possible number of GPU and CPU blocks
+        that can be allocated with the remaining free memory.
+
+        .. tip::
+            You may limit the usage of GPU memory
+            by adjusting the `gpu_memory_utilization` parameter.
+        """
+        # Profile the memory usage of the model and get the maximum number of
+        # cache blocks that can be allocated with the remaining free memory.
+        torch.musa.empty_cache()
+        torch.musa.reset_peak_memory_stats()
+
+        free_memory_pre_profile, total_gpu_memory = torch.musa.mem_get_info()
+
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model.
+        with memory_profiling(
+                self.baseline_snapshot,
+                weights_memory=self.model_runner.model_memory_usage) as result:
+            self.model_runner.profile_run()
+
+        self._assert_memory_footprint_increased_during_profiling()
+
+        memory_for_current_instance = total_gpu_memory * \
+            self.cache_config.gpu_memory_utilization
+        available_kv_cache_memory = (memory_for_current_instance -
+                                     result.non_kv_cache_memory)
+
+        # Calculate the number of blocks that can be allocated with the
+        # profiled peak memory.
+        cache_block_size = self.get_cache_block_size_bytes()
+        if cache_block_size == 0:
+            num_gpu_blocks = 0
+            num_cpu_blocks = 0
+        else:
+            num_gpu_blocks = int(available_kv_cache_memory // cache_block_size)
+            num_cpu_blocks = int(self.cache_config.swap_space_bytes //
+                                 cache_block_size)
+        num_gpu_blocks = max(num_gpu_blocks, 0)
+        num_cpu_blocks = max(num_cpu_blocks, 0)
+
+        msg = (f"Memory profiling takes {result.profile_time:.2f} seconds\n"
+               "the current vLLM instance can use "
+               "total_gpu_memory "
+               f"({(total_gpu_memory / GiB_bytes):.2f}GiB)"
+               " x gpu_memory_utilization "
+               f"({self.cache_config.gpu_memory_utilization:.2f})"
+               f" = {(memory_for_current_instance / GiB_bytes):.2f}GiB\n"
+               "model weights take "
+               f"{(result.weights_memory / GiB_bytes):.2f}GiB;"
+               " non_torch_memory takes "
+               f"{(result.non_torch_increase / GiB_bytes):.2f}GiB;"
+               " PyTorch activation peak memory takes "
+               f"{(result.torch_peak_increase / GiB_bytes):.2f}GiB;"
+               " the rest of the memory reserved for KV Cache is "
+               f"{(available_kv_cache_memory / GiB_bytes):.2f}GiB.")
+
+        logger.info(msg)
+        # Final cleanup
+        gc.collect()
+
+        return num_gpu_blocks, num_cpu_blocks
+
+    def stop_profile(self):
+        if self.profiler is None:
+            raise RuntimeError("Profiler is not enabled.")
+        self.profiler.stop()
+        # print(
+        #     self.profiler.key_averages().table(sort_by="self_musa_time_total"))
+
+    def _assert_memory_footprint_increased_during_profiling(self):
+        # NOTE(woosuk): Here we assume that the other processes using the same
+        # GPU did not change their memory usage during the profiling.
+        free_gpu_memory, total = torch.musa.mem_get_info()
+        cuda_memory = total - free_gpu_memory
+        assert self.baseline_snapshot.cuda_memory < cuda_memory, (
+            "Error in memory profiling. "
+            f"Initial used memory {self.baseline_snapshot.cuda_memory}, "
+            f"currently used memory {cuda_memory}. "
+            f"This happens when the GPU memory was "
+            "not properly cleaned up before initializing the vLLM instance.")
+
+
+def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
+    # Check if the GPU supports the dtype.
+    if torch_dtype == torch.bfloat16:  # noqa: SIM102
+        return True
+
+def init_worker_distributed_environment(
+    vllm_config: VllmConfig,
+    rank: int,
+    distributed_init_method: Optional[str] = None,
+    local_rank: int = -1,
+) -> None:
+    """Initialize the distributed environment."""
+    parallel_config = vllm_config.parallel_config
+    set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
+
+    init_distributed_environment(parallel_config.world_size, rank,
+                                 distributed_init_method, local_rank, backend="mccl")
+    ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
+                                      parallel_config.pipeline_parallel_size)
+
+    ensure_kv_transfer_initialized(vllm_config)
diff --git a/vllm_musa/vllm_musa/v0/paged_attention.py b/vllm_musa/vllm_musa/v0/paged_attention.py
new file mode 100644
index 000000000..441813dfd
--- /dev/null
+++ b/vllm_musa/vllm_musa/v0/paged_attention.py
@@ -0,0 +1,811 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Attention layer MUSA GPUs."""
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
+                                              AttentionLayer,
+                                              AttentionMetadata, AttentionType)
+from vllm.attention.backends.utils import (CommonAttentionState,
+                                           CommonMetadataBuilder)
+from vllm.attention.ops.paged_attn import (PagedAttention,
+                                           PagedAttentionMetadata)
+from vllm.logger import init_logger
+
+if TYPE_CHECKING:
+    from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata
+
+logger = init_logger(__name__)
+
+class MUSAFlashAttentionBackend(AttentionBackend):
+
+    @staticmethod
+    def get_name() -> str:
+        return "MUSA_FLASH"
+
+    @staticmethod
+    def get_impl_cls() -> Type["MUSAFlashAttentionImpl"]:
+        return MUSAFlashAttentionImpl
+
+    @staticmethod
+    def get_metadata_cls() -> Type["AttentionMetadata"]:
+        return MUSAFlashAttentionMetadata
+
+    @staticmethod
+    def get_builder_cls() -> Type["MUSAFlashAttentionMetadataBuilder"]:
+        return MUSAFlashAttentionMetadataBuilder
+
+    @staticmethod
+    def get_state_cls() -> Type["CommonAttentionState"]:
+        return CommonAttentionState
+
+    @staticmethod
+    def get_kv_cache_shape(
+        num_blocks: int,
+        block_size: int,
+        num_kv_heads: int,
+        head_size: int,
+    ) -> Tuple[int, ...]:
+        return PagedAttention.get_kv_cache_shape(num_blocks, block_size,
+                                                 num_kv_heads, head_size)
+
+    @staticmethod
+    def swap_blocks(
+        src_kv_cache: torch.Tensor,
+        dst_kv_cache: torch.Tensor,
+        src_to_dst: torch.Tensor,
+    ) -> None:
+        PagedAttention.swap_blocks(src_kv_cache, dst_kv_cache, src_to_dst)
+
+    @staticmethod
+    def copy_blocks(
+        kv_caches: List[torch.Tensor],
+        src_to_dists: torch.Tensor,
+    ) -> None:
+        PagedAttention.copy_blocks(kv_caches, src_to_dists)
+
+
+@dataclass
+class MUSAFlashAttentionMetadata(AttentionMetadata, PagedAttentionMetadata):
+    """Metadata for FlashAttentionBackend.
+
+    NOTE: Any python object stored here is not updated when it is
+    cuda-graph replayed. If you have values that need to be changed
+    dynamically, it should be stored in tensor. The tensor has to be
+    updated from `CUDAGraphRunner.forward` API.
+    """
+    # (batch_size,). The sequence length per sequence. Sequence length means
+    # the computed tokens + new tokens None if it is a decoding.
+    seq_lens: Optional[List[int]]
+    # seq_lens stored as a tensor.
+    seq_lens_tensor: Optional[torch.Tensor]
+    # Maximum sequence length among prefill batch. 0 if there are decoding
+    # requests only.
+    max_prefill_seq_len: int
+    # Maximum sequence length among decode batch. 0 if there are prefill
+    # requests only.
+    max_decode_seq_len: int
+
+    # Whether or not if cuda graph is enabled.
+    # Cuda-graph is currently enabled for decoding only.
+    # TODO(woosuk): Move `use_cuda_graph` out since it's unrelated to attention.
+    use_cuda_graph: bool
+
+    # NOTE(sang): Definition of context_len, query_len, and seq_len.
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seq_len ----------------------|
+    #                                   |-- query_len ---|
+
+    # Maximum query length in the batch. None for decoding.
+    max_query_len: Optional[int] = None
+    # (batch_size + 1,). The cumulative subquery lengths of the sequences in
+    # the batch, used to index into subquery. E.g., if the subquery length
+    # is [4, 6], it is [0, 4, 10].
+    query_start_loc: Optional[torch.Tensor] = None
+    # (batch_size + 1,). The cumulative sequence lengths of the sequences in
+    # the batch, used to index into sequence. E.g., if the sequence length is
+    # [4, 6], it is [0, 4, 10].
+    seq_start_loc: Optional[torch.Tensor] = None
+    # (batch_size,) A tensor of context lengths (tokens that are computed
+    # so far).
+    context_lens_tensor: Optional[torch.Tensor] = None
+
+    # Max number of query tokens among request in the batch.
+    max_decode_query_len: Optional[int] = None
+
+    _cached_prefill_metadata: Optional["MUSAFlashAttentionMetadata"] = None
+    _cached_decode_metadata: Optional["MUSAFlashAttentionMetadata"] = None
+
+    # Begin encoder attn & enc/dec cross-attn fields...
+
+    # Encoder sequence lengths representation
+    encoder_seq_lens: Optional[List[int]] = None
+    encoder_seq_lens_tensor: Optional[torch.Tensor] = None
+
+    # Maximum sequence length among encoder sequences
+    max_encoder_seq_len: Optional[int] = None
+
+    # Number of tokens input to encoder
+    num_encoder_tokens: Optional[int] = None
+
+    # Cross-attention memory-mapping data structures: slot mapping
+    # and block tables
+    cross_slot_mapping: Optional[torch.Tensor] = None
+    cross_block_tables: Optional[torch.Tensor] = None
+
+    @property
+    def prefill_metadata(self) -> Optional["MUSAFlashAttentionMetadata"]:
+        if self.num_prefills == 0:
+            return None
+
+        if self._cached_prefill_metadata is not None:
+            return self._cached_prefill_metadata
+
+        assert self.seq_lens is not None
+        assert self.seq_lens_tensor is not None
+        assert self.block_tables is not None
+
+        self._cached_prefill_metadata = MUSAFlashAttentionMetadata(
+            num_prefills=self.num_prefills,
+            num_prefill_tokens=self.num_prefill_tokens,
+            num_decode_tokens=0,
+            slot_mapping=self.slot_mapping[:self.num_prefill_tokens],
+            multi_modal_placeholder_index_maps=self.
+            multi_modal_placeholder_index_maps,
+            enable_kv_scales_calculation=self.enable_kv_scales_calculation,
+            seq_lens=self.seq_lens[:self.num_prefills],
+            seq_lens_tensor=self.seq_lens_tensor[:self.num_prefills],
+            max_query_len=self.max_query_len,
+            max_prefill_seq_len=self.max_prefill_seq_len,
+            max_decode_seq_len=0,
+            query_start_loc=None if self.query_start_loc is None else
+            self.query_start_loc[:self.num_prefills + 1],
+            seq_start_loc=None if self.seq_start_loc is None else
+            self.seq_start_loc[:self.num_prefills + 1],
+            context_lens_tensor=None if self.context_lens_tensor is None else
+            self.context_lens_tensor[:self.num_prefills],
+            block_tables=self.block_tables[:self.num_prefills],
+            use_cuda_graph=False,
+            # Begin encoder & cross attn fields below...
+            encoder_seq_lens=self.encoder_seq_lens,
+            encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,
+            max_encoder_seq_len=self.max_encoder_seq_len,
+            cross_slot_mapping=self.cross_slot_mapping,
+            cross_block_tables=self.cross_block_tables)
+        return self._cached_prefill_metadata
+
+    @property
+    def decode_metadata(self) -> Optional["MUSAFlashAttentionMetadata"]:
+        if self.num_decode_tokens == 0:
+            return None
+
+        if self._cached_decode_metadata is not None:
+            return self._cached_decode_metadata
+        assert self.block_tables is not None
+        assert self.seq_lens_tensor is not None
+
+        self._cached_decode_metadata = MUSAFlashAttentionMetadata(
+            num_prefills=0,
+            num_prefill_tokens=0,
+            num_decode_tokens=self.num_decode_tokens,
+            slot_mapping=self.slot_mapping[self.num_prefill_tokens:],
+            multi_modal_placeholder_index_maps=None,
+            enable_kv_scales_calculation=True,
+            seq_lens=None,
+            seq_lens_tensor=self.seq_lens_tensor[self.num_prefills:],
+            max_query_len=None,
+            max_prefill_seq_len=0,
+            max_decode_seq_len=self.max_decode_seq_len,
+            query_start_loc=None,
+            seq_start_loc=None,
+            context_lens_tensor=None,
+            block_tables=self.block_tables[self.num_prefills:],
+            use_cuda_graph=self.use_cuda_graph,
+            # Begin encoder & cross attn fields below...
+            encoder_seq_lens=self.encoder_seq_lens,
+            encoder_seq_lens_tensor=self.encoder_seq_lens_tensor,
+            max_encoder_seq_len=self.max_encoder_seq_len,
+            cross_slot_mapping=self.cross_slot_mapping,
+            cross_block_tables=self.cross_block_tables)
+        # Batch may be composed of prefill|decodes, adjust query start indices
+        # to refer to the start of decodes when the two are split apart.
+        # E.g. in tokens:[3 prefills|6 decodes], query_start_loc=[3,9] => [0,6].
+        if self._cached_decode_metadata.query_start_loc is not None:
+            qs = self._cached_decode_metadata.query_start_loc
+            self._cached_decode_metadata.query_start_loc = qs - qs[0]
+        return self._cached_decode_metadata
+
+    def advance_step(self,
+                     model_input: "ModelInputForGPUWithSamplingMetadata",
+                     sampled_token_ids: Optional[torch.Tensor],
+                     block_size: int,
+                     num_seqs: int,
+                     num_queries: int,
+                     turn_prefills_into_decodes: bool = False):
+        """
+        Update metadata in-place to advance one decode step.
+        """
+
+        assert not turn_prefills_into_decodes, \
+            ("Chunked prefill is not supported with rocm_flash_attn yet."
+             "turn_prefills_into_decodes is a Multi-Step + Chunked-Prefill "
+             "specific parameter.")
+
+        # When using cudagraph, the num_seqs is padded to the next captured
+        # batch sized, but num_queries tracks the actual number of requests in
+        # the batch. For --enforce-eager mode, num_seqs == num_queries
+        if num_seqs != num_queries:
+            assert num_seqs > num_queries
+            assert self.use_cuda_graph
+
+        assert self.num_prefills == 0
+        assert self.num_prefill_tokens == 0
+        assert self.num_decode_tokens == num_seqs
+        assert self.slot_mapping.shape == (num_seqs, )
+
+        assert self.seq_lens is not None
+        assert len(self.seq_lens) == num_seqs
+        assert self.seq_lens_tensor is not None
+        assert self.seq_lens_tensor.shape == (num_seqs, )
+        assert self.max_query_len == 1
+        assert self.max_prefill_seq_len == 0
+        assert self.max_decode_seq_len == max(self.seq_lens)
+
+        assert self.query_start_loc is not None
+        assert self.query_start_loc.shape == (num_queries + 1, )
+        assert self.seq_start_loc is not None
+        assert self.seq_start_loc.shape == (num_seqs + 1, )
+
+        assert self.context_lens_tensor is not None
+        assert self.context_lens_tensor.shape == (num_queries, )
+
+        assert self.block_tables is not None
+        assert self.block_tables.shape[0] == num_seqs
+
+        # Update query lengths. Note that we update only queries and not seqs,
+        # since tensors may be padded due to captured cuda graph batch size
+        for i in range(num_queries):
+            self.seq_lens[i] += 1
+        self.max_decode_seq_len = max(self.seq_lens)
+
+        ops.advance_step_flashattn(num_seqs=num_seqs,
+                                   num_queries=num_queries,
+                                   block_size=block_size,
+                                   input_tokens=model_input.input_tokens,
+                                   sampled_token_ids=sampled_token_ids,
+                                   input_positions=model_input.input_positions,
+                                   seq_lens=self.seq_lens_tensor,
+                                   slot_mapping=self.slot_mapping,
+                                   block_tables=self.block_tables)
+
+
+class MUSAFlashAttentionMetadataBuilder(
+        CommonMetadataBuilder[MUSAFlashAttentionMetadata]):
+
+    _metadata_cls = MUSAFlashAttentionMetadata
+
+
+def _make_alibi_bias(alibi_slopes: torch.Tensor,
+                     dtype: torch.dtype,
+                     seq_lens: Optional[List[int]],
+                     make_attn_mask: bool = True) -> List[torch.Tensor]:
+    attn_biases = []
+    if seq_lens:
+        for seq_len in seq_lens:
+            bias = torch.arange(seq_len, dtype=dtype)
+            # NOTE(zhuohan): HF uses
+            #     `bias = bias[None, :].repeat(seq_len, 1)`
+            # here. We find that both biases give the same results, but
+            # the bias below more accurately follows the original ALiBi
+            # paper.
+            bias = bias[None, :] - bias[:, None]
+
+            num_heads = alibi_slopes.shape[0]
+            bias = bias[None, :].repeat(
+                (num_heads, 1, 1)).to(alibi_slopes.device)
+            bias.mul_(alibi_slopes[:, None, None])
+            if make_attn_mask:
+                inf_mask = torch.empty(
+                    (1, seq_len, seq_len),
+                    dtype=bias.dtype).fill_(-torch.inf).triu_(diagonal=1).to(
+                        alibi_slopes.device)
+                attn_biases.append((bias + inf_mask).to(dtype))
+            else:
+                attn_biases.append(bias.to(dtype))
+
+    return attn_biases
+
+
+def _get_seq_len_block_table_args(
+    attn_metadata: MUSAFlashAttentionMetadata,
+    attn_type: str,
+) -> tuple:
+    '''
+    The particular choice of sequence-length
+    attributes which should be extracted from attn_metadata is dependent
+    on the type of attention operation.
+
+    Decoder attn -> select entirely decoder self-attention-related fields
+    Encoder/decoder cross-attn -> select encoder sequence lengths
+    Encoder attn -> select encoder sequence lengths fields
+
+    Arguments:
+
+    * attn_metadata: Attention metadata structure associated with attention op
+    * attn_type: encoder attention, decoder self-attention,
+                encoder/decoder cross-attention
+
+    Returns:
+
+    * Appropriate sequence-lengths tensors for query and key
+    * Appropriate max sequence-length scalar
+    '''
+
+    partial_prefix_sum = 0
+    if attn_type == AttentionType.ENCODER:
+        assert attn_metadata.encoder_seq_lens is not None
+        assert attn_metadata.encoder_seq_lens_tensor is not None
+        query_seq_start_loc = torch.tensor(
+            [0] + [
+                partial_prefix_sum := partial_prefix_sum + i
+                for i in attn_metadata.encoder_seq_lens
+            ],
+            device=attn_metadata.encoder_seq_lens_tensor.device,
+            dtype=attn_metadata.encoder_seq_lens_tensor.dtype)
+        causal_mask = False
+
+        # No block tables associated with encoder attention
+        return (query_seq_start_loc, attn_metadata.max_encoder_seq_len,
+                query_seq_start_loc, attn_metadata.max_encoder_seq_len,
+                attn_metadata.encoder_seq_lens, causal_mask)
+    elif attn_type == AttentionType.DECODER:
+        # Decoder self-attention
+        # Choose max_seq_len based on whether we are in prompt_run
+        assert attn_metadata.seq_lens is not None
+        assert attn_metadata.seq_lens_tensor is not None
+        query_seq_start_loc = torch.tensor(
+            [0] + [
+                partial_prefix_sum := partial_prefix_sum + i
+                for i in attn_metadata.seq_lens
+            ],
+            device=attn_metadata.seq_lens_tensor.device,
+            dtype=attn_metadata.seq_lens_tensor.dtype)
+        max_seq_len = attn_metadata.max_prefill_seq_len
+        causal_mask = True
+
+        return (query_seq_start_loc, max_seq_len, query_seq_start_loc,
+                max_seq_len, attn_metadata.seq_lens, causal_mask)
+    elif attn_type == AttentionType.ENCODER_DECODER:
+        assert attn_metadata.seq_lens is not None
+        assert attn_metadata.encoder_seq_lens_tensor is not None
+        query_start_loc = torch.tensor(
+            [0] + [
+                partial_prefix_sum := partial_prefix_sum + i
+                for i in attn_metadata.seq_lens
+            ],
+            device=attn_metadata.encoder_seq_lens_tensor.device,
+            dtype=attn_metadata.encoder_seq_lens_tensor.dtype)
+
+        partial_prefix_sum = 0
+        assert attn_metadata.encoder_seq_lens is not None
+        assert attn_metadata.seq_lens_tensor is not None
+        key_seq_start_loc = torch.tensor(
+            [0] + [
+                partial_prefix_sum := partial_prefix_sum + i
+                for i in attn_metadata.encoder_seq_lens
+            ],
+            device=attn_metadata.seq_lens_tensor.device,
+            dtype=attn_metadata.seq_lens_tensor.dtype)
+        causal_mask = False
+
+        # Enc/dec cross-attention KVs match encoder sequence length;
+        # cross-attention utilizes special "cross" block tables
+        return (query_start_loc, attn_metadata.max_prefill_seq_len,
+                key_seq_start_loc, attn_metadata.max_encoder_seq_len,
+                attn_metadata.seq_lens, causal_mask)
+    else:
+        raise AttributeError(f"Invalid attention type {str(attn_type)}")
+
+
+class MUSAFlashAttentionImpl(AttentionImpl):
+    """
+    If the input tensors contain prompt tokens, the layout is as follows:
+    |<--------------- num_prompt_tokens -------------->|
+    |<--prompt_0-->|<--prompt_1-->|...|<--prompt_N-1-->|
+
+    Otherwise, the layout is as follows:
+    |<------------------ num_generation_tokens (M) ----------------->|
+    |<--generation_0-->|..........|<--generation_M-1-->|<--padding-->|
+
+    Generation tokens can contain padding when cuda-graph is used.
+    Currently, prompt tokens don't contain any padding.
+
+    The prompts might have different lengths, while the generation tokens
+    always have length 1.
+
+    If chunked prefill is enabled, prefill tokens and decode tokens can be
+    batched together in a flattened 1D query.
+
+    |<----- num_prefill_tokens ---->|<------- num_decode_tokens ----------->|
+    |<-prompt_0->|...|<-prompt_N-1->|<-generation_0->|...|<-generation_M-1->|
+
+    Currently, cuda graph is disabled for chunked prefill, meaning there's no
+    padding between prefill and decode tokens.
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: Optional[List[float]],
+        sliding_window: Optional[int],
+        kv_cache_dtype: str,
+        blocksparse_params: Optional[Dict[str, Any]] = None,
+        logits_soft_cap: Optional[float] = None,
+        attn_type: str = AttentionType.DECODER,
+        kv_sharing_target_layer_name: Optional[int] = None,
+        **extra_impl_args,
+    ) -> None:
+        if blocksparse_params is not None:
+            raise ValueError(
+                "MUSAFlashAttention does not support blocksparse attention.")
+
+        if logits_soft_cap is None:
+            # In flash-attn, setting logits_soft_cap as 0 means no soft cap.
+            self.logits_soft_cap = 0.0
+        else:
+            self.logits_soft_cap = logits_soft_cap
+        self.attn_type = attn_type
+        self.num_heads = num_heads
+        self.head_size = head_size
+        self.scale = float(scale)
+        self.num_kv_heads = num_kv_heads
+        if alibi_slopes is not None:
+            alibi_slopes = torch.tensor(alibi_slopes, dtype=torch.float32)
+        self.alibi_slopes = alibi_slopes
+        self.sliding_window = ((sliding_window, sliding_window)
+                               if sliding_window is not None else (-1, -1))
+        self.kv_cache_dtype = kv_cache_dtype
+
+        assert self.num_heads % self.num_kv_heads == 0
+        self.num_queries_per_kv = self.num_heads // self.num_kv_heads
+
+        supported_head_sizes = PagedAttention.get_supported_head_sizes()
+        if head_size not in supported_head_sizes:
+            raise ValueError(
+                f"Head size {head_size} is not supported by PagedAttention. "
+                f"Supported head sizes are: {supported_head_sizes}.")
+
+        if logits_soft_cap is not None:
+            raise ValueError(
+                "MUSA Triton FlashAttention does not support attention"
+                "logits soft capping."
+                " please try using the MUSA CK "
+                "FA backend instead by setting the env var "
+                "`VLLM_USE_TRITON_FLASH_ATTN=0`")
+
+        # self.attn_func = _sdpa_attention
+        self.use_naive_attn = True
+        self.use_triton_flash_attn = False
+
+        if self.use_naive_attn:
+            self.attn_func = _sdpa_attention
+            logger.debug("Using naive (SDPA) attention in MUSABackend")
+        else:
+            from vllm_musa.ops.triton_flash_attention import (  # noqa: F401
+                triton_attention)
+            self.attn_func = triton_attention
+            logger.debug("Using Triton FA in MUSABackend")
+
+    def repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor:
+        """torch.repeat_interleave(x, dim=1, repeats=n_rep)"""
+        tokens, n_kv_heads, head_dim = x.shape
+        return (x[:, :,
+                  None, :].expand(tokens, n_kv_heads, n_rep,
+                                  head_dim).reshape(tokens, n_kv_heads * n_rep,
+                                                    head_dim))
+
+    def forward(
+        self,
+        layer: AttentionLayer,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        kv_cache: torch.Tensor,
+        attn_metadata: MUSAFlashAttentionMetadata,
+        output: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Forward pass with FlashAttention and PagedAttention.
+
+        For decoder-only models: query, key and value must be non-None.
+
+        For encoder/decoder models:
+        * MUSAFlashAttentionImpl.forward() may be invoked for both self- and
+            cross-attention layers.
+        * For self-attention: query, key and value must be non-None.
+        * For cross-attention:
+            * Query must be non-None
+            * During prefill, key and value must be non-None; key and value
+              get cached for use during decode.
+            * During decode, key and value may be None, since:
+              (1) key and value tensors were cached during prefill, and
+              (2) cross-attention key and value tensors do not grow during
+                  decode
+
+        A note on how the attn_type (attention type enum) argument impacts
+        attention forward() behavior:
+
+            * DECODER: normal decoder-only behavior;
+                use decoder self-attention block table
+            * ENCODER: no KV caching; pass encoder sequence
+                attributes (encoder_seq_lens/encoder_seq_lens_tensor/
+                max_encoder_seq_len) to kernel, in lieu of decoder
+                sequence attributes (seq_lens/seq_lens_tensor/max_seq_len)
+            * ENCODER_DECODER: cross-attention behavior;
+                use cross-attention block table for caching KVs derived
+                from encoder hidden states; since KV sequence lengths
+                will match encoder sequence lengths, pass encoder sequence
+                attributes to kernel (encoder_seq_lens/encoder_seq_lens_tensor/
+                max_encoder_seq_len)
+
+        Args:
+            query: shape = [num_tokens, num_heads * head_size]
+            key: shape = [num_tokens, num_kv_heads * head_size]
+            value: shape = [num_tokens, num_kv_heads * head_size]
+            kv_cache = [2, num_blocks, block_size * num_kv_heads * head_size]
+                NOTE: kv_cache will be an empty tensor with shape [0]
+                for profiling run.
+            attn_metadata: Metadata for attention.
+            attn_type: Select attention type, between encoder attention,
+                       decoder self-attention, or encoder/decoder cross-
+                       attention. Defaults to decoder self-attention,
+                       which is the vLLM default generally
+        Returns:
+            shape = [num_tokens, num_heads * head_size]
+        """
+        query = query.view(-1, self.num_heads, self.head_size)
+        if key is not None:
+            assert value is not None
+            key = key.view(-1, self.num_kv_heads, self.head_size)
+            value = value.reshape(-1, self.num_kv_heads, self.head_size)
+        else:
+            assert value is None
+
+        if self.attn_type != AttentionType.ENCODER and kv_cache.numel() > 0:
+            key_cache, value_cache = PagedAttention.split_kv_cache(
+                kv_cache, self.num_kv_heads, self.head_size)
+
+            if key is not None and value is not None:
+                # Reshape the input keys and values and store them in the
+                # cache. If kv_cache is not provided, the new key and value
+                # tensors are not cached. This happens during the initial
+                # memory profiling run.
+                PagedAttention.write_to_paged_cache(
+                    key,
+                    value,
+                    key_cache,
+                    value_cache,
+                    attn_metadata.slot_mapping
+                    if self.attn_type != AttentionType.ENCODER_DECODER else
+                    attn_metadata.cross_slot_mapping,
+                    self.kv_cache_dtype,
+                    layer._k_scale,
+                    layer._v_scale,
+                )
+
+        if self.attn_type != AttentionType.ENCODER:
+            num_prefill_tokens = attn_metadata.num_prefill_tokens
+        else:
+            assert attn_metadata.num_encoder_tokens is not None
+            num_prefill_tokens = attn_metadata.num_encoder_tokens
+
+        output = torch.empty_like(query)
+        # Query for decode. KV is not needed because it is already cached.
+        decode_query = query[num_prefill_tokens:]
+        # QKV for prefill.
+        query = query[:num_prefill_tokens]
+
+        if key is not None and value is not None \
+            and self.attn_type != AttentionType.ENCODER_DECODER:
+            key = key[:num_prefill_tokens]
+            value = value[:num_prefill_tokens]
+
+        if prefill_meta := attn_metadata.prefill_metadata:
+            # Prompt run.
+            # normal attention and DECODER
+            if self.attn_type == AttentionType.DECODER and (
+                    kv_cache.numel() == 0 or prefill_meta.block_tables is None
+                    or prefill_meta.block_tables.numel() == 0):
+                (query_seq_start_loc, query_max_seq_len, key_seq_start_loc,
+                 key_max_seq_len, seq_lens,
+                 causal_mask) = (prefill_meta.seq_start_loc,
+                                 prefill_meta.max_prefill_seq_len,
+                                 prefill_meta.seq_start_loc,
+                                 prefill_meta.max_prefill_seq_len,
+                                 attn_metadata.seq_lens, True)
+            # prefix-enabled attention and ENCODER/ENCODER_DECODER
+            else:
+                (query_seq_start_loc, query_max_seq_len, key_seq_start_loc,
+                 key_max_seq_len, seq_lens,
+                 causal_mask) = _get_seq_len_block_table_args(
+                     prefill_meta, self.attn_type)
+            # Prompt run.
+            if kv_cache.numel() == 0 or prefill_meta.block_tables.numel() == 0:
+                # triton attention
+                # When block_tables are not filled, it means q and k are the
+                # prompt, and they have the same length.
+                attn_masks = None
+                if self.use_triton_flash_attn:
+                    if self.alibi_slopes is not None:
+                        attn_masks = _make_alibi_bias(
+                            self.alibi_slopes,
+                            query.dtype,
+                            seq_lens,
+                            make_attn_mask=False)  # type: ignore
+                    out, _ = self.attn_func(
+                        query,
+                        key,
+                        value,
+                        None,
+                        query_seq_start_loc,
+                        key_seq_start_loc,
+                        query_max_seq_len,
+                        key_max_seq_len,
+                        causal_mask,
+                        self.scale,
+                        attn_masks[0][None]
+                        if attn_masks is not None else None,
+                    )
+                elif self.use_naive_attn:
+                    if self.num_kv_heads != self.num_heads:
+                        # Interleave for MQA workaround.
+                        key = self.repeat_kv(key, self.num_queries_per_kv)
+                        value = self.repeat_kv(value, self.num_queries_per_kv)
+                    if self.alibi_slopes is not None:
+                        attn_masks = _make_alibi_bias(
+                            self.alibi_slopes,
+                            query.dtype,
+                            attn_metadata.seq_lens,
+                            make_attn_mask=True)  # type: ignore
+                    query = query.movedim(0, query.dim() - 2)
+                    key = key.movedim(0, key.dim() - 2)
+                    value = value.movedim(0, value.dim() - 2)
+                    # sdpa math backend attention
+                    out = self.attn_func(
+                        query,
+                        key,
+                        value,
+                        prefill_meta.query_start_loc,
+                        num_prefill_tokens,
+                        self.num_heads,
+                        self.head_size,
+                        self.scale,
+                        # causal_mask,
+                        attn_masks,
+                    )
+                else:
+                    out = self.attn_func(
+                        q=query,
+                        k=key,
+                        v=value,
+                        cu_seqlens_q=query_seq_start_loc,
+                        cu_seqlens_k=key_seq_start_loc,
+                        max_seqlen_q=prefill_meta.max_prefill_seq_len,
+                        max_seqlen_k=key_max_seq_len,
+                        softmax_scale=self.scale,
+                        causal=True,
+                        window_size=self.sliding_window,
+                        alibi_slopes=self.alibi_slopes,
+                        softcap=self.logits_soft_cap,
+                    )
+
+                # common code for prefill
+                assert output[:num_prefill_tokens].shape == out.shape
+                if output.shape[0] > num_prefill_tokens:
+                    output[:num_prefill_tokens] = out
+                else:
+                    output = out
+            else:
+
+                # prefix-enabled attention
+                output[:num_prefill_tokens] = PagedAttention.forward_prefix(
+                    query,
+                    key,
+                    value,
+                    self.kv_cache_dtype,
+                    key_cache,
+                    value_cache,
+                    prefill_meta.block_tables,
+                    prefill_meta.query_start_loc,
+                    prefill_meta.seq_lens_tensor,
+                    prefill_meta.context_lens_tensor,
+                    prefill_meta.max_query_len,
+                    self.alibi_slopes,
+                    self.sliding_window[0],
+                    layer._k_scale,
+                    layer._v_scale,
+                )
+
+        if decode_meta := attn_metadata.decode_metadata:
+            # Decoding run.
+            # Whether to use rocm custom paged attention or not
+            output[num_prefill_tokens:] = PagedAttention.forward_decode(
+                decode_query,
+                key_cache,
+                value_cache,
+                decode_meta.block_tables
+                if self.attn_type != AttentionType.ENCODER_DECODER else
+                decode_meta.cross_block_tables,
+                decode_meta.seq_lens_tensor
+                if self.attn_type != AttentionType.ENCODER_DECODER else
+                decode_meta.encoder_seq_lens_tensor,
+                decode_meta.max_decode_seq_len
+                if self.attn_type != AttentionType.ENCODER_DECODER else
+                decode_meta.max_encoder_seq_len,
+                self.kv_cache_dtype,
+                self.num_kv_heads,
+                self.scale,
+                self.alibi_slopes,
+                layer._k_scale,
+                layer._v_scale,
+            )
+
+        # Reshape the output tensor.
+        return output.view(-1, self.num_heads * self.head_size)
+
+def pad_seq(x: torch.Tensor, cu_seqlen:torch.Tensor, max_prefill_seq_len: int, num_heads: int) -> torch.Tensor:
+    bs = cu_seqlen.shape[0] - 1
+    padded = torch.zeros(bs, num_heads, max_prefill_seq_len, x.shape[-1], device=x.device, dtype=x.dtype)
+    for i in range(bs):
+        start = cu_seqlen[i].item()
+        end = cu_seqlen[i + 1].item()
+        seq_len = end - start
+        padded[i, :, :seq_len] = x[:, start:end, :]
+    return padded
+
+def restore_tokens(attn_output: torch.Tensor, cu_seqlen: torch.Tensor, max_prefill_seq_len: int, num_heads: int) -> torch.Tensor:
+    bs = cu_seqlen.shape[0] - 1
+    attn_output_tensor = attn_output.transpose(1, 2).view(bs, max_prefill_seq_len, num_heads, -1)
+    restored_tokens = []
+    for i in range(bs):
+        start = cu_seqlen[i].item()
+        end = cu_seqlen[i + 1].item()
+        seq_len = end - start
+        restored_tokens.append(attn_output_tensor[i, :seq_len])
+    return torch.cat(restored_tokens, dim=0)
+
+# pad 支持 varlen
+def _sdpa_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    seq_lens: List[int],
+    num_tokens: int,
+    num_heads: int,
+    head_size: int,
+    scale: float,
+    attn_masks: Optional[List[torch.Tensor]] = None,
+) -> torch.Tensor:
+    max_prefill_seq_len = torch.diff(seq_lens).max()
+
+    q_pad = pad_seq(query, seq_lens, max_prefill_seq_len, num_heads)
+    k_pad = pad_seq(key, seq_lens, max_prefill_seq_len, num_heads)
+    v_pad = pad_seq(value, seq_lens, max_prefill_seq_len, num_heads)
+
+    output, _, _ = torch.ops.aten._scaled_dot_product_attention_flash_musa(
+        q_pad,
+        k_pad,
+        v_pad,
+        dropout_p=0.0,
+        is_causal=attn_masks is None)
+
+    output = restore_tokens(output, seq_lens, max_prefill_seq_len, num_heads)
+    return output.view(-1, num_heads, head_size)
diff --git a/vllm_musa/vllm_musa/worker.py b/vllm_musa/vllm_musa/worker.py
new file mode 100644
index 000000000..ed8bee937
--- /dev/null
+++ b/vllm_musa/vllm_musa/worker.py
@@ -0,0 +1,283 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# This file is a part of the vllm_musa project.
+# Adapted from vllm-project/vllm/vllm/worker/worker.py
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import gc
+import os
+import time
+from datetime import datetime
+from typing import TYPE_CHECKING, Optional, Type, List, Dict, Tuple
+from distutils.version import StrictVersion
+
+import torch
+import torch.distributed
+from vllm import envs
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor import set_random_seed
+from vllm.sequence import SequenceGroupMetadata
+from vllm.utils import (GiB_bytes, MemorySnapshot, memory_profiling)
+from vllm.worker.cache_engine import CacheEngine
+from vllm.worker.enc_dec_model_runner import EncoderDecoderModelRunner
+from vllm.worker.model_runner_base import ModelRunnerBase
+from vllm.worker.model_runner import GPUModelRunnerBase, ModelRunner
+
+from vllm.worker.pooling_model_runner import PoolingModelRunner
+from vllm.worker.worker_base import WorkerBase
+from vllm.worker.worker import Worker
+from vllm.distributed import (ensure_model_parallel_initialized,
+                              init_distributed_environment,
+                              set_custom_all_reduce)
+from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized
+
+
+logger = init_logger(__name__)
+
+class MTGPUWorker(Worker):
+    """A worker class that executes (a partition of) the model on a MTGPU.
+    Each worker is associated with a single MTGPU. The worker is responsible for
+    maintaining the KV cache and executing the model on the MTGPU. In case of
+    distributed inference, each worker is assigned a partition of the model.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        local_rank: int,
+        rank: int,
+        distributed_init_method: str,
+        is_driver_worker: bool = False,
+        model_runner_cls: Optional[Type[ModelRunnerBase]] = None,
+    ) -> None:
+        # Register ops and patch when worker init.
+        from vllm_musa import ops  # noqa: F401
+        from vllm_musa import patch  # noqa: F401
+
+        WorkerBase.__init__(self, vllm_config)
+        self.parallel_config.rank = rank
+        self.local_rank = local_rank
+        self.rank = rank
+        self.distributed_init_method = distributed_init_method
+        self.is_driver_worker = is_driver_worker
+        if self.model_config.trust_remote_code:
+            # note: lazy import to avoid importing torch before initializing
+            from vllm.utils import init_cached_hf_modules
+            init_cached_hf_modules()
+
+        # Return hidden states from target model if the draft model is an
+        # mlp_speculator
+        speculative_config = self.speculative_config
+        model_config = self.model_config
+        speculative_args = {} if speculative_config is None \
+            or (speculative_config.draft_model_config.model ==
+                model_config.model) \
+            or (speculative_config.draft_model_config.hf_config.model_type
+                not in ["medusa", "mlp_speculator", "eagle"]) \
+                    else {"return_hidden_states": True}
+
+        ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner
+        if model_config.runner_type == "pooling":
+            ModelRunnerClass = PoolingModelRunner
+        elif self.model_config.is_encoder_decoder:
+            ModelRunnerClass = EncoderDecoderModelRunner
+        self.model_runner: GPUModelRunnerBase = ModelRunnerClass(
+            vllm_config=self.vllm_config,
+            kv_cache_dtype=self.cache_config.cache_dtype,
+            is_driver_worker=is_driver_worker,
+            **speculative_args,
+        )
+        if model_runner_cls is not None:
+            self.model_runner = model_runner_cls(self.model_runner)
+
+        # Uninitialized cache engine. Will be initialized by
+        # initialize_cache.
+        self.cache_engine: List[CacheEngine]
+        # Initialize gpu_cache as pooling models don't initialize kv_caches
+        self.gpu_cache: Optional[List[List[torch.Tensor]]] = None
+        self._seq_group_metadata_cache: Dict[str, SequenceGroupMetadata] = {}
+
+        # Torch profiler. Enabled and configured through env vars:
+        # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
+        if envs.VLLM_TORCH_PROFILER_DIR:
+            self.torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
+            logger.info(
+                "Profiling enabled. Traces will be saved to: %s",
+                self.torch_profiler_trace_dir,
+            )
+
+            def trace_handler(prof):
+                begin = time.monotonic()
+                current_time = datetime.now().strftime("%Y.%m.%d-%H:%M:%S")
+                rank = torch.distributed.get_rank()
+                save_dir = os.path.join(self.torch_profiler_trace_dir, current_time)
+                if not os.path.exists(save_dir):
+                    os.makedirs(save_dir, exist_ok=True)
+                prof.export_chrome_trace(f"{save_dir}/profile_rank{rank}.json.gz")
+                print(
+                    f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
+                )
+
+            if StrictVersion(torch.__version__) < StrictVersion('2.5.0'):
+                activities=[
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.MUSA,
+                    ]
+            else:
+                activities=[
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.PrivateUse1,
+                    ]
+
+            self.profiler = torch.profiler.profile(
+                activities=activities,
+                profile_memory=True,
+                record_shapes=True,
+                with_flops=True,
+                with_modules=True,
+                with_stack=True,
+                on_trace_ready=trace_handler,
+            )
+        else:
+            self.profiler = None
+
+
+    def init_device(self) -> None:
+        if self.device_config.device.type == "musa":
+            # os.environ["MCCL_PROTOS"] = "2"
+            self.device = torch.device(f"musa:{self.local_rank}")
+            torch.musa.set_device(self.device)
+
+            _check_if_gpu_supports_dtype(self.model_config.dtype)
+            gc.collect()
+            torch.musa.empty_cache()
+        else:
+            raise RuntimeError(
+                f"Not support device type: {self.device_config.device}")
+
+        self.baseline_snapshot = MemorySnapshot()
+        # Initialize the distributed environment.
+        init_worker_distributed_environment(self.vllm_config, self.rank,
+                                            self.distributed_init_method,
+                                            self.local_rank)
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+
+    @torch.inference_mode()
+    def determine_num_available_blocks(self) -> Tuple[int, int]:
+        """Profiles the peak memory usage of the model to determine how many
+        KV blocks may be allocated without OOMs.
+
+        The engine will first conduct a profiling of the existing memory usage.
+        Then, it calculate the maximum possible number of GPU and CPU blocks
+        that can be allocated with the remaining free memory.
+
+        .. tip::
+            You may limit the usage of GPU memory
+            by adjusting the `gpu_memory_utilization` parameter.
+        """
+        # Profile the memory usage of the model and get the maximum number of
+        # cache blocks that can be allocated with the remaining free memory.
+        torch.musa.empty_cache()
+        torch.musa.reset_peak_memory_stats()
+
+        free_memory_pre_profile, total_gpu_memory = torch.musa.mem_get_info()
+
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model.
+        with memory_profiling(
+                self.baseline_snapshot,
+                weights_memory=self.model_runner.model_memory_usage) as result:
+            self.model_runner.profile_run()
+
+        self._assert_memory_footprint_increased_during_profiling()
+
+        memory_for_current_instance = total_gpu_memory * \
+            self.cache_config.gpu_memory_utilization
+        available_kv_cache_memory = (memory_for_current_instance -
+                                     result.non_kv_cache_memory)
+
+        # Calculate the number of blocks that can be allocated with the
+        # profiled peak memory.
+        cache_block_size = self.get_cache_block_size_bytes()
+        if cache_block_size == 0:
+            num_gpu_blocks = 0
+            num_cpu_blocks = 0
+        else:
+            num_gpu_blocks = int(available_kv_cache_memory // cache_block_size)
+            num_cpu_blocks = int(self.cache_config.swap_space_bytes //
+                                 cache_block_size)
+        num_gpu_blocks = max(num_gpu_blocks, 0)
+        num_cpu_blocks = max(num_cpu_blocks, 0)
+
+        msg = (f"Memory profiling takes {result.profile_time:.2f} seconds\n"
+               "the current vLLM instance can use "
+               "total_gpu_memory "
+               f"({(total_gpu_memory / GiB_bytes):.2f}GiB)"
+               " x gpu_memory_utilization "
+               f"({self.cache_config.gpu_memory_utilization:.2f})"
+               f" = {(memory_for_current_instance / GiB_bytes):.2f}GiB\n"
+               "model weights take "
+               f"{(result.weights_memory / GiB_bytes):.2f}GiB;"
+               " non_torch_memory takes "
+               f"{(result.non_torch_increase / GiB_bytes):.2f}GiB;"
+               " PyTorch activation peak memory takes "
+               f"{(result.torch_peak_increase / GiB_bytes):.2f}GiB;"
+               " the rest of the memory reserved for KV Cache is "
+               f"{(available_kv_cache_memory / GiB_bytes):.2f}GiB.")
+
+        logger.info(msg)
+        # Final cleanup
+        gc.collect()
+
+        return num_gpu_blocks, num_cpu_blocks
+
+    def _assert_memory_footprint_increased_during_profiling(self):
+        # NOTE(woosuk): Here we assume that the other processes using the same
+        # GPU did not change their memory usage during the profiling.
+        free_gpu_memory, total = torch.musa.mem_get_info()
+        cuda_memory = total - free_gpu_memory
+        assert self.baseline_snapshot.cuda_memory < cuda_memory, (
+            "Error in memory profiling. "
+            f"Initial used memory {self.baseline_snapshot.cuda_memory}, "
+            f"currently used memory {cuda_memory}. "
+            f"This happens when the GPU memory was "
+            "not properly cleaned up before initializing the vLLM instance.")
+
+
+def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype):
+    # Check if the GPU supports the dtype.
+    if torch_dtype == torch.bfloat16:  # noqa: SIM102
+        return True
+
+def init_worker_distributed_environment(
+    vllm_config: VllmConfig,
+    rank: int,
+    distributed_init_method: Optional[str] = None,
+    local_rank: int = -1,
+) -> None:
+    """Initialize the distributed environment."""
+    parallel_config = vllm_config.parallel_config
+    set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
+
+    init_distributed_environment(parallel_config.world_size, rank,
+                                 distributed_init_method, local_rank, backend="mccl")
+    ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
+                                      parallel_config.pipeline_parallel_size)
+
+    ensure_kv_transfer_initialized(vllm_config)
\ No newline at end of file
diff --git a/vllm_musa/vllm_musa/worker/__init__.py b/vllm_musa/vllm_musa/worker/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/vllm_musa/vllm_musa/worker/musa_worker.py b/vllm_musa/vllm_musa/worker/musa_worker.py
new file mode 100644
index 000000000..1f84a4bfa
--- /dev/null
+++ b/vllm_musa/vllm_musa/worker/musa_worker.py
@@ -0,0 +1,388 @@
+#
+# Copyright (c) 2025 MooreThreads Co., Ltd. All Rights Reserved.
+# Copyright 2023 The vLLM team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""A MTGPU worker class.
+Adapted from vllm-project/vllm/v1/worker/gpu_worker.py
+"""
+import gc
+import os
+import time
+from datetime import datetime
+from typing import TYPE_CHECKING, Optional
+from distutils.version import StrictVersion
+
+import torch
+import torch_musa
+import torch.distributed
+import torch.nn as nn
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.distributed import (ensure_model_parallel_initialized,
+                              init_distributed_environment,
+                              set_custom_all_reduce)
+from vllm.distributed.kv_transfer import ensure_kv_transfer_initialized
+from vllm.distributed.parallel_state import get_pp_group, get_tp_group
+from vllm.logger import init_logger
+from vllm.lora.request import LoRARequest
+from vllm.model_executor import set_random_seed
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+from vllm.utils import GiB_bytes, MemorySnapshot, memory_profiling
+from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
+from vllm.v1.outputs import ModelRunnerOutput
+from vllm.v1.utils import report_usage_stats
+from vllm.v1.worker.worker_base import WorkerBase
+from vllm_musa.patch.v1.patch_gpu_model_runner import CustomGPUModelRunner
+
+logger = init_logger(__name__)
+
+if TYPE_CHECKING:
+    from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
+    from vllm.v1.core.sched.output import SchedulerOutput
+
+
+class MTGPUWorker(WorkerBase):
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        local_rank: int,
+        rank: int,
+        distributed_init_method: str,
+        is_driver_worker: bool = False,
+    ):
+        # Register ops and patch when worker init.
+        from vllm_musa import ops  # noqa: F401
+        from vllm_musa import patch  # noqa: F401
+        super().__init__(vllm_config=vllm_config,
+                         local_rank=local_rank,
+                         rank=rank,
+                         distributed_init_method=distributed_init_method,
+                         is_driver_worker=is_driver_worker)
+
+        if self.model_config.trust_remote_code:
+            # note: lazy import to avoid importing torch before initializing
+            from vllm.utils import init_cached_hf_modules
+            init_cached_hf_modules()
+
+        # Buffers saved before sleep
+        self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
+
+        # Torch profiler. Enabled and configured through env vars:
+        # VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
+        if envs.VLLM_TORCH_PROFILER_DIR:
+            self.torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
+            logger.info(
+                "Profiling enabled. Traces will be saved to: %s",
+                self.torch_profiler_trace_dir,
+            )
+
+            def trace_handler(prof):
+                begin = time.monotonic()
+                current_time = datetime.now().strftime("%Y.%m.%d-%H:%M:%S")
+                rank = torch.distributed.get_rank()
+                save_dir = os.path.join(self.torch_profiler_trace_dir, current_time)
+                if not os.path.exists(save_dir):
+                    os.makedirs(save_dir, exist_ok=True)
+                prof.export_chrome_trace(f"{save_dir}/profile_rank{rank}.json.gz")
+                print(
+                    f"Finished dumping profiler traces in {time.monotonic() - begin:.2f} seconds"
+                )
+
+            if StrictVersion(torch.__version__) < StrictVersion('2.5.0'):
+                activities=[
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.MUSA,
+                    ]
+            else:
+                activities=[
+                        torch.profiler.ProfilerActivity.CPU,
+                        torch.profiler.ProfilerActivity.PrivateUse1,
+                    ]
+
+            self.profiler = torch.profiler.profile(
+                activities=activities,
+                profile_memory=True,
+                record_shapes=True,
+                with_flops=True,
+                with_modules=True,
+                with_stack=True,
+                on_trace_ready=trace_handler,
+            )
+        else:
+            self.profiler = None
+
+    def sleep(self, level: int = 1) -> None:
+        # ModelConfig.enable_sleep_mode
+        logger.error("Sleep mode is only supported on cuda platform")
+
+    def wake_up(self, tags: Optional[list[str]] = None) -> None:
+        # ModelConfig.enable_sleep_mode
+        logger.error("Sleep mode is only supported on cuda platform")
+
+    def initialize_cache(self, num_gpu_blocks: int,
+                         num_cpu_blocks: int) -> None:
+        self.cache_config.num_gpu_blocks = num_gpu_blocks
+        self.cache_config.num_cpu_blocks = num_cpu_blocks
+
+    def init_device(self):
+        if self.device_config.device.type == "musa":
+            # torch.distributed.all_reduce does not free the input tensor until
+            # the synchronization point. This causes the memory usage to grow
+            # as the number of all_reduce calls increases. This env var disables
+            # this behavior.
+            # Related issue:
+            # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573
+            os.environ["TORCH_MCCL_AVOID_RECORD_STREAMS"] = "1"
+
+            # This env var set by Ray causes exceptions with graph building.
+            os.environ.pop("MCCL_ASYNC_ERROR_HANDLING", None)
+            self.device = torch.device(f"musa:{self.local_rank}")
+            torch.musa.set_device(self.device)
+
+            gc.collect()
+            torch.musa.empty_cache()
+
+            # take current memory snapshot
+            self.init_snapshot = MemorySnapshot()
+            self.requested_memory = (self.init_snapshot.total_memory *
+                                     self.cache_config.gpu_memory_utilization)
+            if self.init_snapshot.free_memory < self.requested_memory:
+                GiB = lambda b: round(b / GiB_bytes, 2)
+                raise ValueError(
+                    f"Free memory on device "
+                    f"({GiB(self.init_snapshot.free_memory)}/"
+                    f"{GiB(self.init_snapshot.total_memory)} GiB) on startup "
+                    f"is less than desired GPU memory utilization "
+                    f"({self.cache_config.gpu_memory_utilization}, "
+                    f"{GiB(self.requested_memory)} GiB). Decrease GPU memory "
+                    f"utilization or reduce GPU memory used by other processes."
+                )
+        else:
+            raise RuntimeError(
+                f"Not support device type: {self.device_config.device}")
+        # Initialize the distributed environment.
+        init_worker_distributed_environment(self.vllm_config, self.rank,
+                                            self.distributed_init_method,
+                                            self.local_rank)
+        # Set random seed.
+        set_random_seed(self.model_config.seed)
+
+        # Construct the model runner
+        self.model_runner: CustomGPUModelRunner = CustomGPUModelRunner(
+            self.vllm_config, self.device)
+
+        if self.rank == 0:
+            # If usage stat is enabled, collect relevant info.
+            report_usage_stats(self.vllm_config)
+
+    def load_model(self) -> None:
+        self.model_runner.load_model()
+
+    @torch.inference_mode()
+    def determine_available_memory(self) -> int:
+        """Profiles the peak memory usage of the model to determine how much
+        memory can be used for KV cache without OOMs.
+
+        The engine will first conduct a profiling of the existing memory usage.
+        Then, it calculate the free memory that can be used for KV cache in
+        bytes.
+
+        Tip:
+            You may limit the usage of GPU memory
+            by adjusting the `gpu_memory_utilization` parameter.
+        """
+        torch.musa.empty_cache()
+        torch.musa.reset_peak_memory_stats()
+        GiB = lambda b: b / GiB_bytes
+
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model.
+        with memory_profiling(
+                self.init_snapshot,
+                weights_memory=int(
+                    self.model_runner.model_memory_usage)) as profile_result:
+            self.model_runner.profile_run()
+
+        free_gpu_memory = profile_result.after_profile.free_memory
+        # NOTE(woosuk): Here we assume that the other processes using the same
+        # GPU did not change their memory usage during the profiling.
+        assert self.init_snapshot.free_memory > free_gpu_memory, (
+            "Error in memory profiling. "
+            f"Initial free memory {GiB(self.init_snapshot.free_memory)} GiB, "
+            f"current free memory {GiB(free_gpu_memory)} GiB. "
+            "This happens when other processes sharing the same container "
+            "release GPU memory while vLLM is profiling during initialization. "
+            "To fix this, ensure consistent GPU memory allocation or "
+            "isolate vLLM in its own container.")
+        available_kv_cache_memory = self.requested_memory \
+            - profile_result.non_kv_cache_memory
+
+        logger.debug(
+            "Initial free memory: %.2f GiB, free memory: %.2f GiB, "
+            "requested GPU memory: %.2f GiB",
+            GiB(self.init_snapshot.free_memory), GiB(free_gpu_memory),
+            GiB(self.requested_memory))
+        logger.debug(profile_result)
+        logger.info("Available KV cache memory: %.2f GiB",
+                    GiB(available_kv_cache_memory))
+        gc.collect()
+
+        return int(available_kv_cache_memory)
+
+    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
+        return self.model_runner.get_kv_cache_spec()
+
+    def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
+        """Allocate GPU KV cache with the specified kv_cache_config."""
+        self.model_runner.initialize_kv_cache(kv_cache_config)
+
+    def compile_or_warm_up_model(self) -> None:
+        # warm up sizes that are not in cudagraph capture sizes,
+        # but users still want to compile for better performance,
+        # e.g. for the max-num-batched token size in chunked prefill.
+        warmup_sizes = self.vllm_config.compilation_config.compile_sizes.copy()
+        if not self.model_config.enforce_eager:
+            warmup_sizes = [
+                x for x in warmup_sizes if x not in
+                self.vllm_config.compilation_config.cudagraph_capture_sizes
+            ]
+        # We skip EPLB here since we don't want to record dummy metrics
+        for size in sorted(warmup_sizes, reverse=True):
+            logger.info("Compile and warming up model for size %d", size)
+            self.model_runner._dummy_run(size, skip_eplb=True)
+        if not self.model_config.enforce_eager:
+            self.model_runner.capture_model()
+
+        # Warm up sampler and preallocate memory buffer for logits and other
+        # sampling related tensors of max possible shape to avoid memory
+        # fragmentation issue.
+        # NOTE: This is called after `capture_model` on purpose to prevent
+        # memory buffers from being cleared by `torch.cuda.empty_cache`.
+        if get_pp_group().is_last_rank:
+            max_num_reqs = min(self.scheduler_config.max_num_seqs,
+                               self.scheduler_config.max_num_batched_tokens)
+            # We skip EPLB here since we don't want to record dummy metrics
+            hidden_states, last_hidden_states = \
+                self.model_runner._dummy_run(
+                    num_tokens=max_num_reqs,
+                    skip_eplb=True,
+                )
+            if self.model_runner.is_pooling_model:
+                self.model_runner._dummy_pooler_run(hidden_states)
+            else:
+                self.model_runner._dummy_sampler_run(
+                    hidden_states=last_hidden_states)
+
+        # Reset the seed to ensure that the random state is not affected by
+        # the model initialization and profiling.
+        set_random_seed(self.model_config.seed)
+
+    def get_model(self) -> nn.Module:
+        return self.model_runner.get_model()
+
+    @torch.inference_mode()
+    def execute_model(
+        self,
+        scheduler_output: "SchedulerOutput",
+    ) -> Optional[ModelRunnerOutput]:
+        intermediate_tensors = None
+        if not get_pp_group().is_first_rank:
+            intermediate_tensors = IntermediateTensors(
+                get_pp_group().recv_tensor_dict(
+                    all_gather_group=get_tp_group()))
+
+        output = self.model_runner.execute_model(scheduler_output,
+                                                 intermediate_tensors)
+        parallel_config = self.vllm_config.parallel_config
+        if parallel_config.distributed_executor_backend != "external_launcher" \
+            and not get_pp_group().is_last_rank:
+            assert isinstance(output, IntermediateTensors)
+            get_pp_group().send_tensor_dict(output.tensors,
+                                            all_gather_group=get_tp_group())
+            return None
+        assert isinstance(output, ModelRunnerOutput)
+        return output if self.is_driver_worker else None
+
+    def profile(self, is_start: bool = True):
+        if self.profiler is None:
+            raise RuntimeError("Profiler is not enabled.")
+        if is_start:
+            self.profiler.start()
+        else:
+            self.profiler.stop()
+
+    def execute_dummy_batch(self) -> None:
+        self.model_runner._dummy_run(1)
+
+    def add_lora(self, lora_request: LoRARequest) -> bool:
+        return self.model_runner.add_lora(lora_request)
+
+    def remove_lora(self, lora_id: int) -> bool:
+        return self.model_runner.remove_lora(lora_id)
+
+    def list_loras(self) -> set[int]:
+        return self.model_runner.list_loras()
+
+    def pin_lora(self, lora_id: int) -> bool:
+        return self.model_runner.pin_lora(lora_id)
+
+    def check_health(self) -> None:
+        # worker will always be healthy as long as it's running.
+        return
+
+    def save_sharded_state(
+        self,
+        path: str,
+        pattern: Optional[str] = None,
+        max_size: Optional[int] = None,
+    ) -> None:
+        from vllm.model_executor.model_loader import ShardedStateLoader
+        ShardedStateLoader.save_model(
+            self.model_runner.model,
+            path,
+            pattern=pattern,
+            max_size=max_size,
+        )
+
+    def save_tensorized_model(
+        self,
+        tensorizer_config: "TensorizerConfig",
+    ) -> None:
+        self.model_runner.save_tensorized_model(
+            tensorizer_config=tensorizer_config, )
+
+
+def init_worker_distributed_environment(
+    vllm_config: VllmConfig,
+    rank: int,
+    distributed_init_method: Optional[str] = None,
+    local_rank: int = -1,
+    backend: str = "mccl",
+) -> None:
+    """Initialize the distributed environment."""
+    parallel_config = vllm_config.parallel_config
+    set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
+
+    init_distributed_environment(parallel_config.world_size, rank,
+                                 distributed_init_method, local_rank, backend)
+
+    ensure_model_parallel_initialized(parallel_config.tensor_parallel_size,
+                                      parallel_config.pipeline_parallel_size)
+
+    ensure_kv_transfer_initialized(vllm_config)
