# omni-ops_1015 **Repository Path**: june86/omni-ops_1015 ## Basic Information - **Project Name**: omni-ops_1015 - **Description**: No description available - **Primary Language**: C++ - **License**: Not specified - **Default Branch**: master - **Homepage**: None - **GVP Project**: No ## Statistics - **Stars**: 0 - **Forks**: 139 - **Created**: 2026-10-15 - **Last Updated**: 2026-10-16 ## Categories & Tags **Categories**: Uncategorized **Tags**: None ## README # omni-ops Ascend-compatible high-performance custom operator repository for large model training and inference scenarios, providing operators implemented via AscendC, PyPTO, and Triton. These operators are exposed through the PyTorch Adapter (PTA) as `torch.ops.custom.npu_*` interfaces, enabling direct invocation within PyTorch, enabling direct invocation within PyTorch. ## Project Structure ``` omni-ops/ ├── inference/ # Inference operators │ └── ascendc/ # AscendC implementation │ ├── cmake/ # CMake build modules and scripts │ ├── scripts/ # Build/deployment helper scripts │ ├── src/ # Source directory │ │ ├── ops-nn/ # NN operators │ │ │ └── matmul/ # Matrix multiplication operators │ │ │ └── ai_infra_matmul/ # AI Infra Matmul │ │ ├── ops-transformer/ # Transformer operators │ │ │ ├── attention/ # Attention operators │ │ │ │ ├── ai_infra_sparse_flash_attention_gqa/ # Sparse Flash Attention (GQA) │ │ │ │ ├── ai_infra_sparse_flash_attention_pioneer/ # Sparse Flash Attention (Pioneer) │ │ │ │ ├── ai_infra_kv_quant_sparse_flash_attention/ # KV Quantized Sparse Flash Attention │ │ │ │ ├── ai_infra_fused_infer_attention_sink/ # Fused Inference Attention Sink │ │ │ │ ├── ai_infra_fused_infer_attention_sink_metadata/ # Attention Sink Metadata │ │ │ │ ├── ai_infra_fused_causal_conv1d/ # Fused Causal Conv1D │ │ │ │ ├── ai_infra_chunk_gated_delta_rule_recurrence/ # Chunked Gated Delta Rule Recurrence │ │ │ │ ├── ai_infra_esa_select_topk/ # ESA TopK Selection │ │ │ │ ├── ai_infra_lower_triangular_inverse/ # Lower Triangular Matrix Inverse │ │ │ │ ├── ai_infra_quant_lightning_indexer/ # Quantized Lightning Indexer │ │ │ │ └── ai_infra_recurrent_gated_delta_rule/ # Recurrent Gated Delta Rule │ │ │ ├── mhc/ # Manifold Constrained Hyper Connection │ │ │ │ ├── ai_infra_mhc_pre_split_post_res/ # MHC Pre Split Post Res │ │ │ │ └── ai_infra_mhc_sandwich_norm_post_preonly/ # MHC Sandwich Norm Post │ │ │ ├── index/ # Index operators │ │ │ │ └── ai_infra_scatter_block_update/ # Scatter Block Update │ │ │ └── posembedding/ # Positional embedding operators │ │ │ └── ai_infra_kv_rms_norm_rope_cache/ # KV RMSNorm RoPE Cache │ │ ├── tests/ # Test framework code │ │ └── utils/ # Common utilities │ └── torch_ops_extension/ # PyTorch operator extension interface │ └── omni_custom_ops/ # Inference custom operator package │ ├── training/ # Training operators │ ├── ascendc/ # AscendC implementation │ │ ├── cmake/ # CMake build modules and scripts │ │ ├── src/ # Source directory │ │ │ ├── ops-transformer/ # Transformer operators │ │ │ │ ├── attention/ # Attention operators │ │ │ │ │ ├── flash_attention_score_enhance/ # Flash Attention Score Enhance │ │ │ │ │ ├── flash_attention_score_grad_enhance/ # Flash Attention Score Grad Enhance │ │ │ │ │ ├── sparse_flash_attention_enhance/ # Sparse Flash Attention Enhance │ │ │ │ │ ├── sparse_flash_attention_grad_enhance/ # Sparse Flash Attention Grad Enhance │ │ │ │ │ ├── lightning_indexer_enhance/ # Lightning Indexer Enhance │ │ │ │ │ ├── sparse_lightning_indexer_grad_kl_loss_enhance/ # Sparse Lightning Indexer KL Loss Grad │ │ │ │ │ ├── ai_infra_attention_pioneer/ # Attention Pioneer Forward (MLA Sink Token) │ │ │ │ │ ├── ai_infra_attention_pioneer_backward/ # Attention Pioneer Backward │ │ │ │ │ ├── ai_infra_attention_pioneer_metadata/ # Attention Pioneer Metadata │ │ │ │ │ └── ai_infra_flash_attention_score/ # Flash Attention Score Fusion │ │ │ │ ├── mhc/ # MHC operators (forward/backward) │ │ │ │ │ ├── ai_infra_mega_mhc/ # Fused Mega MHC │ │ │ │ │ ├── ai_infra_manifold_constrained_hyper_connection_pre/ # MHC Pre-process │ │ │ │ │ ├── ai_infra_manifold_constrained_hyper_connection_pre_grad/ # MHC Pre-process Grad │ │ │ │ │ ├── ai_infra_manifold_constrained_hyper_connection_post/ # MHC Post-process │ │ │ │ │ ├── ai_infra_mhc_post_grad/ # MHC Post Grad │ │ │ │ │ ├── ai_infra_sinkhorn_grad/ # Sinkhorn Grad │ │ │ │ │ └── manifold_constrained_hyper_connection_sinkhorn_enhance/ # MHC Sinkhorn Enhance │ │ │ │ ├── mome/ # MoME Mixture-of-Experts operators │ │ │ │ │ ├── ai_infra_aggregate_hidden/ # Aggregate Hidden State │ │ │ │ │ └── ai_infra_aggregate_hidden_grad/ # Aggregate Hidden State Grad │ │ │ │ ├── posembedding/ # Positional embedding operators │ │ │ │ │ ├── ai_infra_apply_rotary_pos_emb/ # QK Apply Rotary Position Embedding │ │ │ │ │ └── ai_infra_apply_rotary_pos_emb_grad/ # QK Apply Rotary Position Embedding Grad │ │ │ │ └── common/ # Common components │ │ │ ├── ops-communication/ # Collective communication operators (based on HCCL) │ │ │ │ ├── common/ # Common components │ │ │ │ └── ai_infra_all_gather_batch/ # Batched AllGather │ │ │ ├── tests/ # Test framework code │ │ │ └── utils/ # Common utilities │ │ └── torch_ops_extension/ # PyTorch operator extension interface │ │ └── omni_training_custom_ops/ # Training custom operator package │ │ │ ├── pypto/ # PyTorch PTO implementation │ │ └── src/ # Source directory │ │ └── ops-nn/ # NN operators │ │ └── quant/ # Quantization operators │ │ │ └── triton/ # Triton implementation │ └── src/ # Source directory │ └── ops_transformer/ # Transformer operators │ └── attention/ # Attention operators │ └── gated_delta_net/ # Gated Delta Net │ └── .gitee/ # Gitee configuration (Issue templates, etc.) ``` ## Operator List ### Inference Operators (inference/ascendc) | Category | Operator Name | Description | |----------|---------------|-------------| | Attention | ai_infra_sparse_flash_attention_gqa | Sparse Flash Attention (GQA) | | Attention | ai_infra_sparse_flash_attention_pioneer | Sparse Flash Attention (Pioneer) | | Attention | ai_infra_kv_quant_sparse_flash_attention | KV Quantized Sparse Flash Attention | | Attention | ai_infra_fused_infer_attention_sink | Fused Inference Attention Sink | | Attention | ai_infra_fused_causal_conv1d | Fused Causal Conv1D | | Attention | ai_infra_chunk_gated_delta_rule_recurrence | Chunked Gated Delta Rule Recurrence | | Attention | ai_infra_esa_select_topk | ESA TopK Selection | | Attention | ai_infra_lower_triangular_inverse | Lower Triangular Matrix Inverse | | Attention | ai_infra_quant_lightning_indexer | Quantized Lightning Indexer | | Attention | ai_infra_recurrent_gated_delta_rule | Recurrent Gated Delta Rule | | MHC | ai_infra_mhc_pre_split_post_res | MHC Pre Split Post Res | | MHC | ai_infra_mhc_sandwich_norm_post_preonly | MHC Sandwich Norm Post | | Index | ai_infra_scatter_block_update | Scatter Block Update | | PosEmbedding | ai_infra_kv_rms_norm_rope_cache | KV RMSNorm RoPE Cache | | Matmul | ai_infra_matmul | General Matrix Multiplication | ### Training Operators (training/ascendc) | Category | Operator Name | Description | |----------|---------------|-------------| | Attention | flash_attention_score_enhance | Flash Attention Score Enhance | | Attention | flash_attention_score_grad_enhance | Flash Attention Score Grad Enhance | | Attention | sparse_flash_attention_enhance | Sparse Flash Attention Enhance | | Attention | sparse_flash_attention_grad_enhance | Sparse Flash Attention Grad Enhance | | Attention | lightning_indexer_enhance | Lightning Indexer Enhance | | Attention | sparse_lightning_indexer_grad_kl_loss_enhance | Sparse Lightning Indexer KL Loss Grad | | Attention | ai_infra_attention_pioneer | Flash Attention with Sink Token (MLA), supports Prefill/Decode | | Attention | ai_infra_attention_pioneer_backward | Attention Pioneer Backward | | Attention | ai_infra_attention_pioneer_metadata | Attention Pioneer Metadata | | Attention | ai_infra_flash_attention_score | Fused Flash Attention Score | | MHC | ai_infra_mega_mhc | Fused RMSNorm/MHC Post/MHC Pre | | MHC | ai_infra_manifold_constrained_hyper_connection_pre | MHC Pre-process | | MHC | ai_infra_manifold_constrained_hyper_connection_pre_grad | MHC Pre-process Grad | | MHC | ai_infra_manifold_constrained_hyper_connection_post | MHC Post-process | | MHC | ai_infra_mhc_post_grad | MHC Post-process Grad (Cube-Vector fusion op)| | MHC | ai_infra_sinkhorn_grad | Sinkhorn Grad | | MHC | manifold_constrained_hyper_connection_sinkhorn_enhance | MHC Sinkhorn Enhance | | MoME | ai_infra_aggregate_hidden | Aggregate Hidden State | | MoME | ai_infra_aggregate_hidden_grad | Aggregate Hidden State Grad | | PosEmbedding | ai_infra_apply_rotary_pos_emb | QK Apply Rotary Position Embedding | | PosEmbedding | ai_infra_apply_rotary_pos_emb_grad | QK Apply Rotary Position Embedding Grad | | HCCL | ai_infra_all_gather_batch | Batched AllGather | ### Training Operators (training/triton) | Category | Operator Name | Description | |----------|---------------|-------------| | Attention | gated_delta_net | Gated Delta Net | ### Training Operators (training/pypto) | Category | Operator Name | Description | |----------|---------------|-------------| | Quant | quant | Quantization Operator | ## Tech Stack - **AscendC**: Ascend-native C/C++ operator development framework, directly operates NPU hardware resources for optimal performance - **Triton**: Triton-based operator development, suitable for rapid prototyping - **PyPTO**: PyTorch-based operator development approach ## Hardware Support - Ascend 910B (Atlas A2) - Ascend 910C (Atlas A3) - Ascend 950PR (Atlas A5) ## 📝 Related Information - [License](./LICENSE) For operators involved in the omni-ops warehouse, if a License exists under the operator directory, that License shall prevail. If no License is present in the operator directory, the CANN 2.0 protocol shall apply; the corresponding protocol text can be found at [LICENSE](./LICENSE). - [Disclaimer](./DISCLAIMER.md)