Compare commits
4 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 3a5fda8ce3 | |||
| 71b752346b | |||
| ebfeac9863 | |||
| 7f8a1b1f7a |
134
.gitea/workflows/docker-build-push.yml
Normal file
134
.gitea/workflows/docker-build-push.yml
Normal file
@@ -0,0 +1,134 @@
|
||||
name: Docker Build and Push
|
||||
|
||||
on:
|
||||
push:
|
||||
tags:
|
||||
- "v*"
|
||||
|
||||
jobs:
|
||||
docker:
|
||||
runs-on: amd64-ubuntu-24.04
|
||||
|
||||
steps:
|
||||
- name: Clone repository
|
||||
run: |
|
||||
git clone "${{ gitea.server_url }}/${{ gitea.repository }}.git" .
|
||||
git checkout "${{ gitea.ref_name }}"
|
||||
|
||||
- name: Set image metadata
|
||||
run: |
|
||||
IMAGE_NAME="$(echo "${{ gitea.repository }}" | tr '[:upper:]' '[:lower:]' | tr '_' '-')"
|
||||
IMAGE="${DOCKER_REGISTRY}/${DOCKER_USERNAME}/${IMAGE_NAME}:${{ gitea.ref_name }}"
|
||||
|
||||
echo "IMAGE_NAME=${IMAGE_NAME}" >> "$GITEA_ENV"
|
||||
echo "IMAGE=${IMAGE}" >> "$GITEA_ENV"
|
||||
|
||||
- name: Load and Validate Task Info
|
||||
run: |
|
||||
set -a
|
||||
. .gitea/workflows/task_info.env
|
||||
set +a
|
||||
|
||||
for name in FRAMEWORK GPU_TYPE TASK_TYPE; do
|
||||
eval "value=\${${name}:-}"
|
||||
if [ "$name" = "FRAMEWORK" ] && [ -z "$value" ]; then
|
||||
echo "${name} is empty in .gitea/workflows/task_info.env"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "${name}=${value}" >> "$GITEA_ENV"
|
||||
done
|
||||
|
||||
- name: Validate Image Verify Metadata
|
||||
run: |
|
||||
if [ -z "${FIXED_TOKEN:-}" ]; then
|
||||
echo "FIXED_TOKEN is not configured on runner"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if ! response="$(curl --silent --show-error --location --get 'https://modelhub.org.cn/adminApi/image-verify/validate' \
|
||||
--header "Xc-Token: ${FIXED_TOKEN}" \
|
||||
--data-urlencode "gpuType=${GPU_TYPE:-}" \
|
||||
--data-urlencode "taskType=${TASK_TYPE:-}")"; then
|
||||
echo "failed to call image verify validate API"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
VALIDATE_RESPONSE="$response" python3 - <<'PY'
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
raw = os.environ.get("VALIDATE_RESPONSE", "")
|
||||
try:
|
||||
body = json.loads(raw)
|
||||
except json.JSONDecodeError:
|
||||
print("image verify validate API returned invalid JSON")
|
||||
print(raw)
|
||||
sys.exit(1)
|
||||
|
||||
if body.get("code") == 0 and body.get("data") is True:
|
||||
print("image verify metadata validation passed")
|
||||
sys.exit(0)
|
||||
|
||||
message = body.get("message") or "unknown error"
|
||||
print(f"image verify metadata validation failed: {message}")
|
||||
print(raw)
|
||||
sys.exit(1)
|
||||
PY
|
||||
|
||||
- name: Login to Docker Registry
|
||||
run: |
|
||||
echo "$DOCKER_PASSWORD" | docker login "$DOCKER_REGISTRY" \
|
||||
-u "$DOCKER_USERNAME" \
|
||||
--password-stdin
|
||||
|
||||
- name: Build Docker Image
|
||||
run: |
|
||||
docker build -t "$IMAGE" .
|
||||
|
||||
- name: Push Docker Image
|
||||
run: |
|
||||
for attempt in 1 2 3; do
|
||||
echo "Starting docker push attempt ${attempt}/3 for ${IMAGE}"
|
||||
docker push "$IMAGE" &
|
||||
PUSH_PID=$!
|
||||
|
||||
while kill -0 "$PUSH_PID" 2>/dev/null; do
|
||||
echo "docker push is still running at $(date -u '+%Y-%m-%dT%H:%M:%SZ')"
|
||||
sleep 60
|
||||
done
|
||||
|
||||
if wait "$PUSH_PID"; then
|
||||
echo "docker push completed successfully"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
echo "docker push failed on attempt ${attempt}/3"
|
||||
sleep 30
|
||||
done
|
||||
|
||||
echo "docker push failed after 3 attempts"
|
||||
exit 1
|
||||
|
||||
- name: Notify Image Verify
|
||||
run: |
|
||||
if [ -z "${FIXED_TOKEN:-}" ]; then
|
||||
echo "FIXED_TOKEN is not configured on runner"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
curl --silent --show-error --fail-with-body --location --request POST 'https://modelhub.org.cn//adminApi/image-verify' \
|
||||
--header "Xc-Token: ${FIXED_TOKEN}" \
|
||||
--header 'Content-Type: application/json' \
|
||||
--data-raw "{
|
||||
\"framework\": \"${FRAMEWORK}\",
|
||||
\"gpuType\": \"${GPU_TYPE}\",
|
||||
\"imageUrl\": \"${IMAGE}\",
|
||||
\"taskType\": \"${TASK_TYPE}\",
|
||||
\"createBy\": \"${{ gitea.actor }}\",
|
||||
\"repoUrl\": \"${{ gitea.server_url }}/${{ gitea.repository }}\",
|
||||
\"tag\": \"${{ github.ref_name }}\"
|
||||
}"
|
||||
|
||||
|
||||
3
.gitea/workflows/task_info.env
Normal file
3
.gitea/workflows/task_info.env
Normal file
@@ -0,0 +1,3 @@
|
||||
FRAMEWORK=vllm_0.23.0
|
||||
GPU_TYPE=Ascend_910-b3
|
||||
TASK_TYPE=text-generation
|
||||
419
AGENTS.md
Normal file
419
AGENTS.md
Normal file
@@ -0,0 +1,419 @@
|
||||
# vLLM Ascend Development Guidelines
|
||||
|
||||
This document provides instructions for contributors to the vLLM Ascend project. Please read and follow these guidelines to ensure code quality, maintainability, and consistency.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Setup and Environment](#setup-and-environment)
|
||||
- [Environment Variables](#environment-variables)
|
||||
- [Environment Variable Review Requirement](#environment-variable-review-requirement)
|
||||
- [Testing](#testing)
|
||||
- [Unit and System Tests](#unit-and-system-tests)
|
||||
- [Running Tests](#running-tests)
|
||||
- [Code Style](#code-style)
|
||||
- [Python Conventions](#python-conventions)
|
||||
- [Naming Conventions](#naming-conventions)
|
||||
- [NPU-Specific Considerations](#npu-specific-considerations)
|
||||
- [Tensor item() Operations](#tensor-item-operations)
|
||||
- [Memory and Performance](#memory-and-performance)
|
||||
- [Model and Plugin Architecture](#model-and-plugin-architecture)
|
||||
- [vLLM Ascend Plugin Architecture](#vllm-ascend-plugin-architecture)
|
||||
- [Patching Requirement](#patching-requirement)
|
||||
- [Model Runner Changes](#model-runner-changes)
|
||||
- [Commit Messages and Pull Requests](#commit-messages-and-pull-requests)
|
||||
- [Commit Message Format](#commit-message-format)
|
||||
- [Review Checklist](#review-checklist)
|
||||
- [Code Quality](#code-quality)
|
||||
- [Testing](#testing-1)
|
||||
- [Documentation](#documentation)
|
||||
- [NPU Considerations](#npu-considerations)
|
||||
- [Commit and PR](#commit-and-pr)
|
||||
- [Quick Start for Contributors](#quick-start-for-contributors)
|
||||
- [References](#references)
|
||||
|
||||
---
|
||||
|
||||
## Setup and Environment
|
||||
|
||||
### Environment Variables
|
||||
|
||||
All environment variables must be defined in `vllm_ascend/envs.py` using the centralized `env_variables` dictionary.
|
||||
|
||||
**Requirements:**
|
||||
|
||||
- Add documentation for each environment variable in the `env_variables` dict comment
|
||||
- Specify default values and valid ranges
|
||||
- Indicate whether the variable is sensitive (credentials, keys)
|
||||
|
||||
**Example:**
|
||||
|
||||
```python
|
||||
import os
|
||||
|
||||
env_variables = {
|
||||
"VLLM_ASCEND_ENABLE_NZ": lambda: int(os.getenv("VLLM_ASCEND_ENABLE_NZ", 1)),
|
||||
# ...
|
||||
}
|
||||
```
|
||||
|
||||
**Never**: Hardcode environment variable names throughout the codebase. Reference them from the central module using `from vllm_ascend import envs`.
|
||||
|
||||
### Environment Variable Review Requirement
|
||||
|
||||
**Strict Review Required**: All new environment variables must undergo code review.
|
||||
|
||||
Reviewers must verify:
|
||||
|
||||
- The variable name follows the `VLLM_ASCEND_*` naming convention
|
||||
- Default value is appropriate for all supported hardware
|
||||
- Documentation is added to the `env_variables` dict
|
||||
- The variable is used in a performance-critical path
|
||||
|
||||
---
|
||||
|
||||
## Testing
|
||||
|
||||
### Unit and System Tests
|
||||
|
||||
**Requirement**: All new functionality requires corresponding tests.
|
||||
|
||||
- **Unit Tests (UT)**: Located in `tests/ut/`, cover core logic, edge cases, and error conditions
|
||||
- **System Tests (ST)**: Located in `tests/e2e/`, verify end-to-end behavior and integration points
|
||||
- **Nightly Tests**: Include benchmarks for NPU-specific code paths in `tests/e2e/nightly/`
|
||||
|
||||
**Test Coverage Guidelines:**
|
||||
|
||||
- New features: Tests must cover happy path and failure modes
|
||||
- Bug fixes: Tests must include a regression test for the bug
|
||||
- Performance-critical code: Include benchmarks and performance regression tests
|
||||
|
||||
### Running Tests
|
||||
|
||||
```bash
|
||||
|
||||
|
||||
# Run specific unit test file
|
||||
pytest -sv tests/ut/ops/test_prepare_finalize.py
|
||||
|
||||
# Run specific unit test
|
||||
pytest -sv tests/ut/ops/test_prepare_finalize.py::test_prepare_inputs
|
||||
|
||||
# Run NPU-specific tests (requires NPU hardware)
|
||||
pytest -sv tests/e2e/pull_request/one_card/aclgraph/test_aclgraph_accuracy.py::test_default_full_and_piecewise_res_consistency
|
||||
```
|
||||
|
||||
**Requirement**: Run all tests locally before requesting review. Verify tests pass on NPU hardware for NPU-specific changes.
|
||||
|
||||
---
|
||||
|
||||
## Code Style
|
||||
|
||||
### Python Conventions
|
||||
|
||||
- **Imports**: All imports at the top of the file. Valid exceptions:
|
||||
- Circular imports (use inline imports)
|
||||
- Lazy loading for worker/isolation processes
|
||||
- Type-checking imports wrapped in `if TYPE_CHECKING:`
|
||||
|
||||
- **Global Variables**: Avoid new global variables. Pass dependencies explicitly through function parameters.
|
||||
|
||||
**Allowed:**
|
||||
- Constants named `ALL_UPPER_CASE` (e.g., `MAX_BATCH_SIZE` in `envs.py`)
|
||||
- Immutable configuration objects
|
||||
|
||||
**Requires Approval:**
|
||||
- Any new mutable global state
|
||||
|
||||
- **No Magic Numbers**: Use named constants with descriptive names:
|
||||
|
||||
```python
|
||||
# Bad
|
||||
if seq_len > 2048: ...
|
||||
|
||||
# Good
|
||||
MAX_CONTEXT_LENGTH = 2048
|
||||
if seq_len > MAX_CONTEXT_LENGTH: ...
|
||||
```
|
||||
|
||||
- **Descriptive Naming**: Use names that describe functionality, not implementation details.
|
||||
|
||||
```python
|
||||
# Bad
|
||||
is_deepseek_v3_r1
|
||||
flag1
|
||||
tmp_var
|
||||
|
||||
# Good
|
||||
supports_dynamic_temperature
|
||||
uses_speculative_decoding
|
||||
```
|
||||
|
||||
### Naming Conventions
|
||||
|
||||
- **Classes**: `PascalCase` (e.g., `NPUModelRunner`, `AscendSampler`, `ACLGraphManager`)
|
||||
- **Functions/Methods**: `snake_case` (e.g., `forward_pass`, `compute_attention`)
|
||||
- **Constants**: `ALL_UPPER_CASE` (e.g., `MAX_BATCH_SIZE`, `VLLM_ASCEND_ENABLE_NZ`)
|
||||
- **Variables**: `snake_case` (e.g., `token_ids`, `sequence_lengths`)
|
||||
|
||||
---
|
||||
|
||||
## NPU-Specific Considerations
|
||||
|
||||
### Tensor item() Operations
|
||||
|
||||
**Warning**: `tensor.item()` operations cause synchronization overhead on NPU when the `tensor` is on device.
|
||||
|
||||
If the `tensor` is a device tensor, calling `item()` will triggers a synchronous data transfer from NPU to CPU. This can severely degrade performance in hot paths, causing `AsyncScheduler` to block here.
|
||||
|
||||
**Review Requirements:**
|
||||
|
||||
1. Profile performance impact before merging
|
||||
2. Consider alternative patterns:
|
||||
- Keep values on device when possible
|
||||
- Batch operations to reduce sync frequency
|
||||
- Use device-side operations (e.g., `torch.argmax`, `torch.sum`)
|
||||
3. Document when `item()` is unavoidable (e.g., logging, conditional logic)
|
||||
|
||||
**Example Patterns:**
|
||||
|
||||
```python
|
||||
# Bad: In hot loop - causes sync per iteration
|
||||
for tensor in tensors:
|
||||
value = tensor.item()
|
||||
|
||||
# Better: Batch operations - single sync
|
||||
values = [t.item() for t in tensors] # Single batch sync
|
||||
|
||||
# Good: Keep on device when possible
|
||||
max_value = torch.max(tensor) # No sync needed
|
||||
if max_value > threshold: # Comparison can stay on device
|
||||
...
|
||||
```
|
||||
|
||||
### Memory and Performance
|
||||
|
||||
Additional NPU-specific best practices:
|
||||
|
||||
- Avoid CPU-NPU memory transfers in hot paths
|
||||
- Prefer in-place operations where safe (e.g., `x.add_()`, `x.mul_()`)
|
||||
- Monitor memory fragmentation, especially for long-running processes
|
||||
- Test with realistic workloads on actual NPU hardware (Ascend 910B/C)
|
||||
|
||||
---
|
||||
|
||||
## Model and Plugin Architecture
|
||||
|
||||
### vLLM Ascend Plugin Architecture
|
||||
|
||||
vLLM Ascend is a **hardware plugin** that integrates with upstream vLLM via the pluggable hardware interface. It does not add new model files directly.
|
||||
|
||||
**Required Pattern**: Model-specific functionality should be implemented via:
|
||||
|
||||
1. **Patching** (in `vllm_ascend/patch/`):
|
||||
- `vllm_ascend/patch/platform/` - Platform-level patches (distributed, scheduling)
|
||||
- `vllm_ascend/patch/worker/` - Worker-level patches (model-specific behavior)
|
||||
- Example: `patch_deepseek.py` modifies upstream Deepseek model behavior
|
||||
- Patch is not the best solution for all cases. Use it when necessary.
|
||||
|
||||
2. **Inheritance**:
|
||||
- `NPUModelRunner(GPUModelRunner)` - Extend vLLM model runner with NPU-specific behavior
|
||||
- `AscendSampler` - Extend vLLM sampler with NPU-specific operations
|
||||
- Add NPU-specific components via composition (e.g., `AclGraphManager`)
|
||||
- Custom Operators - NPU-specific custom operators (e.g., `AscendRMSNorm`)
|
||||
|
||||
3. **External upstream contributions** where appropriate
|
||||
|
||||
### Patching Requirement
|
||||
|
||||
**Strict Review Required**: All new patches must undergo thorough architectural review.
|
||||
|
||||
Reviewers must verify:
|
||||
|
||||
- The patch targets the correct upstream component
|
||||
- The patch is minimal and focused
|
||||
- Performance implications are understood
|
||||
- A long-term plan exists for upstream contribution
|
||||
|
||||
**Example Patch Pattern:**
|
||||
|
||||
```python
|
||||
# vllm_ascend/patch/worker/patch_deepseek.py
|
||||
from vllm.model_executor.models.deepseek_v2 import DeepseekV2Model
|
||||
|
||||
def forward(self, input_ids, positions, ...):
|
||||
# NPU-specific forward implementation
|
||||
...
|
||||
|
||||
DeepseekV2Model.forward = forward # Patch upstream class
|
||||
```
|
||||
|
||||
### Model Runner Changes
|
||||
|
||||
**Strict Review Required**: All new behaviors added to `model_runner` must undergo thorough architectural review.
|
||||
|
||||
Reviewers must verify:
|
||||
|
||||
- The necessity of the new behavior (why can't this be in a patch?)
|
||||
- Performance implications on NPU hardware
|
||||
- Compatibility with existing model implementations
|
||||
- Long-term maintainability and test coverage
|
||||
|
||||
**NPU Model Runner Files:**
|
||||
|
||||
- `vllm_ascend/worker/model_runner_v1.py` - vLLM v1 model runner
|
||||
- `vllm_ascend/worker/v2/model_runner.py` - vLLM v2 model runner
|
||||
- `vllm_ascend/_310p/model_runner_310p.py` - Ascend 310P model runner
|
||||
|
||||
---
|
||||
|
||||
## Commit Messages and Pull Requests
|
||||
|
||||
### Commit Message Format
|
||||
|
||||
Follow the [Conventional Commits](https://www.conventionalcommits.org/) format and **must include a sign-off**:
|
||||
|
||||
```bash
|
||||
git commit -s -m "<type>: <summary>" -m "<body - explaining what changed and why>"
|
||||
```
|
||||
|
||||
Or using the full message format:
|
||||
|
||||
```txt
|
||||
<type>: <summary>
|
||||
|
||||
<body - explaining what changed and why>
|
||||
|
||||
Signed-off-by: Your Name <your.email@example.com>
|
||||
```
|
||||
|
||||
**Valid Types**: `feat`, `fix`, `perf`, `refactor`, `test`, `docs`, `chore`
|
||||
|
||||
**Good Examples:**
|
||||
|
||||
```txt
|
||||
feat(npu): add flash attention support for Ascend CANN
|
||||
|
||||
- Implements FlashAttention-2 kernel for NPU backend
|
||||
- Reduces memory usage by 30% compared to baseline
|
||||
|
||||
fix(model_runner): correct padding token handling
|
||||
|
||||
- Fixes token padding that caused incorrect attention masks
|
||||
- Addresses issue #1234
|
||||
|
||||
perf: avoid CPU-NPU sync in attention computation
|
||||
|
||||
- Inline computation to avoid tensor.item() calls
|
||||
- Improves throughput by 15%
|
||||
```
|
||||
|
||||
**Bad Examples:**
|
||||
|
||||
```txt
|
||||
fix bug
|
||||
add feature
|
||||
update code
|
||||
```
|
||||
|
||||
### Pull Request Title Format
|
||||
|
||||
PR titles should follow the format: `[Type][Module] Description`
|
||||
|
||||
- **Type**: The type of change (e.g., `CI`, `Doc`, `BugFix`, `Feat`, `Platform`, `Refactor`)
|
||||
- **Module**: The affected module (optional, e.g., `Misc`, `Model`, `Worker`)
|
||||
- **Description**: Brief description of the change
|
||||
|
||||
**Examples:**
|
||||
|
||||
- `[Doc][Misc] Update contribution guidelines`
|
||||
- `[BugFix] Fix CPU binding logic`
|
||||
- `[CI] Update image build workflow`
|
||||
|
||||
### Pull Request Template
|
||||
|
||||
When creating a PR, please follow the template in `.github/PULL_REQUEST_TEMPLATE.md` and ensure the following sections are completed:
|
||||
|
||||
> **Note**: The PR description will be automatically updated by GitHub Actions to include vLLM version info at the bottom. If you update the PR description via API or CLI, make sure to preserve the `- vLLM version:` and `- vLLM main:` lines.
|
||||
|
||||
- **What this PR does / why we need it?** - Clearly describe the changes and their purpose
|
||||
- **Does this PR introduce _any_ user-facing change?** - Indicate if there are any user-visible changes
|
||||
- **How was this patch tested?** - Describe how you tested the changes. Examples:
|
||||
- Unit tests added/updated: list the test files
|
||||
- Manual testing: provide the test steps and commands
|
||||
- CI testing: indicate if only CI verification is needed
|
||||
|
||||
---
|
||||
|
||||
## Review Checklist
|
||||
|
||||
Before merging, verify:
|
||||
|
||||
### Code Quality
|
||||
|
||||
- [ ] Code follows style guidelines (naming, imports, no magic numbers)
|
||||
- [ ] No global state added without justification
|
||||
- [ ] Patching pattern used correctly (if applicable)
|
||||
- [ ] No direct model file additions
|
||||
|
||||
### Testing
|
||||
|
||||
- [ ] New tests added for new functionality (`tests/ut/` or `tests/e2e/`)
|
||||
- [ ] Existing tests pass
|
||||
- [ ] NPU-specific tests verified on actual hardware
|
||||
- [ ] Performance benchmarks included where applicable
|
||||
|
||||
### Documentation
|
||||
|
||||
- [ ] Environment variables documented
|
||||
- [ ] Public APIs documented
|
||||
- [ ] User-facing changes reflected in docs
|
||||
|
||||
### NPU Considerations
|
||||
|
||||
- [ ] `tensor.item()` usage reviewed for performance impact
|
||||
- [ ] No unnecessary CPU-NPU transfers in hot paths
|
||||
- [ ] Memory usage verified on NPU hardware
|
||||
|
||||
### Commit and PR
|
||||
|
||||
- [ ] Commit messages are clear and descriptive, following Conventional Commits format
|
||||
- [ ] **All commits are signed off** (`git commit -s`)
|
||||
- [ ] PR is created from your fork repository, not directly from the main repository
|
||||
- [ ] PR description is complete, following the PR template
|
||||
- [ ] All review comments addressed
|
||||
|
||||
---
|
||||
|
||||
## Quick Start for Contributors
|
||||
|
||||
1. Install development dependencies: `pip install -e .[dev]`
|
||||
2. Run tests: `pytest tests/`
|
||||
3. Check linting: `ruff check vllm_ascend/`
|
||||
4. Format code: `ruff format vllm_ascend/`
|
||||
5. Make your changes following guidelines in this document
|
||||
6. Add tests for new behavior
|
||||
7. Run full test suite before committing
|
||||
8. Commit with sign-off: `git commit -s`
|
||||
9. Run linting check before pushing:
|
||||
```bash
|
||||
bash format.sh ci
|
||||
```
|
||||
> **Note**: This check is required for **all file types**, including markdown files. If `markdownlint` modifies files, re-add them with `git add` and commit again.
|
||||
10. Push to your fork repository (NOT the main repository):
|
||||
|
||||
```bash
|
||||
git remote add myfork https://github.com/YOUR_USERNAME/vllm-ascend.git
|
||||
git push -u myfork your-branch-name
|
||||
```
|
||||
|
||||
11. Create a PR from your fork to the main repository with clear description
|
||||
|
||||
---
|
||||
|
||||
## References
|
||||
|
||||
- [vLLM Hardware Plugin RFC](https://github.com/vllm-project/vllm/issues/11162)
|
||||
- [Documentation](https://docs.vllm.ai/projects/ascend/en/latest/)
|
||||
- [Contributors Guide](https://docs.vllm.ai/projects/ascend/en/latest/community/contributors.html)
|
||||
1
CLAUDE.md
Normal file
1
CLAUDE.md
Normal file
@@ -0,0 +1 @@
|
||||
IMPORTANT: Ensure you've thoroughly reviewed the [AGENTS.md](AGENTS.md) file before beginning any work.
|
||||
136
CMakeLists.txt
136
CMakeLists.txt
@@ -10,8 +10,7 @@ include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake)
|
||||
# Suppress potential warnings about unused manually-specified variables
|
||||
set(ignoreMe "${VLLM_PYTHON_PATH}")
|
||||
|
||||
# TODO: Add 3.12 back when torch-npu support 3.12
|
||||
set(PYTHON_SUPPORTED_VERSIONS "3.9" "3.10" "3.11")
|
||||
set(PYTHON_SUPPORTED_VERSIONS "3.9" "3.10" "3.11" "3.12")
|
||||
|
||||
find_package(pybind11 REQUIRED)
|
||||
|
||||
@@ -20,6 +19,13 @@ set(VLLM_ASCEND_INSTALL_PATH "${CMAKE_INSTALL_PREFIX}")
|
||||
|
||||
find_package(Torch REQUIRED)
|
||||
|
||||
run_python(TORCH_VERSION
|
||||
"import torch; print(torch.__version__)" "Failed to locate torch path")
|
||||
# check torch version is 2.10.0
|
||||
if(NOT ${TORCH_VERSION} VERSION_EQUAL "2.10.0")
|
||||
message(FATAL_ERROR "Expected PyTorch version 2.10.0, but found ${TORCH_VERSION}")
|
||||
endif()
|
||||
|
||||
set(RUN_MODE "npu" CACHE STRING "cpu/sim/npu")
|
||||
set(SOC_VERSION ${SOC_VERSION})
|
||||
message(STATUS "Detected SOC version: ${SOC_VERSION}")
|
||||
@@ -44,17 +50,47 @@ else()
|
||||
endif()
|
||||
|
||||
include(${ASCENDC_CMAKE_DIR}/ascendc.cmake)
|
||||
|
||||
file(GLOB KERNEL_FILES
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/kernels/*.cpp)
|
||||
|
||||
ascendc_library(vllm_ascend_kernels SHARED
|
||||
set(VLLM_ASCEND_CUSTOM_OP
|
||||
${KERNEL_FILES}
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/mla_preprocess/op_kernel/mla_preprocess_kernel.cpp
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/batch_matmul_transpose/op_kernel/batch_matmul_transpose_kernel.cpp
|
||||
)
|
||||
|
||||
set(VLLM_ASCEND_CUSTOM_OP_EXCLUDE_ASCEND950
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/mla_preprocess/op_kernel/mla_preprocess_kernel.cpp
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/batch_matmul_transpose/op_kernel/batch_matmul_transpose_kernel.cpp
|
||||
)
|
||||
|
||||
if(SOC_VERSION MATCHES "ascend950")
|
||||
message(STATUS "A5 hardware detected: disabling MLAPO operators")
|
||||
message(STATUS "A5 hardware detected: excluding batch_matmul_transpose operators")
|
||||
list(REMOVE_ITEM VLLM_ASCEND_CUSTOM_OP ${VLLM_ASCEND_CUSTOM_OP_EXCLUDE_ASCEND950})
|
||||
endif()
|
||||
|
||||
if(SOC_VERSION MATCHES "ascend310p.*|ascend950")
|
||||
message(STATUS "Hardware ${SOC_VERSION} detected: skip vllm_ascend_kernels compile")
|
||||
else()
|
||||
ascendc_library(vllm_ascend_kernels SHARED
|
||||
${VLLM_ASCEND_CUSTOM_OP}
|
||||
)
|
||||
endif()
|
||||
|
||||
message("TORCH_NPU_PATH is ${TORCH_NPU_PATH}")
|
||||
|
||||
file(GLOB VLLM_ASCEND_SRC
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/*.cpp)
|
||||
if(SOC_VERSION MATCHES "ascend310p.*")
|
||||
file(GLOB VLLM_ASCEND_SRC
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/*.cpp
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/aclnn_torch_adapter/*.cpp)
|
||||
else()
|
||||
file(GLOB VLLM_ASCEND_SRC
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/*.cpp
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/aclnn_torch_adapter/*.cpp
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/batch_matmul_transpose/op_host/tiling/tiling_data.cpp)
|
||||
endif()
|
||||
|
||||
include_directories(
|
||||
${pybind11_INCLUDE_DIRS}
|
||||
@@ -62,8 +98,7 @@ include_directories(
|
||||
${TORCH_INCLUDE_DIRS}
|
||||
${TORCH_NPU_PATH}/include
|
||||
${ASCEND_HOME_PATH}/include
|
||||
${ASCEND_HOME_PATH}/aarch64-linux/include/experiment/platform
|
||||
${ASCEND_HOME_PATH}/x86_64-linux/include/experiment/platform
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/csrc/batch_matmul_transpose/op_host
|
||||
)
|
||||
|
||||
set(
|
||||
@@ -71,28 +106,99 @@ set(
|
||||
${TORCH_INCLUDE_DIRS}
|
||||
${TORCH_NPU_INCLUDE_DIRS}
|
||||
${ASCEND_HOME_PATH}/include
|
||||
${ASCEND_HOME_PATH}/aarch64-linux/include/experiment/platform
|
||||
)
|
||||
|
||||
pybind11_add_module(vllm_ascend_C ${VLLM_ASCEND_SRC})
|
||||
|
||||
# Detect aclrtMemcpyBatchAsync availability (CANN 8.5+)
|
||||
# Can be overridden via VLLM_ASCEND_ENABLE_BATCH_MEMCPY env var (registered
|
||||
# in vllm_ascend/envs.py, forwarded by setup.py as a CMake variable):
|
||||
# VLLM_ASCEND_ENABLE_BATCH_MEMCPY=1 -> force enable
|
||||
# VLLM_ASCEND_ENABLE_BATCH_MEMCPY=0 -> force disable
|
||||
# unset -> auto-detect from CANN headers
|
||||
include(CheckCXXSourceCompiles)
|
||||
set(CMAKE_REQUIRED_INCLUDES ${ASCEND_HOME_PATH}/include)
|
||||
set(CMAKE_REQUIRED_LIBRARIES ascendcl)
|
||||
set(CMAKE_REQUIRED_LINK_OPTIONS "-L${ASCEND_HOME_PATH}/lib64")
|
||||
|
||||
if(DEFINED VLLM_ASCEND_ENABLE_BATCH_MEMCPY)
|
||||
if("${VLLM_ASCEND_ENABLE_BATCH_MEMCPY}" STREQUAL "1")
|
||||
message(STATUS "aclrtMemcpyBatchAsync: force enabled via VLLM_ASCEND_ENABLE_BATCH_MEMCPY=1")
|
||||
target_compile_definitions(vllm_ascend_C PRIVATE CANN_MEMCPY_BATCH_ASYNC)
|
||||
else()
|
||||
message(STATUS "aclrtMemcpyBatchAsync: force disabled via VLLM_ASCEND_ENABLE_BATCH_MEMCPY=0")
|
||||
endif()
|
||||
else()
|
||||
# Test the full code pattern we actually use, including struct member access.
|
||||
# This ensures the macro is only defined when the API is fully compatible.
|
||||
check_cxx_source_compiles("
|
||||
#include <acl/acl_rt.h>
|
||||
int main() {
|
||||
aclrtMemLocation loc = {};
|
||||
loc.type = ACL_MEM_LOCATION_TYPE_HOST;
|
||||
loc.id = 0;
|
||||
aclrtMemcpyBatchAttr attr = {};
|
||||
attr.srcLoc = loc;
|
||||
attr.dstLoc = loc;
|
||||
(void)aclrtMemcpyBatchAsync;
|
||||
return 0;
|
||||
}
|
||||
" HAVE_ACLRT_MEMCPY_BATCH_ASYNC)
|
||||
if(HAVE_ACLRT_MEMCPY_BATCH_ASYNC)
|
||||
message(STATUS "aclrtMemcpyBatchAsync: detected in CANN headers, enabling batch memcpy path")
|
||||
target_compile_definitions(vllm_ascend_C PRIVATE CANN_MEMCPY_BATCH_ASYNC)
|
||||
else()
|
||||
message(STATUS "aclrtMemcpyBatchAsync: not found in CANN headers, using fallback aclrtMemcpyAsync loop")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
if(SOC_VERSION MATCHES "ascend310p.*")
|
||||
target_compile_definitions(vllm_ascend_C PRIVATE -DASCEND_PLATFORM_310P)
|
||||
endif()
|
||||
|
||||
if(NOT (SOC_VERSION MATCHES "ascend310p.*|ascend950"))
|
||||
target_compile_definitions(vllm_ascend_C PRIVATE -DVLLM_ENABLE_ATB_AND_DIRECT_KERNELS)
|
||||
endif()
|
||||
|
||||
target_link_directories(
|
||||
vllm_ascend_C
|
||||
PRIVATE
|
||||
${TORCH_LIBRARY_DIRS}
|
||||
${TORCH_NPU_PATH}/lib/
|
||||
${ASCEND_HOME_PATH}/lib64
|
||||
)
|
||||
|
||||
target_link_libraries(
|
||||
vllm_ascend_C
|
||||
PUBLIC
|
||||
set(VLLM_ASCEND_C_COMMON_LIBS
|
||||
${TORCH_LIBRARIES}
|
||||
libtorch_npu.so
|
||||
vllm_ascend_kernels
|
||||
torch_npu
|
||||
ascendcl
|
||||
tiling_api
|
||||
register
|
||||
platform
|
||||
ascendalog
|
||||
dl
|
||||
opapi
|
||||
)
|
||||
|
||||
target_link_options(vllm_ascend_C PRIVATE "-Wl,-rpath,$ORIGIN:$ORIGIN/lib")
|
||||
if(SOC_VERSION MATCHES "ascend310p.*|ascend950")
|
||||
target_link_libraries(
|
||||
vllm_ascend_C
|
||||
PUBLIC
|
||||
${VLLM_ASCEND_C_COMMON_LIBS}
|
||||
)
|
||||
else()
|
||||
target_link_libraries(
|
||||
vllm_ascend_C
|
||||
PUBLIC
|
||||
vllm_ascend_kernels
|
||||
${VLLM_ASCEND_C_COMMON_LIBS}
|
||||
)
|
||||
endif()
|
||||
|
||||
install(TARGETS vllm_ascend_C vllm_ascend_kernels DESTINATION ${VLLM_ASCEND_INSTALL_PATH})
|
||||
target_link_options(vllm_ascend_C PRIVATE "-Wl,-rpath,$ORIGIN:$ORIGIN/lib:$ORIGIN/_cann_ops_custom/vendors/custom_transformer/op_api/lib")
|
||||
|
||||
if(SOC_VERSION MATCHES "ascend310p.*|ascend950")
|
||||
install(TARGETS vllm_ascend_C DESTINATION ${VLLM_ASCEND_INSTALL_PATH})
|
||||
else()
|
||||
install(TARGETS vllm_ascend_C vllm_ascend_kernels DESTINATION ${VLLM_ASCEND_INSTALL_PATH})
|
||||
endif()
|
||||
|
||||
@@ -59,9 +59,7 @@ representative at an online or offline/IRL event.
|
||||
|
||||
## Enforcement
|
||||
|
||||
Instances of abusive, harassing, or otherwise unacceptable behavior may be
|
||||
reported to the community leaders responsible for enforcement in the #code-of-conduct
|
||||
channel in the [vLLM Discord](https://discord.com/invite/jz7wjKhh6g).
|
||||
If you experience any abusive, harassing, or otherwise unacceptable behavior, please feel free to report it to us via [Ascend opensource assistant](docs/source/community/images/ascend_assistant.png) or <contact@ascend.osinfra.cn>.
|
||||
All complaints will be reviewed and investigated promptly and fairly.
|
||||
|
||||
All community leaders are obligated to respect the privacy and security of the
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
# Contributing to vLLM Ascend
|
||||
|
||||
You may find information about contributing to vLLM Ascend on [Developer Guide - Contributing](https://vllm-ascend.readthedocs.io/en/latest/developer_guide/contribution/index.html), including step-by-step guide to help you setup development environment, contribute first PR and test locally.
|
||||
You may find information about contributing to vLLM Ascend on [Developer Guide - Contributing](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/contribution/index.html), including step-by-step guide to help you setup development environment, contribute first PR and test locally.
|
||||
|
||||
59
Dockerfile
59
Dockerfile
@@ -1,60 +1,3 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies 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.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:8.2.rc1-910b-ubuntu22.04-py3.11
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
|
||||
# Define environments
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV COMPILE_CUSTOM_KERNELS=${COMPILE_CUSTOM_KERNELS}
|
||||
|
||||
RUN apt-get update -y && \
|
||||
apt-get install -y python3-pip git vim wget net-tools gcc g++ cmake libnuma-dev && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL}
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.11.0rc3
|
||||
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -v -e /vllm-workspace/vllm/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN python3 -m pip install modelscope 'ray>=2.47.1' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
FROM quay.io/ascend/vllm-ascend:v0.23.0
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -15,47 +15,97 @@
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:8.2.rc1-310p-ubuntu22.04-py3.11
|
||||
FROM quay.io/ascend/cann:9.1.0-310p-ubuntu22.04-py3.12
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
|
||||
# Define environments
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV COMPILE_CUSTOM_KERNELS=${COMPILE_CUSTOM_KERNELS}
|
||||
|
||||
RUN apt-get update -y && \
|
||||
apt-get install -y python3-pip git vim wget net-tools gcc g++ cmake libnuma-dev && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
RUN apt-get update -y && \
|
||||
apt-get install -y python3-pip git vim wget net-tools gcc g++ cmake numactl libnuma-dev libjemalloc2 pciutils && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL}
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL} && \
|
||||
python3 -m pip install modelscope 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.11.0rc3
|
||||
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
|
||||
ARG VLLM_TAG=v0.23.0
|
||||
ARG VLLM_COMMIT=""
|
||||
RUN if [ -n "$VLLM_COMMIT" ]; then \
|
||||
git init /vllm-workspace/vllm && \
|
||||
git -C /vllm-workspace/vllm fetch --depth 1 $VLLM_REPO "$VLLM_COMMIT" && \
|
||||
git -C /vllm-workspace/vllm checkout FETCH_HEAD; \
|
||||
else \
|
||||
git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm; \
|
||||
fi
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -v -e /vllm-workspace/vllm/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
ARG SOC_VERSION="ascend310p1"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV SOC_VERSION=$SOC_VERSION \
|
||||
TASK_QUEUE_ENABLE=1 \
|
||||
OMP_NUM_THREADS=1
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
export SOC_VERSION=ASCEND310P3 && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton-ascend triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN python3 -m pip install modelscope 'ray>=2.47.1' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
|
||||
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc
|
||||
RUN echo 'export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/lib64:/root/aicpu_kernels/0/aicpu_kernels_device/sand_box:$LD_LIBRARY_PATH' >> /root/.bashrc
|
||||
|
||||
# Initialize driver directories, users and file groups
|
||||
RUN mkdir -m 750 /var/driver -m 750 /var/dmp -m 750 /usr/slog && \
|
||||
mkdir -m 755 /home/drv && \
|
||||
mkdir -m 750 /home/drv/hdc_ppc -m 750 -p /var/log/npu/slog && \
|
||||
ldconfig && \
|
||||
ln -sf /lib /lib64 && \
|
||||
userdel -r ubuntu || true && \
|
||||
groupadd -g 1000 HwHiAiUser && useradd -u 1000 -g HwHiAiUser -d /home/HwHiAiUser -m HwHiAiUser && \
|
||||
groupadd -g 1101 HwDmUser && useradd -u 1101 -g HwDmUser -d /home/HwDmUser -m HwDmUser && \
|
||||
groupadd -g 1102 HwBaseUser && useradd -u 1102 -g HwBaseUser -d /home/HwBaseUser -m HwBaseUser && \
|
||||
groupadd -g 1100 HwSysUser && useradd -u 1100 -g HwSysUser -d /home/HwSysUser -m HwSysUser && \
|
||||
usermod -a -G HwBaseUser HwHiAiUser && \
|
||||
usermod -a -G HwDmUser HwHiAiUser && \
|
||||
usermod -a -G HwBaseUser HwDmUser && \
|
||||
usermod -a -G HwHiAiUser HwDmUser && \
|
||||
usermod -a -G HwSysUser HwSysUser && \
|
||||
chown HwDmUser:HwDmUser /var/dmp && \
|
||||
chown HwHiAiUser:HwHiAiUser /var/driver && \
|
||||
chown HwHiAiUser:HwHiAiUser /usr/slog && \
|
||||
chown HwHiAiUser:HwHiAiUser /home/drv/hdc_ppc && \
|
||||
chown HwHiAiUser:HwHiAiUser /var/log/npu/slog
|
||||
|
||||
# Dynamically generate entrypoint.sh for dual mode compatibility
|
||||
RUN echo '#!/bin/bash' > /usr/local/bin/entrypoint.sh && \
|
||||
echo 'if lspci 2>/dev/null | grep -qi "accelerators"; then' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' :' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo 'else' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' su - HwHiAiUser -c "export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/:/usr/lib64 && /lib/ld-linux-aarch64.so.1 /var/slogd -d >/dev/null &"' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' su - HwDmUser -c "export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/:/usr/lib64 && /var/dmp_daemon -I -M -U 8087 &"' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' ' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' rm -rf /lib64/ld-linux-aarch64.so.1' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' ln -sf /lib/ld-linux-aarch64.so.1 /lib64/ld-linux-aarch64.so.1' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo 'fi' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo '' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo 'exec "$@"' >> /usr/local/bin/entrypoint.sh && \
|
||||
chmod +x /usr/local/bin/entrypoint.sh
|
||||
|
||||
# Set entrypoint and default command
|
||||
ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -15,45 +15,93 @@
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:8.2.rc1-310p-openeuler24.03-py3.11
|
||||
FROM quay.io/ascend/cann:9.1.0-310p-openeuler24.03-py3.12
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
|
||||
ENV COMPILE_CUSTOM_KERNELS=${COMPILE_CUSTOM_KERNELS}
|
||||
|
||||
RUN yum update -y && \
|
||||
yum install -y python3-pip git vim wget net-tools gcc gcc-c++ make cmake numactl-devel && \
|
||||
rm -rf /var/cache/yum
|
||||
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL}
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
RUN yum update -y && \
|
||||
yum install -y python3-pip git vim wget net-tools gcc gcc-c++ make cmake numactl numactl-devel jemalloc patch && \
|
||||
rm -rf /var/cache/yum
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL} && \
|
||||
python3 -m pip install modelscope 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.11.0rc3
|
||||
|
||||
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
|
||||
ARG VLLM_TAG=v0.23.0
|
||||
ARG VLLM_COMMIT=""
|
||||
RUN if [ -n "$VLLM_COMMIT" ]; then \
|
||||
git init /vllm-workspace/vllm && \
|
||||
git -C /vllm-workspace/vllm fetch --depth 1 $VLLM_REPO "$VLLM_COMMIT" && \
|
||||
git -C /vllm-workspace/vllm checkout FETCH_HEAD; \
|
||||
else \
|
||||
git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm; \
|
||||
fi
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
ARG SOC_VERSION="ascend310p1"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
ENV SOC_VERSION=$SOC_VERSION \
|
||||
TASK_QUEUE_ENABLE=1 \
|
||||
OMP_NUM_THREADS=1
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:/usr/include/c++/12:/usr/include/c++/12/`uname -i`-openEuler-linux && \
|
||||
export SOC_VERSION=ASCEND310P3 && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton-ascend triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN python3 -m pip install modelscope 'ray>=2.47.1' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
RUN echo "export LD_PRELOAD=/usr/lib64/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
|
||||
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc
|
||||
RUN echo 'export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64:/usr/lib64:/root/aicpu_kernels/0/aicpu_kernels_device/sand_box:$LD_LIBRARY_PATH' >> /root/.bashrc
|
||||
|
||||
# Initialize driver directories, users and file groups
|
||||
RUN mkdir -m 750 /var/driver -m 750 /var/dmp -m 750 /usr/slog && \
|
||||
mkdir -m 755 /home/drv && \
|
||||
mkdir -m 750 /home/drv/hdc_ppc -m 750 -p /var/log/npu/slog && \
|
||||
ldconfig && \
|
||||
ln -sf /lib /lib64 && \
|
||||
groupadd -g 1000 HwHiAiUser && useradd -u 1000 -g HwHiAiUser -d /home/HwHiAiUser -m HwHiAiUser && \
|
||||
groupadd -g 1101 HwDmUser && useradd -u 1101 -g HwDmUser -d /home/HwDmUser -m HwDmUser && \
|
||||
groupadd -g 1102 HwBaseUser && useradd -u 1102 -g HwBaseUser -d /home/HwBaseUser -m HwBaseUser && \
|
||||
groupadd -g 1100 HwSysUser && useradd -u 1100 -g HwSysUser -d /home/HwSysUser -m HwSysUser && \
|
||||
usermod -a -G HwBaseUser HwHiAiUser && \
|
||||
usermod -a -G HwDmUser HwHiAiUser && \
|
||||
usermod -a -G HwBaseUser HwDmUser && \
|
||||
usermod -a -G HwHiAiUser HwDmUser && \
|
||||
usermod -a -G HwSysUser HwSysUser && \
|
||||
chown HwDmUser:HwDmUser /var/dmp && \
|
||||
chown HwHiAiUser:HwHiAiUser /var/driver && \
|
||||
chown HwHiAiUser:HwHiAiUser /usr/slog && \
|
||||
chown HwHiAiUser:HwHiAiUser /home/drv/hdc_ppc && \
|
||||
chown HwHiAiUser:HwHiAiUser /var/log/npu/slog
|
||||
|
||||
# Dynamically generate entrypoint.sh for dual mode compatibility
|
||||
RUN echo '#!/bin/bash' > /usr/local/bin/entrypoint.sh && \
|
||||
echo 'if lspci 2>/dev/null | grep -qi "accelerators"; then' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' :' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo 'else' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' su - HwHiAiUser -c "export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/:/usr/lib64 && /lib/ld-linux-aarch64.so.1 /var/slogd -d >/dev/null &"' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' su - HwDmUser -c "export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/:/usr/lib64 && /lib/ld-linux-aarch64.so.1 /var/dmp_daemon -I -M -U 8087 &"' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' ' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' rm -rf /lib64/ld-linux-aarch64.so.1' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo ' ln -sf /lib/ld-linux-aarch64.so.1 /lib64/ld-linux-aarch64.so.1' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo 'fi' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo '' >> /usr/local/bin/entrypoint.sh && \
|
||||
echo 'exec "$@"' >> /usr/local/bin/entrypoint.sh && \
|
||||
chmod +x /usr/local/bin/entrypoint.sh
|
||||
|
||||
# Set entrypoint and default command
|
||||
ENTRYPOINT ["/usr/local/bin/entrypoint.sh"]
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -15,46 +15,66 @@
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:8.2.rc1-a3-ubuntu22.04-py3.11
|
||||
FROM quay.io/ascend/cann:9.1.0-a3-ubuntu22.04-py3.12
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
|
||||
# Define environments
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV COMPILE_CUSTOM_KERNELS=${COMPILE_CUSTOM_KERNELS}
|
||||
|
||||
RUN apt-get update -y && \
|
||||
apt-get install -y python3-pip git vim wget net-tools gcc g++ cmake libnuma-dev && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
# Install clang-15 (for triton-ascend) and Mooncake
|
||||
ARG MOONCAKE_TAG=0.3.11.post1
|
||||
RUN apt-get update -y && \
|
||||
apt-get install -y git vim wget net-tools gcc g++ cmake numactl libnuma-dev libibverbs-dev libjemalloc2 libhiredis-dev clang-15 && \
|
||||
update-alternatives --install /usr/bin/clang clang /usr/bin/clang-15 20 && \
|
||||
update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-15 20 && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
python3 -m pip install --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple --extra-index-url https://pypi.org/simple mooncake-transfer-engine-npu==${MOONCAKE_TAG} && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL}
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL} && \
|
||||
python3 -m pip install modelscope 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.11.0rc3
|
||||
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
|
||||
ARG VLLM_TAG=v0.23.0
|
||||
ARG VLLM_COMMIT=""
|
||||
RUN if [ -n "$VLLM_COMMIT" ]; then \
|
||||
git init /vllm-workspace/vllm && \
|
||||
git -C /vllm-workspace/vllm fetch --depth 1 $VLLM_REPO "$VLLM_COMMIT" && \
|
||||
git -C /vllm-workspace/vllm checkout FETCH_HEAD; \
|
||||
else \
|
||||
git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm; \
|
||||
fi
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -v -e /vllm-workspace/vllm/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
ARG SOC_VERSION="ascend910_9391"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV SOC_VERSION=$SOC_VERSION \
|
||||
TASK_QUEUE_ENABLE=1 \
|
||||
OMP_NUM_THREADS=1
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" && \
|
||||
export VLLM_BATCH_INVARIANT=1 && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton triton-ascend && \
|
||||
python3 -m pip install triton-ascend==3.2.2 --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN python3 -m pip install modelscope 'ray>=2.47.1' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
|
||||
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
@@ -15,44 +15,61 @@
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:8.2.rc1-a3-openeuler24.03-py3.11
|
||||
FROM quay.io/ascend/cann:9.1.0-a3-openeuler24.03-py3.12
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
|
||||
ENV COMPILE_CUSTOM_KERNELS=${COMPILE_CUSTOM_KERNELS}
|
||||
|
||||
RUN yum update -y && \
|
||||
yum install -y python3-pip git vim wget net-tools gcc gcc-c++ make cmake numactl-devel && \
|
||||
rm -rf /var/cache/yum
|
||||
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL}
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
SHELL ["/bin/bash", "-c"]
|
||||
|
||||
# Install clang (for triton-ascend) and Mooncake
|
||||
ARG MOONCAKE_TAG=0.3.11.post1
|
||||
RUN yum update -y && \
|
||||
yum install -y git vim wget net-tools gcc gcc-c++ make cmake numactl numactl-devel libibverbs-devel jemalloc hiredis-devel clang patch && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
python3 -m pip install --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple --extra-index-url https://pypi.org/simple mooncake-transfer-engine-npu==${MOONCAKE_TAG} && \
|
||||
rm -rf /var/cache/yum/*
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL} && \
|
||||
python3 -m pip install modelscope 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.11.0rc3
|
||||
|
||||
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
|
||||
ARG VLLM_TAG=v0.23.0
|
||||
ARG VLLM_COMMIT=""
|
||||
RUN if [ -n "$VLLM_COMMIT" ]; then \
|
||||
git init /vllm-workspace/vllm && \
|
||||
git -C /vllm-workspace/vllm fetch --depth 1 $VLLM_REPO "$VLLM_COMMIT" && \
|
||||
git -C /vllm-workspace/vllm checkout FETCH_HEAD; \
|
||||
else \
|
||||
git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm; \
|
||||
fi
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
ARG SOC_VERSION="ascend910_9391"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
ENV SOC_VERSION=$SOC_VERSION \
|
||||
TASK_QUEUE_ENABLE=1 \
|
||||
OMP_NUM_THREADS=1
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" && \
|
||||
export VLLM_BATCH_INVARIANT=1 && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:/usr/include/c++/12:/usr/include/c++/12/`uname -i`-openEuler-linux && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton triton-ascend && \
|
||||
python3 -m pip install triton-ascend==3.2.2 --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN python3 -m pip install modelscope 'ray>=2.47.1' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
RUN echo "export LD_PRELOAD=/usr/lib64/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
|
||||
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
72
Dockerfile.a5
Normal file
72
Dockerfile.a5
Normal file
@@ -0,0 +1,72 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies 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.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:9.1.0-950-ubuntu22.04-py3.12
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Install clang-15 (for triton-ascend) and Mooncake
|
||||
ARG MOONCAKE_TAG=0.3.11.post1
|
||||
RUN apt-get update -y && \
|
||||
apt-get install -y git vim wget net-tools gcc g++ cmake numactl libnuma-dev libibverbs-dev libjemalloc2 libhiredis-dev clang-15 && \
|
||||
update-alternatives --install /usr/bin/clang clang /usr/bin/clang-15 20 && \
|
||||
update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-15 20 && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
python3 -m pip install --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple --extra-index-url https://pypi.org/simple mooncake-transfer-engine-npu==${MOONCAKE_TAG} && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL} && \
|
||||
python3 -m pip install modelscope 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.23.0
|
||||
RUN git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
ARG SOC_VERSION="ascend950dt_9582"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV SOC_VERSION=$SOC_VERSION \
|
||||
TASK_QUEUE_ENABLE=1 \
|
||||
OMP_NUM_THREADS=1
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton triton-ascend && \
|
||||
python3 -m pip install triton-ascend==3.2.2 --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
|
||||
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
67
Dockerfile.a5.openEuler
Normal file
67
Dockerfile.a5.openEuler
Normal file
@@ -0,0 +1,67 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies 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.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:9.1.0-950-openeuler24.03-py3.12
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
SHELL ["/bin/bash", "-c"]
|
||||
|
||||
# Install clang (for triton-ascend) and Mooncake
|
||||
ARG MOONCAKE_TAG=0.3.11.post1
|
||||
RUN yum update -y && \
|
||||
yum install -y git vim wget net-tools gcc gcc-c++ make cmake numactl numactl-devel libibverbs-devel jemalloc hiredis-devel clang patch && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
python3 -m pip install --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple --extra-index-url https://pypi.org/simple mooncake-transfer-engine-npu==${MOONCAKE_TAG} && \
|
||||
rm -rf /var/cache/yum/*
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL} && \
|
||||
python3 -m pip install modelscope 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.23.0
|
||||
RUN git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
ARG SOC_VERSION="ascend950dt_9582"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
ENV SOC_VERSION=$SOC_VERSION \
|
||||
TASK_QUEUE_ENABLE=1 \
|
||||
OMP_NUM_THREADS=1
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton triton-ascend && \
|
||||
python3 -m pip install triton-ascend==3.2.2 --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
RUN echo "export LD_PRELOAD=/usr/lib64/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
|
||||
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
@@ -15,44 +15,61 @@
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:8.2.rc1-910b-openeuler24.03-py3.11
|
||||
FROM quay.io/ascend/cann:9.1.0-910b-openeuler24.03-py3.12
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
|
||||
ENV COMPILE_CUSTOM_KERNELS=${COMPILE_CUSTOM_KERNELS}
|
||||
|
||||
RUN yum update -y && \
|
||||
yum install -y python3-pip git vim wget net-tools gcc gcc-c++ make cmake numactl-devel && \
|
||||
rm -rf /var/cache/yum
|
||||
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL}
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
SHELL ["/bin/bash", "-c"]
|
||||
|
||||
# Install clang (for triton-ascend) and Mooncake
|
||||
ARG MOONCAKE_TAG=0.3.11.post1
|
||||
RUN yum update -y && \
|
||||
yum install -y git vim wget net-tools gcc gcc-c++ make cmake numactl numactl-devel libibverbs-devel jemalloc hiredis-devel clang patch && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
python3 -m pip install --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple --extra-index-url https://pypi.org/simple mooncake-transfer-engine-npu==${MOONCAKE_TAG} && \
|
||||
rm -rf /var/cache/yum/*
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL} && \
|
||||
python3 -m pip install modelscope 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.11.0rc3
|
||||
|
||||
RUN git clone --depth 1 $VLLM_REPO --branch $VLLM_TAG /vllm-workspace/vllm
|
||||
ARG VLLM_TAG=v0.23.0
|
||||
ARG VLLM_COMMIT=""
|
||||
RUN if [ -n "$VLLM_COMMIT" ]; then \
|
||||
git init /vllm-workspace/vllm && \
|
||||
git -C /vllm-workspace/vllm fetch --depth 1 $VLLM_REPO "$VLLM_COMMIT" && \
|
||||
git -C /vllm-workspace/vllm checkout FETCH_HEAD; \
|
||||
else \
|
||||
git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm; \
|
||||
fi
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
ARG SOC_VERSION="ascend910b1"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
ENV SOC_VERSION=$SOC_VERSION \
|
||||
TASK_QUEUE_ENABLE=1 \
|
||||
OMP_NUM_THREADS=1
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" && \
|
||||
export VLLM_BATCH_INVARIANT=1 && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
export CPLUS_INCLUDE_PATH=$CPLUS_INCLUDE_PATH:/usr/include/c++/12:/usr/include/c++/12/`uname -i`-openEuler-linux && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton triton-ascend && \
|
||||
python3 -m pip install triton-ascend==3.2.2 --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN python3 -m pip install modelscope 'ray>=2.47.1' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
RUN echo "export LD_PRELOAD=/usr/lib64/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
|
||||
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
|
||||
78
Dockerfile.origin
Normal file
78
Dockerfile.origin
Normal file
@@ -0,0 +1,78 @@
|
||||
#
|
||||
# Copyright (c) 2025 Huawei Technologies 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.
|
||||
# This file is a part of the vllm-ascend project.
|
||||
#
|
||||
|
||||
FROM quay.io/ascend/cann:9.1.0-910b-ubuntu22.04-py3.12
|
||||
|
||||
ARG PIP_INDEX_URL="https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple"
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Install clang-15 (for triton-ascend) and Mooncake
|
||||
ARG MOONCAKE_TAG=0.3.11.post1
|
||||
RUN apt-get update -y && \
|
||||
apt-get install -y git vim wget net-tools gcc g++ cmake numactl libnuma-dev libibverbs-dev libjemalloc2 libhiredis-dev clang-15 && \
|
||||
update-alternatives --install /usr/bin/clang clang /usr/bin/clang-15 20 && \
|
||||
update-alternatives --install /usr/bin/clang++ clang++ /usr/bin/clang++-15 20 && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
python3 -m pip install --index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple --extra-index-url https://pypi.org/simple mooncake-transfer-engine-npu==${MOONCAKE_TAG} && \
|
||||
rm -rf /var/cache/apt/* && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install modelscope (for fast download) and ray (for multinode)
|
||||
RUN pip config set global.index-url ${PIP_INDEX_URL} && \
|
||||
python3 -m pip install modelscope 'ray>=2.47.1,<=2.48.0' 'protobuf>3.20.0' && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vLLM
|
||||
ARG VLLM_REPO=https://github.com/vllm-project/vllm.git
|
||||
ARG VLLM_TAG=v0.23.0
|
||||
ARG VLLM_COMMIT=""
|
||||
RUN if [ -n "$VLLM_COMMIT" ]; then \
|
||||
git init /vllm-workspace/vllm && \
|
||||
git -C /vllm-workspace/vllm fetch --depth 1 $VLLM_REPO "$VLLM_COMMIT" && \
|
||||
git -C /vllm-workspace/vllm checkout FETCH_HEAD; \
|
||||
else \
|
||||
git clone --depth 1 -b $VLLM_TAG $VLLM_REPO /vllm-workspace/vllm; \
|
||||
fi
|
||||
# In x86, triton will be installed by vllm. But in Ascend, triton doesn't work correctly. we need to uninstall it.
|
||||
RUN VLLM_TARGET_DEVICE="empty" python3 -m pip install -e /vllm-workspace/vllm/[audio] --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Install vllm-ascend
|
||||
ARG SOC_VERSION="ascend910b1"
|
||||
ARG COMPILE_CUSTOM_KERNELS=1
|
||||
ENV DEBIAN_FRONTEND=noninteractive
|
||||
ENV SOC_VERSION=$SOC_VERSION \
|
||||
TASK_QUEUE_ENABLE=1 \
|
||||
OMP_NUM_THREADS=1
|
||||
COPY . /vllm-workspace/vllm-ascend/
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL="https://mirrors.huaweicloud.com/ascend/repos/pypi" && \
|
||||
export VLLM_BATCH_INVARIANT=1 && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
python3 -m pip install -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip uninstall -y triton triton-ascend && \
|
||||
python3 -m pip install triton-ascend==3.2.2 --extra-index-url https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
|
||||
RUN echo "export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib" >> ~/.bashrc
|
||||
|
||||
CMD ["/bin/bash"]
|
||||
118
README.origin.md
Normal file
118
README.origin.md
Normal file
@@ -0,0 +1,118 @@
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-ascend/main/docs/source/logos/vllm-ascend-logo-text-dark.png">
|
||||
<img alt="vllm-ascend" src="https://raw.githubusercontent.com/vllm-project/vllm-ascend/main/docs/source/logos/vllm-ascend-logo-text-light.png" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
<h3 align="center">
|
||||
vLLM Ascend Plugin
|
||||
</h3>
|
||||
|
||||
<div align="center">
|
||||
|
||||
[](https://deepwiki.com/vllm-project/vllm-ascend)
|
||||
|
||||
</div>
|
||||
|
||||
<p align="center">
|
||||
| <a href="https://www.hiascend.com/en/"><b>About Ascend</b></a> | <a href="https://docs.vllm.ai/projects/ascend/en/latest/"><b>Documentation</b></a> | <a href="https://slack.vllm.ai"><b>#SIG-Ascend</b></a> | <a href="https://discuss.vllm.ai/c/hardware-support/vllm-ascend-support"><b>Users Forum</b></a> | <a href="https://tinyurl.com/vllm-ascend-meeting"><b>Weekly Meeting</b></a> |
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a ><b>English</b></a> | <a href="README.zh.md"><b>中文</b></a>
|
||||
</p>
|
||||
|
||||
---
|
||||
*Latest News* 🔥
|
||||
|
||||
- [2026/07] We released the new official version [v0.23.0](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.23.0)! Please follow the [official guide](https://docs.vllm.ai/projects/ascend/en/v0.23.0/) to start using vLLM Ascend Plugin on Ascend.
|
||||
- [2026/05] We released the new official version [v0.18.0](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.18.0)! Please follow the [official guide](https://docs.vllm.ai/projects/ascend/en/v0.18.0/) to start using vLLM Ascend Plugin on Ascend.
|
||||
- [2026/02] We released the new official version [v0.13.0](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.13.0)! Please follow the [official guide](https://docs.vllm.ai/projects/ascend/en/v0.13.0/) to start using vLLM Ascend Plugin on Ascend.
|
||||
|
||||
<details>
|
||||
<summary>More</summary>
|
||||
|
||||
- [2025/12] We released the new official version [v0.11.0](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.11.0)! Please follow the [official guide](https://docs.vllm.ai/projects/ascend/en/v0.11.0/) to start using vLLM Ascend Plugin on Ascend.
|
||||
- [2025/09] We released the new official version [v0.9.1](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.9.1)! Please follow the [official guide](https://docs.vllm.ai/projects/ascend/en/v0.9.1/tutorials/large_scale_ep.html) to start deploying large-scale Expert Parallelism (EP) on Ascend.
|
||||
- [2025/08] We hosted the [vLLM Beijing Meetup](https://mp.weixin.qq.com/s/7n8OYNrCC_I9SJaybHA_-Q) with vLLM and Tencent! Please find the [meetup slides](https://drive.google.com/drive/folders/1Pid6NSFLU43DZRi0EaTcPgXsAzDvbBqF).
|
||||
- [2025/06] [User stories](https://docs.vllm.ai/projects/ascend/en/latest/community/user_stories/index.html) page is now live! It kicks off with LLaMA-Factory/verl/TRL/GPUStack to demonstrate how vLLM Ascend assists Ascend users in enhancing their experience across fine-tuning, evaluation, reinforcement learning (RL), and deployment scenarios.
|
||||
- [2025/06] [Contributors](https://docs.vllm.ai/projects/ascend/en/latest/community/contributors.html) page is now live! All contributions deserve to be recorded, thanks for all contributors.
|
||||
- [2025/05] We've released the first official version [v0.7.3](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.7.3)! We collaborated with the vLLM community to publish a blog post sharing our practice: [Introducing vLLM Hardware Plugin, Best Practice from Ascend NPU](https://blog.vllm.ai/2025/05/12/hardware-plugin.html).
|
||||
- [2025/03] We hosted the [vLLM Beijing Meetup](https://mp.weixin.qq.com/s/VtxO9WXa5fC-mKqlxNUJUQ) with vLLM team! Please find the [meetup slides](https://drive.google.com/drive/folders/1Pid6NSFLU43DZRi0EaTcPgXsAzDvbBqF).
|
||||
- [2025/02] vLLM community officially created [vllm-project/vllm-ascend](https://github.com/vllm-project/vllm-ascend) repo for running vLLM seamlessly on the Ascend NPU.
|
||||
- [2024/12] We are working with the vLLM community to support [[RFC]: Hardware pluggable](https://github.com/vllm-project/vllm/issues/11162).
|
||||
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
vLLM Ascend (`vllm-ascend`) is a community maintained hardware plugin for running vLLM seamlessly on the Ascend NPU.
|
||||
|
||||
It is the recommended approach for supporting the Ascend backend within the vLLM community. It adheres to the principles outlined in the [[RFC]: Hardware pluggable](https://github.com/vllm-project/vllm/issues/11162), providing a hardware-pluggable interface that decouples the integration of the Ascend NPU with vLLM.
|
||||
|
||||
By using vLLM Ascend plugin, popular open-source models, including Transformer-like, Mixture-of-Experts (MoE), Embedding, Multi-modal LLMs can run seamlessly on the Ascend NPU.
|
||||
|
||||
For detailed information on supported models, please refer to [supported models](https://docs.vllm.ai/projects/ascend/en/latest/user_guide/support_matrix/supported_models.html).
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Hardware: Atlas 800I A2 Inference series, Atlas A2 Training series, Atlas 800I A3 Inference series, Atlas A3 Training series, Atlas 300I Duo (Experimental)
|
||||
- OS: Linux
|
||||
- Software:
|
||||
- Python >= 3.10, < 3.13
|
||||
- CANN == 9.1.0 (For Ascend HDK version, please refer to the [CANN 9.1.0 Release Notes](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/910/softwareinst/releasenote/9.1.0/release-notes.md))
|
||||
- PyTorch == 2.10.0, TorchNPU == 2.10.0.post4
|
||||
- vLLM (the same version as vllm-ascend)
|
||||
|
||||
## Getting Started
|
||||
|
||||
Please use the following recommended versions to get started quickly:
|
||||
|
||||
| Version | Release type | Doc |
|
||||
|------------|--------------|--------------------------------------|
|
||||
| v0.23.0 | Latest stable version | See [QuickStart](https://docs.vllm.ai/projects/ascend/en/v0.23.0/quick_start.html) and [Installation](https://docs.vllm.ai/projects/ascend/en/v0.23.0/installation.html) for more details |
|
||||
|
||||
## Branch
|
||||
|
||||
vllm-ascend has a main branch and a dev branch.
|
||||
|
||||
- **main**: main branch, corresponds to the vLLM main branch, and is continuously monitored for quality through Ascend CI.
|
||||
- **releases/vX.Y.Z**: development branch, created alongside new releases of vLLM. For example, `releases/v0.13.0` is the dev branch for vLLM `v0.13.0` version.
|
||||
|
||||
Below are the maintained branches:
|
||||
|
||||
| Branch | Status | Note |
|
||||
|------------------|--------------|--------------------------------------|
|
||||
| main | Maintained | CI commitment for vLLM main branch and vLLM v0.23.0 tag |
|
||||
| v0.7.1-dev | Unmaintained | Outdated, no longer maintained. |
|
||||
| v0.7.3-dev | Unmaintained | Only bug fixes are allowed, and no new release tags anymore. |
|
||||
| v0.9.1-dev | Unmaintained | Only bug fixes are allowed, and no new release tags anymore. |
|
||||
| v0.11.0-dev | Unmaintained | Only bug fixes are allowed, and no new release tags anymore. |
|
||||
| releases/v0.13.0 | Maintained | CI commitment for vLLM 0.13.0 version |
|
||||
| releases/v0.18.0 | Maintained | CI commitment for vLLM 0.18.0 version |
|
||||
| releases/v0.20.2rc | Maintained | CI commitment for vLLM 0.20.2 version |
|
||||
| rfc/feature-name | Maintained | [Feature branches](https://docs.vllm.ai/projects/ascend/en/latest/community/versioning_policy.html#feature-branches) for collaboration |
|
||||
| releases/v0.23.0 | Maintained | CI commitment for vLLM 0.23.0 version |
|
||||
|
||||
Please refer to [Versioning policy](https://docs.vllm.ai/projects/ascend/en/latest/community/versioning_policy.html) for more details.
|
||||
|
||||
## Contributing
|
||||
|
||||
See [CONTRIBUTING](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/contribution/index.html) for more details, which is a step-by-step guide to help you set up the development environment, build and test.
|
||||
|
||||
We welcome and value any contributions and collaborations:
|
||||
|
||||
- Please let us know if you encounter a bug by [filing an issue](https://github.com/vllm-project/vllm-ascend/issues)
|
||||
- Please use [User forum](https://discuss.vllm.ai/c/hardware-support/vllm-ascend-support) for usage questions and help.
|
||||
|
||||
## Weekly Meeting
|
||||
|
||||
- vLLM Ascend Weekly Meeting: <https://tinyurl.com/vllm-ascend-meeting>
|
||||
- Wednesday, 15:00 - 16:00 (UTC+8, [Convert to your timezone](https://dateful.com/convert/gmt8?t=15))
|
||||
|
||||
## License
|
||||
|
||||
Apache License 2.0, as found in the [LICENSE](./LICENSE) file.
|
||||
112
README.zh.md
Normal file
112
README.zh.md
Normal file
@@ -0,0 +1,112 @@
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-ascend/main/docs/source/logos/vllm-ascend-logo-text-dark.png">
|
||||
<img alt="vllm-ascend" src="https://raw.githubusercontent.com/vllm-project/vllm-ascend/main/docs/source/logos/vllm-ascend-logo-text-light.png" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
<h3 align="center">
|
||||
vLLM Ascend Plugin
|
||||
</h3>
|
||||
|
||||
<p align="center">
|
||||
| <a href="https://www.hiascend.com/en/"><b>关于昇腾</b></a> | <a href="https://docs.vllm.ai/projects/ascend/en/latest/"><b>官方文档</b></a> | <a href="https://slack.vllm.ai"><b>#sig-ascend</b></a> | <a href="https://discuss.vllm.ai/c/hardware-support/vllm-ascend-support"><b>用户论坛</b></a> | <a href="https://tinyurl.com/vllm-ascend-meeting"><b>社区例会</b></a> |
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="README.md"><b>English</b></a> | <a><b>中文</b></a>
|
||||
</p>
|
||||
|
||||
---
|
||||
*最新消息* 🔥
|
||||
|
||||
- [2026/07] 我们发布了新的正式版本 [v0.23.0](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.23.0)! 请按照[官方指南](https://docs.vllm.ai/projects/ascend/en/v0.23.0/)开始在 Ascend 上部署 vLLM Ascend Plugin。
|
||||
- [2026/05] 我们发布了新的正式版本 [v0.18.0](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.18.0)! 请按照[官方指南](https://docs.vllm.ai/projects/ascend/en/v0.18.0/)开始在Ascend上部署vLLM Ascend Plugin。
|
||||
- [2026/02] 我们发布了新的正式版本 [v0.13.0](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.13.0)! 请按照[官方指南](https://docs.vllm.ai/projects/ascend/en/v0.13.0/)开始在Ascend上部署vLLM Ascend Plugin。
|
||||
|
||||
<details>
|
||||
<summary>更多内容</summary>
|
||||
|
||||
- [2025/12] 我们发布了新的正式版本 [v0.11.0](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.11.0)! 请按照[官方指南](https://docs.vllm.ai/projects/ascend/en/v0.11.0/)开始在Ascend上部署vLLM Ascend Plugin。
|
||||
- [2025/09] 我们发布了新的正式版本 [v0.9.1](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.9.1)! 请按照[官方指南](https://docs.vllm.ai/projects/ascend/en/v0.9.1/tutorials/large_scale_ep.html)开始在Ascend上部署大型专家并行 (EP)。
|
||||
- [2025/08] 我们与vLLM和腾讯合作举办了[vLLM北京Meetup](https://mp.weixin.qq.com/s/7n8OYNrCC_I9SJaybHA_-Q),!请查阅[活动幻灯片](https://drive.google.com/drive/folders/1Pid6NSFLU43DZRi0EaTcPgXsAzDvbBqF)。
|
||||
- [2025/06] [用户案例](https://docs.vllm.ai/projects/ascend/en/latest/community/user_stories/index.html)现已上线!展示了LLaMA-Factory/verl/TRL/GPUStack等用户案例,展示了vLLM Ascend如何帮助昇腾用户在模型微调、评估、强化学习 (RL) 以及部署等场景中提升体验。
|
||||
- [2025/06] [贡献者](https://docs.vllm.ai/projects/ascend/en/latest/community/contributors.html)页面现已上线!所有的贡献都值得被记录,感谢所有的贡献者。
|
||||
- [2025/05] 我们发布了首个正式版本 [v0.7.3](https://github.com/vllm-project/vllm-ascend/releases/tag/v0.7.3)!我们与 vLLM 社区合作发布了一篇博客文章,分享了我们的实践:[Introducing vLLM Hardware Plugin, Best Practice from Ascend NPU](https://blog.vllm.ai/2025/05/12/hardware-plugin.html)。
|
||||
- [2025/03] 我们和vLLM团队举办了[vLLM Beijing Meetup](https://mp.weixin.qq.com/s/CGDuMoB301Uytnrkc2oyjg)! 请查阅[活动幻灯片](https://drive.google.com/drive/folders/1Pid6NSFLU43DZRi0EaTcPgXsAzDvbBqF).
|
||||
- [2025/02] vLLM社区正式创建了[vllm-project/vllm-ascend](https://github.com/vllm-project/vllm-ascend)仓库,让vLLM可以无缝运行在Ascend NPU。
|
||||
- [2024/12] 我们正在与 vLLM 社区合作,以支持 [[RFC]: Hardware pluggable](https://github.com/vllm-project/vllm/issues/11162).
|
||||
|
||||
</details>
|
||||
|
||||
---
|
||||
|
||||
## 总览
|
||||
|
||||
vLLM 昇腾插件 (`vllm-ascend`) 是一个由社区维护的让vLLM在Ascend NPU无缝运行的后端插件。
|
||||
|
||||
此插件是 vLLM 社区中支持昇腾后端的推荐方式。它遵循[[RFC]: Hardware pluggable](https://github.com/vllm-project/vllm/issues/11162)所述原则:通过解耦的方式提供了vLLM对Ascend NPU的支持。
|
||||
|
||||
使用 vLLM 昇腾插件,可以让类Transformer、混合专家(MOE)、嵌入、多模态等流行的大语言模型在 Ascend NPU 上无缝运行。
|
||||
|
||||
支持的模型详细信息,请参考[模型支持列表](https://docs.vllm.ai/projects/ascend/en/latest/user_guide/support_matrix/supported_models.html)。
|
||||
|
||||
## 准备
|
||||
|
||||
- 硬件:Atlas 800I A2 Inference系列、Atlas A2 Training系列、Atlas 800I A3 Inference系列、Atlas A3 Training系列、Atlas 300I Duo(实验性支持)
|
||||
- 操作系统:Linux
|
||||
- 软件:
|
||||
- Python >= 3.10, < 3.13
|
||||
- CANN == 9.1.0 (Ascend HDK 版本详见 [CANN 9.1.0 版本说明](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/910/softwareinst/releasenote/9.1.0/release-notes.md))
|
||||
- PyTorch == 2.10.0, TorchNPU == 2.10.0.post4
|
||||
- vLLM (与vllm-ascend版本一致)
|
||||
|
||||
## 开始使用
|
||||
|
||||
推荐您使用以下版本快速开始使用:
|
||||
|
||||
| Version | Release type | Doc |
|
||||
|------------|--------------|--------------------------------------|
|
||||
| v0.23.0 | 最新正式/稳定版本 | 请查看[快速开始](https://docs.vllm.ai/projects/ascend/en/v0.23.0/quick_start.html)和[安装指南](https://docs.vllm.ai/projects/ascend/en/v0.23.0/installation.html)了解更多 |
|
||||
|
||||
## 分支策略
|
||||
|
||||
vllm-ascend有主干分支和开发分支。
|
||||
|
||||
- **main**: 主干分支,与vLLM的主干分支对应,并通过昇腾CI持续进行质量看护。
|
||||
- **releases/vX.Y.Z**: 开发分支,随vLLM部分新版本发布而创建,比如`releases/v0.13.0`是vllm-ascend针对vLLM `v0.13.0` 版本的开发分支。
|
||||
|
||||
下面是维护中的分支:
|
||||
|
||||
| 分支 | 状态 | 备注 |
|
||||
|------------------|--------------|----------------------|
|
||||
| main | Maintained | 基于vLLM main分支和vLLM最新版本(v0.23.0)CI看护 |
|
||||
| v0.7.1-dev | Unmaintained | 不再维护 |
|
||||
| v0.7.3-dev | Unmaintained | 只允许Bug修复,不会再发布新版本 |
|
||||
| v0.9.1-dev | Unmaintained | 只允许Bug修复,不会再发布新版本 |
|
||||
| v0.11.0-dev | Unmaintained | 只允许Bug修复,不会再发布新版本 |
|
||||
| releases/v0.13.0 | Maintained | 基于vLLM v0.13.0版本CI看护 |
|
||||
| releases/v0.18.0 | Maintained | 基于vLLM v0.18.0版本CI看护 |
|
||||
| releases/v0.20.2rc | Maintained | 基于vLLM v0.20.2版本CI看护 |
|
||||
| rfc/feature-name | Maintained | 为协作创建的[特性分支](https://docs.vllm.ai/projects/ascend/en/latest/community/versioning_policy.html#feature-branches) |
|
||||
| releases/v0.23.0 | Maintained | 基于vLLM v0.23.0版本CI看护 |
|
||||
|
||||
请参阅[版本策略](https://docs.vllm.ai/projects/ascend/en/latest/community/versioning_policy.html)了解更多详细信息。
|
||||
|
||||
## 贡献
|
||||
|
||||
请参考[CONTRIBUTING](https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/contribution/index.html)文档了解更多关于开发环境搭建、功能测试以及 PR 提交规范的信息。
|
||||
|
||||
我们欢迎并重视任何形式的贡献与合作:
|
||||
|
||||
- 请通过[Issue](https://github.com/vllm-project/vllm-ascend/issues)来告知我们您遇到的任何Bug。
|
||||
- 请通过[用户论坛](https://discuss.vllm.ai/c/hardware-support/vllm-ascend-support)来交流使用问题和寻求帮助。
|
||||
|
||||
## 社区例会
|
||||
|
||||
- vLLM Ascend 每周社区例会: <https://tinyurl.com/vllm-ascend-meeting>
|
||||
- 每周三下午,15:00 - 16:00 (UTC+8, [查看您的时区](https://dateful.com/convert/gmt8?t=15))
|
||||
|
||||
## 许可证
|
||||
|
||||
Apache 许可证 2.0,如 [LICENSE](./LICENSE) 文件中所示。
|
||||
@@ -1,8 +1,13 @@
|
||||
# Introduction
|
||||
# vLLM Ascend Benchmarks
|
||||
|
||||
## Introduction
|
||||
|
||||
This document outlines the benchmarking methodology for vllm-ascend, aimed at evaluating the performance under a variety of workloads. The primary goal is to help developers assess whether their pull requests improve or degrade vllm-ascend's performance.
|
||||
|
||||
# Overview
|
||||
## Overview
|
||||
|
||||
**Benchmarking Coverage**: We measure latency, throughput, and fixed-QPS serving on the Atlas800I A2 (see [quick_start](../docs/source/quick_start.md) to learn more supported devices list), with different models(coming soon).
|
||||
|
||||
- Latency tests
|
||||
- Input length: 32 tokens.
|
||||
- Output length: 128 tokens.
|
||||
@@ -24,10 +29,12 @@ This document outlines the benchmarking methodology for vllm-ascend, aimed at ev
|
||||
- Models: Qwen2.5-VL-7B-Instruct, Qwen2.5-7B-Instruct, Qwen3-8B.
|
||||
- Evaluation metrics: throughput, TTFT (time to the first token, with mean, median and p99), ITL (inter-token latency, with mean, median and p99).
|
||||
|
||||
**Benchmarking Duration**: about 800 senond for single model.
|
||||
**Benchmarking Duration**: about 800 seconds for single model.
|
||||
|
||||
## Quick Use
|
||||
|
||||
### Prerequisites
|
||||
|
||||
# Quick Use
|
||||
## Prerequisites
|
||||
Before running the benchmarks, ensure the following:
|
||||
|
||||
- vllm and vllm-ascend are installed and properly set up in an NPU environment, as these scripts are specifically designed for NPU devices.
|
||||
@@ -39,9 +46,9 @@ Before running the benchmarks, ensure the following:
|
||||
```
|
||||
|
||||
- For performance benchmark, it is recommended to set the [load-format](https://github.com/vllm-project/vllm-ascend/blob/5897dc5bbe321ca90c26225d0d70bff24061d04b/benchmarks/tests/latency-tests.json#L7) as `dummy`, It will construct random weights based on the passed model without downloading the weights from internet, which can greatly reduce the benchmark time.
|
||||
- If you want to run benchmark customized, feel free to add your own models and parameters in the [JSON](https://github.com/vllm-project/vllm-ascend/tree/main/benchmarks/tests), let's take `Qwen2.5-VL-7B-Instruct`as an example:
|
||||
- If you want to run a customized benchmark, feel free to add your own models and parameters in the [JSON](https://github.com/vllm-project/vllm-ascend/tree/main/benchmarks/tests), let's take `Qwen2.5-VL-7B-Instruct`as an example:
|
||||
|
||||
```shell
|
||||
```json
|
||||
[
|
||||
{
|
||||
"test_name": "serving_qwen2_5vl_7B_tp1",
|
||||
@@ -75,45 +82,46 @@ Before running the benchmarks, ensure the following:
|
||||
|
||||
this Json will be structured and parsed into server parameters and client parameters by the benchmark script. This configuration defines a test case named `serving_qwen2_5vl_7B_tp1`, designed to evaluate the performance of the `Qwen/Qwen2.5-VL-7B-Instruct` model under different request rates. The test includes both server and client parameters, for more parameters details, see vllm benchmark [cli](https://github.com/vllm-project/vllm/tree/main/vllm/benchmarks).
|
||||
|
||||
- **Test Overview**
|
||||
- Test Name: serving_qwen2_5vl_7B_tp1
|
||||
- **Test Overview**
|
||||
- Test Name: serving_qwen2_5vl_7B_tp1
|
||||
|
||||
- Queries Per Second (QPS): The test is run at four different QPS levels: 1, 4, 16, and inf (infinite load, typically used for stress testing).
|
||||
- Queries Per Second (QPS): The test is run at four different QPS levels: 1, 4, 16, and inf (infinite load, typically used for stress testing).
|
||||
|
||||
- Server Parameters
|
||||
- Model: Qwen/Qwen2.5-VL-7B-Instruct
|
||||
- Server Parameters
|
||||
- Model: Qwen/Qwen2.5-VL-7B-Instruct
|
||||
|
||||
- Tensor Parallelism: 1 (no model parallelism is used; the model runs on a single device or node)
|
||||
- Tensor Parallelism: 1 (no model parallelism is used; the model runs on a single device or node)
|
||||
|
||||
- Swap Space: 16 GB (used to handle memory overflow by swapping to disk)
|
||||
- Swap Space: 16 GB (used to handle memory overflow by swapping to disk)
|
||||
|
||||
- disable_log_stats: disables logging of performance statistics.
|
||||
- disable_log_stats: disables logging of performance statistics.
|
||||
|
||||
- disable_log_requests: disables logging of individual requests.
|
||||
- disable_log_requests: disables logging of individual requests.
|
||||
|
||||
- Trust Remote Code: enabled (allows execution of model-specific custom code)
|
||||
- Trust Remote Code: enabled (allows execution of model-specific custom code)
|
||||
|
||||
- Max Model Length: 16,384 tokens (maximum context length supported by the model)
|
||||
- Max Model Length: 16,384 tokens (maximum context length supported by the model)
|
||||
|
||||
- Client Parameters
|
||||
- Client Parameters
|
||||
|
||||
- Model: Qwen/Qwen2.5-VL-7B-Instruct (same as the server)
|
||||
- Model: Qwen/Qwen2.5-VL-7B-Instruct (same as the server)
|
||||
|
||||
- Backend: openai-chat (suggests the client uses the OpenAI-compatible chat API format)
|
||||
- Backend: openai-chat (suggests the client uses the OpenAI-compatible chat API format)
|
||||
|
||||
- Dataset Source: Hugging Face (hf)
|
||||
- Dataset Source: Hugging Face (hf)
|
||||
|
||||
- Dataset Split: train
|
||||
- Dataset Split: train
|
||||
|
||||
- Endpoint: /v1/chat/completions (the REST API endpoint to which chat requests are sent)
|
||||
- Endpoint: /v1/chat/completions (the REST API endpoint to which chat requests are sent)
|
||||
|
||||
- Dataset Path: lmarena-ai/vision-arena-bench-v0.1 (the benchmark dataset used for evaluation, hosted on Hugging Face)
|
||||
- Dataset Path: lmarena-ai/vision-arena-bench-v0.1 (the benchmark dataset used for evaluation, hosted on Hugging Face)
|
||||
|
||||
- Number of Prompts: 200 (the total number of prompts used during the test)
|
||||
- Number of Prompts: 200 (the total number of prompts used during the test)
|
||||
|
||||
## Run benchmarks
|
||||
### Run benchmarks
|
||||
|
||||
#### Use benchmark script
|
||||
|
||||
### Use benchmark script
|
||||
The provided scripts automatically execute performance tests for serving, throughput, and latency. To start the benchmarking process, run command in the vllm-ascend root directory:
|
||||
|
||||
```shell
|
||||
@@ -124,21 +132,22 @@ Once the script completes, you can find the results in the benchmarks/results fo
|
||||
|
||||
```shell
|
||||
.
|
||||
|-- serving_qwen2_5_7B_tp1_qps_1.json
|
||||
|-- serving_qwen2_5_7B_tp1_qps_16.json
|
||||
|-- serving_qwen2_5_7B_tp1_qps_4.json
|
||||
|-- serving_qwen2_5_7B_tp1_qps_inf.json
|
||||
|-- latency_qwen2_5_7B_tp1.json
|
||||
|-- throughput_qwen2_5_7B_tp1.json
|
||||
|-- serving_qwen2_5_7Bvl_tp1_qps_1.json
|
||||
|-- serving_qwen2_5_7Bvl_tp1_qps_16.json
|
||||
|-- serving_qwen2_5_7Bvl_tp1_qps_4.json
|
||||
|-- serving_qwen2_5_7Bvl_tp1_qps_inf.json
|
||||
|-- throughput_qwen2_5_7Bvl_tp1.json
|
||||
```
|
||||
|
||||
These files contain detailed benchmarking results for further analysis.
|
||||
|
||||
### Use benchmark cli
|
||||
#### Use benchmark cli
|
||||
|
||||
For more flexible and customized use, benchmark cli is also provided to run online/offline benchmarks
|
||||
Similarly, let’s take `Qwen2.5-VL-7B-Instruct` benchmark as an example:
|
||||
#### Online serving
|
||||
Similarly, let's take `Qwen2.5-VL-7B-Instruct` benchmark as an example:
|
||||
|
||||
##### Online serving
|
||||
|
||||
1. Launch the server:
|
||||
|
||||
```shell
|
||||
@@ -156,7 +165,8 @@ Similarly, let’s take `Qwen2.5-VL-7B-Instruct` benchmark as an example:
|
||||
--request-rate 16
|
||||
```
|
||||
|
||||
#### Offline
|
||||
##### Offline
|
||||
|
||||
- **Throughput**
|
||||
|
||||
```shell
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
from typing import Tuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
@@ -47,20 +45,12 @@ def get_masked_input_and_mask_ref(
|
||||
num_org_vocab_padding: int,
|
||||
added_vocab_start_index: int,
|
||||
added_vocab_end_index: int,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Reference implementation for verification"""
|
||||
org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ < org_vocab_end_index)
|
||||
added_vocab_mask = (input_ >= added_vocab_start_index) & (
|
||||
input_ < added_vocab_end_index
|
||||
)
|
||||
added_offset = (
|
||||
added_vocab_start_index
|
||||
- (org_vocab_end_index - org_vocab_start_index)
|
||||
- num_org_vocab_padding
|
||||
)
|
||||
valid_offset = (org_vocab_start_index * org_vocab_mask) + (
|
||||
added_offset * added_vocab_mask
|
||||
)
|
||||
added_vocab_mask = (input_ >= added_vocab_start_index) & (input_ < added_vocab_end_index)
|
||||
added_offset = added_vocab_start_index - (org_vocab_end_index - org_vocab_start_index) - num_org_vocab_padding
|
||||
valid_offset = (org_vocab_start_index * org_vocab_mask) + (added_offset * added_vocab_mask)
|
||||
vocab_mask = org_vocab_mask | added_vocab_mask
|
||||
masked_input = vocab_mask * (input_ - valid_offset)
|
||||
return masked_input, ~vocab_mask
|
||||
@@ -78,7 +68,7 @@ SEEDS = [0]
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_get_masked_input_and_mask(
|
||||
shape: Tuple[int, ...],
|
||||
shape: tuple[int, ...],
|
||||
dtype: torch.dtype,
|
||||
device: str,
|
||||
seed: int,
|
||||
|
||||
@@ -59,9 +59,7 @@ def results_to_json(latency, throughput, serving):
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Process the results of the benchmark tests."
|
||||
)
|
||||
parser = argparse.ArgumentParser(description="Process the results of the benchmark tests.")
|
||||
parser.add_argument(
|
||||
"--results_folder",
|
||||
type=str,
|
||||
@@ -80,12 +78,8 @@ if __name__ == "__main__":
|
||||
default="./perf_result_template.md",
|
||||
help="The template file for the markdown report.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tag", default="main", help="Tag to be used for release message."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--commit_id", default="", help="Commit ID to be used for release message."
|
||||
)
|
||||
parser.add_argument("--tag", default="main", help="Tag to be used for release message.")
|
||||
parser.add_argument("--commit_id", default="", help="Commit ID to be used for release message.")
|
||||
|
||||
args = parser.parse_args()
|
||||
results_folder = (CUR_PATH / args.results_folder).resolve()
|
||||
@@ -116,9 +110,7 @@ if __name__ == "__main__":
|
||||
# get different percentiles
|
||||
for perc in [10, 25, 50, 75, 90, 99]:
|
||||
# Multiply 1000 to convert the time unit from s to ms
|
||||
raw_result.update(
|
||||
{f"P{perc}": 1000 * raw_result["percentiles"][str(perc)]}
|
||||
)
|
||||
raw_result.update({f"P{perc}": 1000 * raw_result["percentiles"][str(perc)]})
|
||||
raw_result["avg_latency"] = raw_result["avg_latency"] * 1000
|
||||
|
||||
# add the result to raw_result
|
||||
@@ -142,38 +134,24 @@ if __name__ == "__main__":
|
||||
serving_results = pd.DataFrame.from_dict(serving_results)
|
||||
throughput_results = pd.DataFrame.from_dict(throughput_results)
|
||||
|
||||
raw_results_json = results_to_json(
|
||||
latency_results, throughput_results, serving_results
|
||||
)
|
||||
raw_results_json = results_to_json(latency_results, throughput_results, serving_results)
|
||||
|
||||
# remapping the key, for visualization purpose
|
||||
if not latency_results.empty:
|
||||
latency_results = latency_results[list(latency_column_mapping.keys())].rename(
|
||||
columns=latency_column_mapping
|
||||
)
|
||||
latency_results = latency_results[list(latency_column_mapping.keys())].rename(columns=latency_column_mapping)
|
||||
if not serving_results.empty:
|
||||
serving_results = serving_results[list(serving_column_mapping.keys())].rename(
|
||||
columns=serving_column_mapping
|
||||
)
|
||||
serving_results = serving_results[list(serving_column_mapping.keys())].rename(columns=serving_column_mapping)
|
||||
if not throughput_results.empty:
|
||||
throughput_results = throughput_results[
|
||||
list(throughput_results_column_mapping.keys())
|
||||
].rename(columns=throughput_results_column_mapping)
|
||||
throughput_results = throughput_results[list(throughput_results_column_mapping.keys())].rename(
|
||||
columns=throughput_results_column_mapping
|
||||
)
|
||||
|
||||
processed_results_json = results_to_json(
|
||||
latency_results, throughput_results, serving_results
|
||||
)
|
||||
processed_results_json = results_to_json(latency_results, throughput_results, serving_results)
|
||||
|
||||
# get markdown tables
|
||||
latency_md_table = tabulate(
|
||||
latency_results, headers="keys", tablefmt="pipe", showindex=False
|
||||
)
|
||||
serving_md_table = tabulate(
|
||||
serving_results, headers="keys", tablefmt="pipe", showindex=False
|
||||
)
|
||||
throughput_md_table = tabulate(
|
||||
throughput_results, headers="keys", tablefmt="pipe", showindex=False
|
||||
)
|
||||
latency_md_table = tabulate(latency_results, headers="keys", tablefmt="pipe", showindex=False)
|
||||
serving_md_table = tabulate(serving_results, headers="keys", tablefmt="pipe", showindex=False)
|
||||
throughput_md_table = tabulate(throughput_results, headers="keys", tablefmt="pipe", showindex=False)
|
||||
|
||||
# document the result
|
||||
print(output_folder)
|
||||
|
||||
@@ -10,6 +10,7 @@
|
||||
{serving_tests_markdown_table}
|
||||
|
||||
## Offline tests
|
||||
|
||||
### Latency tests
|
||||
|
||||
- Input length: 32 tokens.
|
||||
|
||||
@@ -18,7 +18,7 @@
|
||||
},
|
||||
"client_parameters": {
|
||||
"model": "Qwen/Qwen2.5-VL-7B-Instruct",
|
||||
"endpoint_type": "openai-chat",
|
||||
"backend": "openai-chat",
|
||||
"dataset_name": "hf",
|
||||
"hf_split": "train",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
@@ -45,7 +45,7 @@
|
||||
},
|
||||
"client_parameters": {
|
||||
"model": "Qwen/Qwen3-8B",
|
||||
"endpoint_type": "vllm",
|
||||
"backend": "vllm",
|
||||
"dataset_name": "sharegpt",
|
||||
"dataset_path": "/github/home/.cache/datasets/ShareGPT_V3_unfiltered_cleaned_split.json",
|
||||
"num_prompts": 200
|
||||
@@ -69,7 +69,7 @@
|
||||
},
|
||||
"client_parameters": {
|
||||
"model": "Qwen/Qwen2.5-7B-Instruct",
|
||||
"endpoint_type": "vllm",
|
||||
"backend": "vllm",
|
||||
"dataset_name": "sharegpt",
|
||||
"dataset_path": "/github/home/.cache/datasets/ShareGPT_V3_unfiltered_cleaned_split.json",
|
||||
"num_prompts": 200
|
||||
|
||||
@@ -1,25 +1,3 @@
|
||||
#
|
||||
# Attempt to find the python package that uses the same python executable as
|
||||
# `EXECUTABLE` and is one of the `SUPPORTED_VERSIONS`.
|
||||
#
|
||||
macro (find_python_from_executable EXECUTABLE SUPPORTED_VERSIONS)
|
||||
file(REAL_PATH ${EXECUTABLE} EXECUTABLE)
|
||||
set(Python_EXECUTABLE ${EXECUTABLE})
|
||||
find_package(Python COMPONENTS Interpreter Development.Module Development.SABIModule)
|
||||
if (NOT Python_FOUND)
|
||||
message(FATAL_ERROR "Unable to find python matching: ${EXECUTABLE}.")
|
||||
endif()
|
||||
set(_VER "${Python_VERSION_MAJOR}.${Python_VERSION_MINOR}")
|
||||
set(_SUPPORTED_VERSIONS_LIST ${SUPPORTED_VERSIONS} ${ARGN})
|
||||
if (NOT _VER IN_LIST _SUPPORTED_VERSIONS_LIST)
|
||||
message(FATAL_ERROR
|
||||
"Python version (${_VER}) is not one of the supported versions: "
|
||||
"${_SUPPORTED_VERSIONS_LIST}.")
|
||||
endif()
|
||||
message(STATUS "Found python matching: ${EXECUTABLE}.")
|
||||
endmacro()
|
||||
|
||||
#
|
||||
# Run `EXPR` in python. The standard output of python is stored in `OUT` and
|
||||
# has trailing whitespace stripped. If an error is encountered when running
|
||||
# python, a fatal message `ERR_MSG` is issued.
|
||||
@@ -46,88 +24,3 @@ macro (append_cmake_prefix_path PKG EXPR)
|
||||
"import ${PKG}; print(${EXPR})" "Failed to locate ${PKG} path")
|
||||
list(APPEND CMAKE_PREFIX_PATH ${_PREFIX_PATH})
|
||||
endmacro()
|
||||
|
||||
|
||||
# This cmake function is adapted from vllm /Users/ganyi/workspace/vllm-ascend/cmake/utils.cmake
|
||||
# Define a target named `GPU_MOD_NAME` for a single extension. The
|
||||
# arguments are:
|
||||
#
|
||||
# DESTINATION <dest> - Module destination directory.
|
||||
# LANGUAGE <lang> - The GPU language for this module, e.g CUDA, HIP,
|
||||
# etc.
|
||||
# SOURCES <sources> - List of source files relative to CMakeLists.txt
|
||||
# directory.
|
||||
#
|
||||
# Optional arguments:
|
||||
#
|
||||
# ARCHITECTURES <arches> - A list of target GPU architectures in cmake
|
||||
# format.
|
||||
# Refer `CMAKE_CUDA_ARCHITECTURES` documentation
|
||||
# and `CMAKE_HIP_ARCHITECTURES` for more info.
|
||||
# ARCHITECTURES will use cmake's defaults if
|
||||
# not provided.
|
||||
# COMPILE_FLAGS <flags> - Extra compiler flags passed to NVCC/hip.
|
||||
# INCLUDE_DIRECTORIES <dirs> - Extra include directories.
|
||||
# LIBRARIES <libraries> - Extra link libraries.
|
||||
# WITH_SOABI - Generate library with python SOABI suffix name.
|
||||
# USE_SABI <version> - Use python stable api <version>
|
||||
#
|
||||
# Note: optimization level/debug info is set via cmake build type.
|
||||
#
|
||||
function (define_gpu_extension_target GPU_MOD_NAME)
|
||||
cmake_parse_arguments(PARSE_ARGV 1
|
||||
GPU
|
||||
"WITH_SOABI"
|
||||
"DESTINATION;LANGUAGE;USE_SABI"
|
||||
"SOURCES;ARCHITECTURES;COMPILE_FLAGS;INCLUDE_DIRECTORIES;LIBRARIES")
|
||||
|
||||
# Add hipify preprocessing step when building with HIP/ROCm.
|
||||
if (GPU_LANGUAGE STREQUAL "HIP")
|
||||
hipify_sources_target(GPU_SOURCES ${GPU_MOD_NAME} "${GPU_SOURCES}")
|
||||
endif()
|
||||
|
||||
if (GPU_WITH_SOABI)
|
||||
set(GPU_WITH_SOABI WITH_SOABI)
|
||||
else()
|
||||
set(GPU_WITH_SOABI)
|
||||
endif()
|
||||
|
||||
if (GPU_USE_SABI)
|
||||
Python_add_library(${GPU_MOD_NAME} MODULE USE_SABI ${GPU_USE_SABI} ${GPU_WITH_SOABI} "${GPU_SOURCES}")
|
||||
else()
|
||||
Python_add_library(${GPU_MOD_NAME} MODULE ${GPU_WITH_SOABI} "${GPU_SOURCES}")
|
||||
endif()
|
||||
|
||||
if (GPU_LANGUAGE STREQUAL "HIP")
|
||||
# Make this target dependent on the hipify preprocessor step.
|
||||
add_dependencies(${GPU_MOD_NAME} hipify${GPU_MOD_NAME})
|
||||
endif()
|
||||
|
||||
if (GPU_ARCHITECTURES)
|
||||
set_target_properties(${GPU_MOD_NAME} PROPERTIES
|
||||
${GPU_LANGUAGE}_ARCHITECTURES "${GPU_ARCHITECTURES}")
|
||||
endif()
|
||||
|
||||
set_property(TARGET ${GPU_MOD_NAME} PROPERTY CXX_STANDARD 17)
|
||||
|
||||
target_compile_options(${GPU_MOD_NAME} PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:${GPU_LANGUAGE}>:${GPU_COMPILE_FLAGS}>)
|
||||
|
||||
target_compile_definitions(${GPU_MOD_NAME} PRIVATE
|
||||
"-DTORCH_EXTENSION_NAME=${GPU_MOD_NAME}")
|
||||
|
||||
target_include_directories(${GPU_MOD_NAME} PRIVATE csrc
|
||||
${GPU_INCLUDE_DIRECTORIES})
|
||||
|
||||
target_link_libraries(${GPU_MOD_NAME} PRIVATE torch ${GPU_LIBRARIES})
|
||||
|
||||
# Don't use `TORCH_LIBRARIES` for CUDA since it pulls in a bunch of
|
||||
# dependencies that are not necessary and may not be installed.
|
||||
if (GPU_LANGUAGE STREQUAL "CUDA")
|
||||
target_link_libraries(${GPU_MOD_NAME} PRIVATE CUDA::cudart CUDA::cuda_driver)
|
||||
else()
|
||||
target_link_libraries(${GPU_MOD_NAME} PRIVATE ${TORCH_LIBRARIES})
|
||||
endif()
|
||||
|
||||
install(TARGETS ${GPU_MOD_NAME} LIBRARY DESTINATION ${GPU_DESTINATION} COMPONENT ${GPU_MOD_NAME})
|
||||
endfunction()
|
||||
|
||||
247
collect_env.py
247
collect_env.py
@@ -18,42 +18,44 @@
|
||||
import datetime
|
||||
import locale
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from collections import namedtuple
|
||||
|
||||
import regex as re
|
||||
from vllm.envs import environment_variables
|
||||
|
||||
try:
|
||||
import torch
|
||||
|
||||
TORCH_AVAILABLE = True
|
||||
except (ImportError, NameError, AttributeError, OSError):
|
||||
TORCH_AVAILABLE = False
|
||||
|
||||
# System Environment Information
|
||||
SystemEnv = namedtuple(
|
||||
'SystemEnv',
|
||||
"SystemEnv",
|
||||
[
|
||||
'torch_version',
|
||||
'is_debug_build',
|
||||
'gcc_version',
|
||||
'clang_version',
|
||||
'cmake_version',
|
||||
'os',
|
||||
'libc_version',
|
||||
'python_version',
|
||||
'python_platform',
|
||||
'pip_version', # 'pip' or 'pip3'
|
||||
'pip_packages',
|
||||
'conda_packages',
|
||||
'cpu_info',
|
||||
'vllm_version', # vllm specific field
|
||||
'vllm_ascend_version', # vllm ascend specific field
|
||||
'env_vars',
|
||||
'npu_info', # ascend specific field
|
||||
'cann_info', # ascend specific field
|
||||
])
|
||||
"torch_version",
|
||||
"is_debug_build",
|
||||
"gcc_version",
|
||||
"clang_version",
|
||||
"cmake_version",
|
||||
"os",
|
||||
"libc_version",
|
||||
"python_version",
|
||||
"python_platform",
|
||||
"pip_version", # 'pip' or 'pip3'
|
||||
"pip_packages",
|
||||
"conda_packages",
|
||||
"cpu_info",
|
||||
"vllm_version", # vllm specific field
|
||||
"vllm_ascend_version", # vllm ascend specific field
|
||||
"env_vars",
|
||||
"npu_info", # ascend specific field
|
||||
"cann_info", # ascend specific field
|
||||
],
|
||||
)
|
||||
|
||||
DEFAULT_CONDA_PATTERNS = {
|
||||
"torch",
|
||||
@@ -65,6 +67,7 @@ DEFAULT_CONDA_PATTERNS = {
|
||||
"transformers",
|
||||
"zmq",
|
||||
"pynvml",
|
||||
"triton-ascend",
|
||||
}
|
||||
|
||||
DEFAULT_PIP_PATTERNS = {
|
||||
@@ -77,20 +80,18 @@ DEFAULT_PIP_PATTERNS = {
|
||||
"transformers",
|
||||
"zmq",
|
||||
"pynvml",
|
||||
"triton-ascend",
|
||||
}
|
||||
|
||||
|
||||
def run(command):
|
||||
"""Return (return-code, stdout, stderr)."""
|
||||
shell = True if type(command) is str else False
|
||||
p = subprocess.Popen(command,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
shell=shell)
|
||||
shell = isinstance(command, str)
|
||||
p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell)
|
||||
raw_output, raw_err = p.communicate()
|
||||
rc = p.returncode
|
||||
if get_platform() == 'win32':
|
||||
enc = 'oem'
|
||||
if get_platform() == "win32":
|
||||
enc = "oem"
|
||||
else:
|
||||
enc = locale.getpreferredencoding()
|
||||
output = raw_output.decode(enc)
|
||||
@@ -122,42 +123,40 @@ def run_and_return_first_line(run_lambda, command):
|
||||
rc, out, _ = run_lambda(command)
|
||||
if rc != 0:
|
||||
return None
|
||||
return out.split('\n')[0]
|
||||
return out.split("\n")[0]
|
||||
|
||||
|
||||
def get_conda_packages(run_lambda, patterns=None):
|
||||
if patterns is None:
|
||||
patterns = DEFAULT_CONDA_PATTERNS
|
||||
conda = os.environ.get('CONDA_EXE', 'conda')
|
||||
conda = os.environ.get("CONDA_EXE", "conda")
|
||||
out = run_and_read_all(run_lambda, "{} list".format(conda))
|
||||
if out is None:
|
||||
return out
|
||||
|
||||
return "\n".join(line for line in out.splitlines()
|
||||
if not line.startswith("#") and any(name in line
|
||||
for name in patterns))
|
||||
return "\n".join(
|
||||
line for line in out.splitlines() if not line.startswith("#") and any(name in line for name in patterns)
|
||||
)
|
||||
|
||||
|
||||
def get_gcc_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, 'gcc --version', r'gcc (.*)')
|
||||
return run_and_parse_first_match(run_lambda, "gcc --version", r"gcc (.*)")
|
||||
|
||||
|
||||
def get_clang_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, 'clang --version',
|
||||
r'clang version (.*)')
|
||||
return run_and_parse_first_match(run_lambda, "clang --version", r"clang version (.*)")
|
||||
|
||||
|
||||
def get_cmake_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, 'cmake --version',
|
||||
r'cmake (.*)')
|
||||
return run_and_parse_first_match(run_lambda, "cmake --version", r"cmake (.*)")
|
||||
|
||||
|
||||
def _parse_version(version, version_tuple):
|
||||
version_str = version_tuple[-1]
|
||||
if isinstance(version_str, str) and version_str.startswith('g'):
|
||||
if '.' in version_str:
|
||||
git_sha = version_str.split('.')[0][1:]
|
||||
date = version_str.split('.')[-1][1:]
|
||||
if isinstance(version_str, str) and version_str.startswith("g"):
|
||||
if "." in version_str:
|
||||
git_sha = version_str.split(".")[0][1:]
|
||||
date = version_str.split(".")[-1][1:]
|
||||
return f"{version} (git sha: {git_sha}, date: {date})"
|
||||
else:
|
||||
git_sha = version_str[1:] # type: ignore
|
||||
@@ -167,26 +166,28 @@ def _parse_version(version, version_tuple):
|
||||
|
||||
def get_vllm_version():
|
||||
from vllm import __version__, __version_tuple__
|
||||
|
||||
return _parse_version(__version__, __version_tuple__)
|
||||
|
||||
|
||||
def get_vllm_ascend_version():
|
||||
from vllm_ascend._version import __version__, __version_tuple__
|
||||
|
||||
return _parse_version(__version__, __version_tuple__)
|
||||
|
||||
|
||||
def get_cpu_info(run_lambda):
|
||||
rc, out, err = 0, '', ''
|
||||
if get_platform() == 'linux':
|
||||
rc, out, err = run_lambda('lscpu')
|
||||
elif get_platform() == 'win32':
|
||||
rc, out, err = 0, "", ""
|
||||
if get_platform() == "linux":
|
||||
rc, out, err = run_lambda("lscpu")
|
||||
elif get_platform() == "win32":
|
||||
rc, out, err = run_lambda(
|
||||
'wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID, \
|
||||
CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE'
|
||||
"wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID, \
|
||||
CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE"
|
||||
)
|
||||
elif get_platform() == 'darwin':
|
||||
elif get_platform() == "darwin":
|
||||
rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string")
|
||||
cpu_info = 'None'
|
||||
cpu_info = "None"
|
||||
if rc == 0:
|
||||
cpu_info = out
|
||||
else:
|
||||
@@ -195,67 +196,63 @@ def get_cpu_info(run_lambda):
|
||||
|
||||
|
||||
def get_platform():
|
||||
if sys.platform.startswith('linux'):
|
||||
return 'linux'
|
||||
elif sys.platform.startswith('win32'):
|
||||
return 'win32'
|
||||
elif sys.platform.startswith('cygwin'):
|
||||
return 'cygwin'
|
||||
elif sys.platform.startswith('darwin'):
|
||||
return 'darwin'
|
||||
if sys.platform.startswith("linux"):
|
||||
return "linux"
|
||||
elif sys.platform.startswith("win32"):
|
||||
return "win32"
|
||||
elif sys.platform.startswith("cygwin"):
|
||||
return "cygwin"
|
||||
elif sys.platform.startswith("darwin"):
|
||||
return "darwin"
|
||||
else:
|
||||
return sys.platform
|
||||
|
||||
|
||||
def get_mac_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, 'sw_vers -productVersion',
|
||||
r'(.*)')
|
||||
return run_and_parse_first_match(run_lambda, "sw_vers -productVersion", r"(.*)")
|
||||
|
||||
|
||||
def get_windows_version(run_lambda):
|
||||
system_root = os.environ.get('SYSTEMROOT', 'C:\\Windows')
|
||||
wmic_cmd = os.path.join(system_root, 'System32', 'Wbem', 'wmic')
|
||||
findstr_cmd = os.path.join(system_root, 'System32', 'findstr')
|
||||
return run_and_read_all(
|
||||
run_lambda,
|
||||
'{} os get Caption | {} /v Caption'.format(wmic_cmd, findstr_cmd))
|
||||
system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
|
||||
wmic_cmd = os.path.join(system_root, "System32", "Wbem", "wmic")
|
||||
findstr_cmd = os.path.join(system_root, "System32", "findstr")
|
||||
return run_and_read_all(run_lambda, "{} os get Caption | {} /v Caption".format(wmic_cmd, findstr_cmd))
|
||||
|
||||
|
||||
def get_lsb_version(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, 'lsb_release -a',
|
||||
r'Description:\t(.*)')
|
||||
return run_and_parse_first_match(run_lambda, "lsb_release -a", r"Description:\t(.*)")
|
||||
|
||||
|
||||
def check_release_file(run_lambda):
|
||||
return run_and_parse_first_match(run_lambda, 'cat /etc/*-release',
|
||||
r'PRETTY_NAME="(.*)"')
|
||||
return run_and_parse_first_match(run_lambda, "cat /etc/*-release", r'PRETTY_NAME="(.*)"')
|
||||
|
||||
|
||||
def get_os(run_lambda):
|
||||
from platform import machine
|
||||
|
||||
platform = get_platform()
|
||||
|
||||
if platform == 'win32' or platform == 'cygwin':
|
||||
if platform == "win32" or platform == "cygwin":
|
||||
return get_windows_version(run_lambda)
|
||||
|
||||
if platform == 'darwin':
|
||||
if platform == "darwin":
|
||||
version = get_mac_version(run_lambda)
|
||||
if version is None:
|
||||
return None
|
||||
return 'macOS {} ({})'.format(version, machine())
|
||||
return "macOS {} ({})".format(version, machine())
|
||||
|
||||
if platform == 'linux':
|
||||
if platform == "linux":
|
||||
# Ubuntu/Debian based
|
||||
desc = get_lsb_version(run_lambda)
|
||||
if desc is not None:
|
||||
return '{} ({})'.format(desc, machine())
|
||||
return "{} ({})".format(desc, machine())
|
||||
|
||||
# Try reading /etc/*-release
|
||||
desc = check_release_file(run_lambda)
|
||||
if desc is not None:
|
||||
return '{} ({})'.format(desc, machine())
|
||||
return "{} ({})".format(desc, machine())
|
||||
|
||||
return '{} ({})'.format(platform, machine())
|
||||
return "{} ({})".format(platform, machine())
|
||||
|
||||
# Unknown platform
|
||||
return platform
|
||||
@@ -263,14 +260,16 @@ def get_os(run_lambda):
|
||||
|
||||
def get_python_platform():
|
||||
import platform
|
||||
|
||||
return platform.platform()
|
||||
|
||||
|
||||
def get_libc_version():
|
||||
import platform
|
||||
if get_platform() != 'linux':
|
||||
return 'N/A'
|
||||
return '-'.join(platform.libc_ver())
|
||||
|
||||
if get_platform() != "linux":
|
||||
return "N/A"
|
||||
return "-".join(platform.libc_ver())
|
||||
|
||||
|
||||
def get_pip_packages(run_lambda, patterns=None):
|
||||
@@ -282,31 +281,29 @@ def get_pip_packages(run_lambda, patterns=None):
|
||||
# But here it is invoked as `python -mpip`
|
||||
def run_with_pip(pip):
|
||||
out = run_and_read_all(run_lambda, pip + ["list", "--format=freeze"])
|
||||
return "\n".join(line for line in out.splitlines()
|
||||
if any(name in line for name in patterns))
|
||||
return "\n".join(line for line in out.splitlines() if any(name in line for name in patterns))
|
||||
|
||||
pip_version = 'pip3' if sys.version[0] == '3' else 'pip'
|
||||
out = run_with_pip([sys.executable, '-mpip'])
|
||||
pip_version = "pip3" if sys.version[0] == "3" else "pip"
|
||||
out = run_with_pip([sys.executable, "-mpip"])
|
||||
|
||||
return pip_version, out
|
||||
|
||||
|
||||
def get_npu_info(run_lambda):
|
||||
return run_and_read_all(run_lambda, 'npu-smi info')
|
||||
return run_and_read_all(run_lambda, "npu-smi info")
|
||||
|
||||
|
||||
def get_cann_info(run_lambda):
|
||||
out = run_and_read_all(run_lambda, 'lscpu | grep Architecture:')
|
||||
out = run_and_read_all(run_lambda, "lscpu | grep Architecture:")
|
||||
cpu_arch = str(out).split()[-1]
|
||||
return run_and_read_all(
|
||||
run_lambda,
|
||||
'cat /usr/local/Ascend/ascend-toolkit/latest/{}-linux/ascend_toolkit_install.info'
|
||||
.format(cpu_arch))
|
||||
run_lambda, "cat /usr/local/Ascend/ascend-toolkit/latest/{}-linux/ascend_toolkit_install.info".format(cpu_arch)
|
||||
)
|
||||
|
||||
|
||||
def get_env_vars():
|
||||
env_vars = ''
|
||||
secret_terms = ('secret', 'token', 'api', 'access', 'password')
|
||||
env_vars = ""
|
||||
secret_terms = ("secret", "token", "api", "access", "password")
|
||||
report_prefix = ("TORCH", "PYTORCH", "ASCEND_", "ATB_")
|
||||
for k, v in os.environ.items():
|
||||
if any(term in k.lower() for term in secret_terms):
|
||||
@@ -327,7 +324,7 @@ def get_env_info():
|
||||
version_str = torch.__version__
|
||||
debug_mode_str = str(torch.version.debug)
|
||||
else:
|
||||
version_str = debug_mode_str = 'N/A'
|
||||
version_str = debug_mode_str = "N/A"
|
||||
|
||||
sys_version = sys.version.replace("\n", " ")
|
||||
|
||||
@@ -336,9 +333,7 @@ def get_env_info():
|
||||
return SystemEnv(
|
||||
torch_version=version_str,
|
||||
is_debug_build=debug_mode_str,
|
||||
python_version='{} ({}-bit runtime)'.format(
|
||||
sys_version,
|
||||
sys.maxsize.bit_length() + 1),
|
||||
python_version="{} ({}-bit runtime)".format(sys_version, sys.maxsize.bit_length() + 1),
|
||||
python_platform=get_python_platform(),
|
||||
pip_version=pip_version,
|
||||
pip_packages=pip_list_output,
|
||||
@@ -399,36 +394,35 @@ CANN:
|
||||
|
||||
|
||||
def pretty_str(envinfo):
|
||||
|
||||
def replace_nones(dct, replacement='Could not collect'):
|
||||
for key in dct.keys():
|
||||
def replace_nones(dct, replacement="Could not collect"):
|
||||
for key in dct:
|
||||
if dct[key] is not None:
|
||||
continue
|
||||
dct[key] = replacement
|
||||
return dct
|
||||
|
||||
def replace_bools(dct, true='Yes', false='No'):
|
||||
for key in dct.keys():
|
||||
def replace_bools(dct, true="Yes", false="No"):
|
||||
for key in dct:
|
||||
if dct[key] is True:
|
||||
dct[key] = true
|
||||
elif dct[key] is False:
|
||||
dct[key] = false
|
||||
return dct
|
||||
|
||||
def prepend(text, tag='[prepend]'):
|
||||
lines = text.split('\n')
|
||||
def prepend(text, tag="[prepend]"):
|
||||
lines = text.split("\n")
|
||||
updated_lines = [tag + line for line in lines]
|
||||
return '\n'.join(updated_lines)
|
||||
return "\n".join(updated_lines)
|
||||
|
||||
def replace_if_empty(text, replacement='No relevant packages'):
|
||||
def replace_if_empty(text, replacement="No relevant packages"):
|
||||
if text is not None and len(text) == 0:
|
||||
return replacement
|
||||
return text
|
||||
|
||||
def maybe_start_on_next_line(string):
|
||||
# If `string` is multiline, prepend a \n to it.
|
||||
if string is not None and len(string.split('\n')) > 1:
|
||||
return '\n{}\n'.format(string)
|
||||
if string is not None and len(string.split("\n")) > 1:
|
||||
return "\n{}\n".format(string)
|
||||
return string
|
||||
|
||||
mutable_dict = envinfo._asdict()
|
||||
@@ -440,22 +434,18 @@ def pretty_str(envinfo):
|
||||
mutable_dict = replace_nones(mutable_dict)
|
||||
|
||||
# If either of these are '', replace with 'No relevant packages'
|
||||
mutable_dict['pip_packages'] = replace_if_empty(
|
||||
mutable_dict['pip_packages'])
|
||||
mutable_dict['conda_packages'] = replace_if_empty(
|
||||
mutable_dict['conda_packages'])
|
||||
mutable_dict["pip_packages"] = replace_if_empty(mutable_dict["pip_packages"])
|
||||
mutable_dict["conda_packages"] = replace_if_empty(mutable_dict["conda_packages"])
|
||||
|
||||
# Tag conda and pip packages with a prefix
|
||||
# If they were previously None, they'll show up as ie '[conda] Could not collect'
|
||||
if mutable_dict['pip_packages']:
|
||||
mutable_dict['pip_packages'] = prepend(
|
||||
mutable_dict['pip_packages'], '[{}] '.format(envinfo.pip_version))
|
||||
if mutable_dict['conda_packages']:
|
||||
mutable_dict['conda_packages'] = prepend(
|
||||
mutable_dict['conda_packages'], '[conda] ')
|
||||
mutable_dict['cpu_info'] = envinfo.cpu_info
|
||||
mutable_dict['npu_info'] = envinfo.npu_info
|
||||
mutable_dict['cann_info'] = envinfo.cann_info
|
||||
if mutable_dict["pip_packages"]:
|
||||
mutable_dict["pip_packages"] = prepend(mutable_dict["pip_packages"], "[{}] ".format(envinfo.pip_version))
|
||||
if mutable_dict["conda_packages"]:
|
||||
mutable_dict["conda_packages"] = prepend(mutable_dict["conda_packages"], "[conda] ")
|
||||
mutable_dict["cpu_info"] = envinfo.cpu_info
|
||||
mutable_dict["npu_info"] = envinfo.npu_info
|
||||
mutable_dict["cann_info"] = envinfo.cann_info
|
||||
return env_info_fmt.format(**mutable_dict)
|
||||
|
||||
|
||||
@@ -468,22 +458,19 @@ def main():
|
||||
output = get_pretty_env_info()
|
||||
print(output)
|
||||
|
||||
if TORCH_AVAILABLE and hasattr(torch, 'utils') and hasattr(
|
||||
torch.utils, '_crash_handler'):
|
||||
if TORCH_AVAILABLE and hasattr(torch, "utils") and hasattr(torch.utils, "_crash_handler"):
|
||||
minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR
|
||||
if sys.platform == "linux" and os.path.exists(minidump_dir):
|
||||
dumps = [
|
||||
os.path.join(minidump_dir, dump)
|
||||
for dump in os.listdir(minidump_dir)
|
||||
]
|
||||
dumps = [os.path.join(minidump_dir, dump) for dump in os.listdir(minidump_dir)]
|
||||
latest = max(dumps, key=os.path.getctime)
|
||||
ctime = os.path.getctime(latest)
|
||||
creation_time = datetime.datetime.fromtimestamp(ctime).strftime(
|
||||
'%Y-%m-%d %H:%M:%S')
|
||||
msg = "\n*** Detected a minidump at {} created on {}, ".format(latest, creation_time) + \
|
||||
"if this is related to your bug please include it when you file a report ***"
|
||||
creation_time = datetime.datetime.fromtimestamp(ctime).strftime("%Y-%m-%d %H:%M:%S")
|
||||
msg = (
|
||||
"\n*** Detected a minidump at {} created on {}, ".format(latest, creation_time)
|
||||
+ "if this is related to your bug please include it when you file a report ***"
|
||||
)
|
||||
print(msg, file=sys.stderr)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
621
csrc/CMakeLists.txt
Normal file
621
csrc/CMakeLists.txt
Normal file
@@ -0,0 +1,621 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
|
||||
cmake_minimum_required(VERSION 3.16)
|
||||
|
||||
project(cann_ops-transformer)
|
||||
|
||||
option(BUILD_OPEN_PROJECT "Build open ascend ops project." ON)
|
||||
option(BUILD_OPS_RTY_KERNEL "Build return yellow kernel." OFF)
|
||||
option(ENABLE_CCACHE "Enable ccache capability" ON)
|
||||
option(ENABLE_BUILT_IN "Enable built-in package" OFF)
|
||||
option(ENABLE_STATIC "Enable Static" OFF)
|
||||
option(ENABLE_EXPERIMENTAL "Enable experimental module" OFF)
|
||||
option(ENABLE_TEST "Enable test" OFF)
|
||||
option(ENABLE_UT_EXEC "Enable exec ut" OFF)
|
||||
option(ENABLE_ASAN "Enable asan" OFF)
|
||||
option(ENABLE_VALGRIND "Enable valgrind" OFF)
|
||||
option(OP_HOST_UT "Enable ophost ut" OFF)
|
||||
option(OP_API_UT "Enable opapi ut" OFF)
|
||||
option(OP_GRAPH_UT "Enable graph ut" OFF)
|
||||
option(OP_KERNEL_UT "Enable kernel ut" OFF)
|
||||
option(OP_KERNEL_AICPU_UT "Enable aicpu kernel ut" OFF)
|
||||
option(UT_TEST_ALL "Enable all ut" OFF)
|
||||
option(ENABLE_OOM "Enable kernel oom" OFF)
|
||||
|
||||
set(ASCEND_COMPUTE_UNIT "ascend910b" CACHE STRING "soc that need to be compiled")
|
||||
set(ASCEND_OP_NAME "ALL" CACHE STRING "operators that need to be compiled")
|
||||
set(ARCH_DIRECTORY "" CACHE STRING "arch directory that need to be compiled")
|
||||
set(VENDOR_NAME "custom" CACHE STRING "vendor name")
|
||||
set(ASCEND_ALL_COMPUTE_UNIT "ascend310p;ascend910b;ascend910_93;ascend950;kirinx90" CACHE STRING "all soc list")
|
||||
|
||||
set(SOC_VERSION_LIST ascend310p ascend910b ascend910_93 ascend950 kirinx90)
|
||||
set(ARCH_DIRECTORY_LIST arch22 arch32 arch32 arch35 arch32)
|
||||
|
||||
if ("ascend950" IN_LIST ASCEND_COMPUTE_UNIT)
|
||||
message(STATUS "build with 3~8 packages........")
|
||||
set(BUILD_WITH_3_8_PACKAGE ON CACHE BOOL "build with 3~8 package and opsbase")
|
||||
endif()
|
||||
|
||||
foreach(SOC_VERSION ${ASCEND_COMPUTE_UNIT})
|
||||
list(FIND SOC_VERSION_LIST ${SOC_VERSION} INDEX)
|
||||
if(NOT INDEX EQUAL -1)
|
||||
list(GET ARCH_DIRECTORY_LIST ${INDEX} VAL)
|
||||
list(APPEND ARCH_DIRECTORY ${VAL})
|
||||
else()
|
||||
message(STATUS "unsupported chip type")
|
||||
if ((NOT BUILD_OPS_RTY_KERNEL) AND (BUILD_OPEN_PROJECT))
|
||||
include(cmake/build_empty_package.cmake)
|
||||
cpack_empty_package()
|
||||
return()
|
||||
endif()
|
||||
endif()
|
||||
endforeach()
|
||||
|
||||
list(FIND ARCH_DIRECTORY "arch32" INDEX)
|
||||
if(NOT INDEX EQUAL -1)
|
||||
list(APPEND ARCH_DIRECTORY "arch22")
|
||||
endif()
|
||||
|
||||
if(PROJECT_SOURCE_DIR STREQUAL CMAKE_SOURCE_DIR)
|
||||
message(STATUS "compile project with library")
|
||||
option(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG "Build ops-transformer with cann pkg" ON)
|
||||
else()
|
||||
message(STATUS "compile project with src")
|
||||
option(BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG "Build ops-transformer with cann source" OFF)
|
||||
endif()
|
||||
|
||||
if(UNIX)
|
||||
set(SYSTEM_PREFIX ${CMAKE_SYSTEM_PROCESSOR}-linux)
|
||||
endif()
|
||||
|
||||
#外部传参
|
||||
if(NOT ${CMAKE_BUILD_MODE} STREQUAL "FALSE")
|
||||
if(ENABLE_DEBUG)
|
||||
set(CMAKE_BUILD_MODE "${CMAKE_BUILD_MODE} -g")
|
||||
endif()
|
||||
set(COMPILE_OP_MODE ${CMAKE_BUILD_MODE})
|
||||
else()
|
||||
if(ENABLE_TEST)
|
||||
set(COMPILE_OP_MODE "-O0 -g")
|
||||
endif()
|
||||
if(ENABLE_DEBUG)
|
||||
set(CMAKE_BUILD_MODE "-g")
|
||||
endif()
|
||||
endif()
|
||||
|
||||
set(PKG_NAME transformer)
|
||||
set(OPS_TRANSFORMER_DIR ${CMAKE_CURRENT_SOURCE_DIR})
|
||||
set(CMAKE_CXX_STANDARD 17 CACHE STRING "c++17 is needed for this project")
|
||||
set_directory_properties(PROPERTIES
|
||||
ADDITIONAL_MAKE_CLEAN_FILES "${CMAKE_BINARY_DIR}/_CPack_Packages"
|
||||
)
|
||||
|
||||
# Suppress warnings from catlass/tla third-party headers for CANN kernel compilation
|
||||
set(VLLM_ASCEND_CANN_COMPAT_HEADER "${OPS_TRANSFORMER_DIR}/common/include/cann_compat.h")
|
||||
list(APPEND OPS_COMPILE_OPTIONS -Wno-ignored-attributes)
|
||||
list(APPEND OPS_COMPILE_OPTIONS -include${VLLM_ASCEND_CANN_COMPAT_HEADER})
|
||||
add_compile_options(
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-include${VLLM_ASCEND_CANN_COMPAT_HEADER}>
|
||||
)
|
||||
|
||||
include(cmake/config.cmake)
|
||||
include(cmake/func.cmake)
|
||||
include(cmake/third_party/json.cmake)
|
||||
if (ENABLE_TEST)
|
||||
include(${PROJECT_SOURCE_DIR}/cmake/third_party/gtest.cmake)
|
||||
endif()
|
||||
include(${OPS_ADV_CMAKE_DIR}/ut.cmake)
|
||||
|
||||
if (BUILD_OPEN_PROJECT)
|
||||
include(cmake/intf.cmake)
|
||||
add_definitions(-DBUILD_OPEN_PROJECT)
|
||||
if (BUILD_OPS_RTY_KERNEL)
|
||||
message(STATUS "Build return yellow kernel.")
|
||||
include(cmake/rty_obj_func.cmake)
|
||||
else()
|
||||
message(STATUS "Start building custom package.")
|
||||
include(ExternalProject)
|
||||
include(cmake/dependencies.cmake)
|
||||
include(cmake/variables.cmake)
|
||||
include(cmake/obj_func.cmake)
|
||||
include(cmake/third_party/abseil-cpp.cmake)
|
||||
include(cmake/third_party/ascend_protobuf.cmake)
|
||||
|
||||
include(cmake/third_party/makeself-fetch.cmake)
|
||||
include(cmake/opbuild.cmake)
|
||||
include(cmake/custom_build.cmake)
|
||||
message(STATUS "End building custom package.")
|
||||
if (ENABLE_OPS_HOST)
|
||||
gen_aclnn_with_opdef()
|
||||
endif()
|
||||
if (ENABLE_STATIC)
|
||||
include(cmake/static.cmake)
|
||||
endif()
|
||||
if (ENABLE_AICPU)
|
||||
include(cmake/symbol.cmake)
|
||||
gen_cust_aicpu_json_symbol()
|
||||
gen_cust_aicpu_kernel_symbol()
|
||||
endif()
|
||||
if (ENABLE_BUILT_IN)
|
||||
message(STATUS "Start building built-in package.")
|
||||
include(cmake/symbol.cmake)
|
||||
gen_norm_symbol()
|
||||
include(cmake/package.cmake)
|
||||
pack_built_in()
|
||||
else()
|
||||
include(cmake/package.cmake)
|
||||
pack_tiling_sink()
|
||||
endif()
|
||||
return()
|
||||
endif()
|
||||
else()
|
||||
include(cmake/dependencies.cmake)
|
||||
include(cmake/variables.cmake)
|
||||
include(cmake/opbuild.cmake)
|
||||
include(cmake/rty_obj_func.cmake)
|
||||
include(cmake/intf_pub_linux.cmake)
|
||||
endif()
|
||||
|
||||
if (BUILD_OPS_RTY_KERNEL)
|
||||
set(CMAKE_MODULE_PATH
|
||||
${CMAKE_MODULE_PATH}
|
||||
${CMAKE_CURRENT_LIST_DIR}/cmake/modules
|
||||
)
|
||||
|
||||
set(CMAKE_PREFIX_PATH
|
||||
${CMAKE_PREFIX_PATH}
|
||||
${ASCEND_CANN_PACKAGE_PATH}
|
||||
)
|
||||
|
||||
set(_op_host_aclnn_link
|
||||
$<BUILD_INTERFACE:intf_pub>
|
||||
exe_graph
|
||||
register
|
||||
c_sec
|
||||
)
|
||||
|
||||
find_package(alog MODULE)
|
||||
|
||||
if(NOT ${alog_FOUND})
|
||||
add_definitions(-DALOG_NOT_FOUND)
|
||||
endif()
|
||||
|
||||
add_library(op_host_aclnn SHARED EXCLUDE_FROM_ALL)
|
||||
target_link_libraries(op_host_aclnn PRIVATE
|
||||
${_op_host_aclnn_link}
|
||||
)
|
||||
target_compile_options(op_host_aclnn PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-std=gnu++1z>
|
||||
)
|
||||
|
||||
add_library(op_host_aclnnInner SHARED EXCLUDE_FROM_ALL)
|
||||
target_link_libraries(op_host_aclnnInner PRIVATE
|
||||
${_op_host_aclnn_link}
|
||||
)
|
||||
target_compile_options(op_host_aclnnInner PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-std=gnu++1z>
|
||||
)
|
||||
|
||||
add_library(op_host_aclnnExc SHARED EXCLUDE_FROM_ALL)
|
||||
target_link_libraries(op_host_aclnnExc PRIVATE
|
||||
${_op_host_aclnn_link}
|
||||
)
|
||||
target_compile_options(op_host_aclnnExc PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-std=gnu++1z>
|
||||
)
|
||||
|
||||
# op proto
|
||||
add_library(opsproto SHARED)
|
||||
target_compile_options(opsproto PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-std=c++11>
|
||||
-fvisibility=hidden
|
||||
)
|
||||
target_compile_definitions(opsproto PRIVATE
|
||||
LOG_CPP
|
||||
PROCESS_LOG
|
||||
)
|
||||
target_link_libraries(opsproto PRIVATE
|
||||
$<BUILD_INTERFACE:intf_pub>
|
||||
$<BUILD_INTERFACE:ops_transformer_utils_proto_headers>
|
||||
$<$<BOOL:${alog_FOUND}>:$<BUILD_INTERFACE:alog_headers>>
|
||||
-Wl,--whole-archive
|
||||
rt2_registry
|
||||
-Wl,--no-whole-archive
|
||||
-Wl,--no-as-needed
|
||||
exe_graph
|
||||
graph
|
||||
graph_base
|
||||
register
|
||||
ascendalog
|
||||
error_manager
|
||||
platform
|
||||
-Wl,--as-needed
|
||||
c_sec
|
||||
)
|
||||
set_target_properties(opsproto PROPERTIES OUTPUT_NAME
|
||||
cust_opsproto_rt2.0
|
||||
)
|
||||
install(TARGETS opsproto
|
||||
LIBRARY DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_proto/lib/linux/${CMAKE_SYSTEM_PROCESSOR}
|
||||
)
|
||||
|
||||
add_ops_tiling_keys(
|
||||
OP_NAME "ALL"
|
||||
TILING_KEYS ${TILING_KEY}
|
||||
)
|
||||
|
||||
add_opc_config(
|
||||
OP_NAME "ALL"
|
||||
CONFIG ${OP_DEBUG_CONFIG}
|
||||
)
|
||||
|
||||
if(ADD_OPS_COMPILE_OPTION_V2)
|
||||
add_ops_compile_options(
|
||||
OP_NAME "ALL"
|
||||
OPTIONS ${OPS_COMPILE_OPTIONS}
|
||||
)
|
||||
endif()
|
||||
endif ()
|
||||
|
||||
add_subdirectory(common)
|
||||
if (NOT BUILD_OPS_RTY_KERNEL)
|
||||
add_subdirectory(mc2)
|
||||
add_subdirectory(posembedding)
|
||||
endif()
|
||||
|
||||
if(NOT DEFINED COMPILED_OPS)
|
||||
set(COMPILED_OPS ${COMPILED_OPS} CACHE STRING "Comp")
|
||||
set(COMPILED_OPS CACHE STRING "Compiled Ops" FORCE)
|
||||
set(COMPILED_OP_DIRS CACHE STRING "Compiled Ops Dirs" FORCE)
|
||||
endif()
|
||||
|
||||
if(NOT DEFINED COMPILED_OP_DIRS)
|
||||
set(COMPILED_OP_DIRS CACHE STRING "Compiled Ops Dirs" FORCE)
|
||||
endif()
|
||||
|
||||
set(OP_LIST)
|
||||
set(OP_DIR_LIST)
|
||||
op_add_subdirectory(OP_LIST OP_DIR_LIST)
|
||||
|
||||
foreach (OP_DIR ${OP_DIR_LIST})
|
||||
if (EXISTS "${OP_DIR}/op_host")
|
||||
add_subdirectory(${OP_DIR}/op_host)
|
||||
else()
|
||||
add_subdirectory(${OP_DIR})
|
||||
endif()
|
||||
endforeach ()
|
||||
|
||||
add_subdirectory(moe)
|
||||
list(APPEND OP_LIST "moe_init_routing_v2")
|
||||
list(APPEND OP_LIST "moe_token_unpermute_with_ep_grad")
|
||||
list(APPEND OP_DIR_LIST ${CMAKE_CURRENT_SOURCE_DIR}/moe/moe_init_routing_v2)
|
||||
list(APPEND OP_DIR_LIST ${CMAKE_CURRENT_SOURCE_DIR}/moe/moe_token_unpermute_with_ep_grad)
|
||||
add_subdirectory(ffn)
|
||||
list(APPEND OP_LIST "ffn")
|
||||
list(APPEND OP_DIR_LIST ${CMAKE_CURRENT_SOURCE_DIR}/ffn/ffn)
|
||||
add_subdirectory(attention)
|
||||
list(APPEND OP_LIST ${COMPILED_OPS})
|
||||
list(REMOVE_DUPLICATES OP_LIST)
|
||||
list(APPEND OP_DIR_LIST ${COMPILED_OP_DIRS})
|
||||
list(REMOVE_DUPLICATES OP_DIR_LIST)
|
||||
add_subdirectory(gmm)
|
||||
list(REMOVE_DUPLICATES OP_LIST)
|
||||
list(APPEND OP_DIR_LIST ${COMPILED_OP_DIRS})
|
||||
list(REMOVE_DUPLICATES OP_DIR_LIST)
|
||||
add_subdirectory(mc2)
|
||||
list(REMOVE_DUPLICATES OP_LIST)
|
||||
list(APPEND OP_DIR_LIST ${COMPILED_OP_DIRS})
|
||||
list(REMOVE_DUPLICATES OP_DIR_LIST)
|
||||
list(APPEND OP_LIST "fused_gdn_gating")
|
||||
list(APPEND OP_DIR_LIST ${CMAKE_CURRENT_SOURCE_DIR}/attention/fused_gdn_gating)
|
||||
|
||||
set(OP_DEPEND_DIR_LIST)
|
||||
op_add_depend_directory(
|
||||
OP_LIST ${OP_LIST}
|
||||
OP_DIR_LIST OP_DEPEND_DIR_LIST
|
||||
)
|
||||
|
||||
foreach (OP_DEPEND_DIR ${OP_DEPEND_DIR_LIST})
|
||||
if (EXISTS "${OP_DEPEND_DIR}/op_host")
|
||||
add_subdirectory(${OP_DEPEND_DIR}/op_host)
|
||||
else()
|
||||
add_subdirectory(${OP_DEPEND_DIR})
|
||||
endif()
|
||||
endforeach ()
|
||||
|
||||
install(DIRECTORY ${OPS_ADV_ACT}/
|
||||
DESTINATION ${IMPL_INSTALL_DIR}/ascendc/common/act
|
||||
)
|
||||
|
||||
install(DIRECTORY ${OPS_GROUPEDMATMUL_ACT}/
|
||||
DESTINATION ${IMPL_INSTALL_DIR}/ascendc/common/groupedmatmul_act
|
||||
)
|
||||
|
||||
|
||||
if (BUILD_OPS_RTY_KERNEL)
|
||||
get_target_property(base_aclnn_srcs op_host_aclnn SOURCES)
|
||||
get_target_property(base_aclnn_inner_srcs op_host_aclnnInner SOURCES)
|
||||
get_target_property(base_aclnn_exclude_srcs op_host_aclnnExc SOURCES)
|
||||
set(base_aclnn_binary_dir ${ASCEND_AUTOGEN_DIR})
|
||||
|
||||
set(generate_aclnn_srcs)
|
||||
set(generate_aclnn_inner_srcs)
|
||||
set(generate_aclnn_headers)
|
||||
set(generate_proto_dir ${base_aclnn_binary_dir})
|
||||
set(generate_exclude_proto_srcs)
|
||||
set(generate_proto_srcs)
|
||||
set(generate_proto_headers)
|
||||
|
||||
if (base_aclnn_srcs)
|
||||
foreach (_src ${base_aclnn_srcs})
|
||||
string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}")
|
||||
if (is_match)
|
||||
get_filename_component(name_without_ext ${_src} NAME_WE)
|
||||
|
||||
string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext})
|
||||
list(APPEND generate_aclnn_srcs ${base_aclnn_binary_dir}/aclnn_${_op_name}.cpp)
|
||||
list(APPEND generate_aclnn_headers ${base_aclnn_binary_dir}/aclnn_${_op_name}.h)
|
||||
list(APPEND generate_proto_srcs ${generate_proto_dir}/${_op_name}_proto.cpp)
|
||||
list(APPEND generate_proto_headers ${generate_proto_dir}/${_op_name}_proto.h)
|
||||
endif ()
|
||||
endforeach ()
|
||||
else ()
|
||||
add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp
|
||||
COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp
|
||||
)
|
||||
|
||||
target_sources(op_host_aclnn PRIVATE
|
||||
${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_stub.cpp
|
||||
)
|
||||
endif ()
|
||||
|
||||
if (base_aclnn_inner_srcs)
|
||||
foreach (_src ${base_aclnn_inner_srcs})
|
||||
string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}")
|
||||
if (is_match)
|
||||
get_filename_component(name_without_ext ${_src} NAME_WE)
|
||||
string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext})
|
||||
list(APPEND generate_aclnn_inner_srcs ${base_aclnn_binary_dir}/inner/aclnnInner_${_op_name}.cpp)
|
||||
list(APPEND generate_proto_srcs ${generate_proto_dir}/inner/${_op_name}_proto.cpp)
|
||||
list(APPEND generate_proto_headers ${generate_proto_dir}/inner/${_op_name}_proto.h)
|
||||
endif ()
|
||||
endforeach ()
|
||||
else ()
|
||||
add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp
|
||||
COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp
|
||||
)
|
||||
|
||||
target_sources(op_host_aclnnInner PRIVATE
|
||||
${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_inner_stub.cpp
|
||||
)
|
||||
endif ()
|
||||
|
||||
if (base_aclnn_exclude_srcs)
|
||||
foreach (_src ${base_aclnn_exclude_srcs})
|
||||
string(REGEX MATCH "^${CMAKE_CURRENT_SOURCE_DIR}" is_match "${_src}")
|
||||
if (is_match)
|
||||
get_filename_component(name_without_ext ${_src} NAME_WE)
|
||||
string(REGEX REPLACE "_def$" "" _op_name ${name_without_ext})
|
||||
list(APPEND generate_exclude_proto_srcs ${generate_proto_dir}/exc/${_op_name}_proto.cpp)
|
||||
list(APPEND generate_proto_srcs ${generate_proto_dir}/exc/${_op_name}_proto.cpp)
|
||||
list(APPEND generate_proto_headers ${generate_proto_dir}/exc/${_op_name}_proto.h)
|
||||
endif ()
|
||||
endforeach ()
|
||||
else()
|
||||
add_custom_command(OUTPUT ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp
|
||||
COMMAND touch ${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp
|
||||
)
|
||||
|
||||
target_sources(op_host_aclnnExc PRIVATE
|
||||
${CMAKE_CURRENT_BINARY_DIR}/op_host_aclnn_exc_stub.cpp
|
||||
)
|
||||
endif ()
|
||||
|
||||
|
||||
if (generate_aclnn_srcs OR generate_aclnn_inner_srcs)
|
||||
set(ops_aclnn_src ${generate_aclnn_srcs} ${generate_aclnn_inner_srcs})
|
||||
else ()
|
||||
set(ops_aclnn_src ${CMAKE_CURRENT_BINARY_DIR}/ops_aclnn_src_stub.cpp)
|
||||
|
||||
add_custom_command(OUTPUT ${ops_aclnn_src}
|
||||
COMMAND touch ${ops_aclnn_src}
|
||||
)
|
||||
endif ()
|
||||
|
||||
set_source_files_properties(${ops_aclnn_src}
|
||||
PROPERTIES GENERATED TRUE
|
||||
)
|
||||
add_library(ops_aclnn STATIC
|
||||
${ops_aclnn_src}
|
||||
)
|
||||
target_compile_options(ops_aclnn PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CXX>:-std=gnu++1z>
|
||||
)
|
||||
target_link_libraries(ops_aclnn PRIVATE
|
||||
$<BUILD_INTERFACE:intf_pub>
|
||||
)
|
||||
add_dependencies(ops_aclnn opbuild_gen_default opbuild_gen_inner)
|
||||
|
||||
set_source_files_properties(${generate_proto_srcs}
|
||||
PROPERTIES GENERATED TRUE
|
||||
)
|
||||
target_sources(opsproto PRIVATE
|
||||
${generate_proto_srcs}
|
||||
)
|
||||
add_dependencies(opsproto ops_transformer_proto_headers)
|
||||
|
||||
install(FILES ${generate_proto_headers}
|
||||
DESTINATION packages/vendors/${VENDOR_NAME}_transformer/op_proto/inc OPTIONAL
|
||||
)
|
||||
|
||||
add_library(ops_transformer_proto_headers INTERFACE)
|
||||
|
||||
target_include_directories(ops_transformer_proto_headers INTERFACE
|
||||
$<BUILD_INTERFACE:${generate_proto_dir}>
|
||||
$<BUILD_INTERFACE:${generate_proto_dir}/inner>
|
||||
$<BUILD_INTERFACE:${generate_proto_dir}/exc>
|
||||
$<INSTALL_INTERFACE:include/ops_adv/proto>
|
||||
)
|
||||
|
||||
add_dependencies(ops_transformer_proto_headers opbuild_gen_default opbuild_gen_inner opbuild_gen_exc)
|
||||
|
||||
if (NOT BUILD_OPEN_PROJECT)
|
||||
if (generate_proto_srcs)
|
||||
install_package(
|
||||
PACKAGE ops_adv
|
||||
TARGETS ops_proto_headers
|
||||
FILES ${generate_proto_headers}
|
||||
DESTINATION include/ops_adv/proto
|
||||
)
|
||||
endif ()
|
||||
endif ()
|
||||
|
||||
if (generate_aclnn_srcs)
|
||||
add_custom_command(OUTPUT ${generate_aclnn_srcs} ${generate_aclnn_headers}
|
||||
COMMAND mkdir -p ${base_aclnn_binary_dir}
|
||||
COMMAND OPS_PROTO_SEPARATE=1
|
||||
OPS_ACLNN_GEN=1
|
||||
OPS_PROJECT_NAME=aclnn
|
||||
${OP_BUILD_TOOL}
|
||||
$<TARGET_FILE:op_host_aclnn>
|
||||
${base_aclnn_binary_dir}
|
||||
)
|
||||
endif ()
|
||||
|
||||
add_custom_target(opbuild_gen_default
|
||||
DEPENDS ${generate_aclnn_srcs} ${generate_aclnn_headers} op_host_aclnn
|
||||
)
|
||||
|
||||
if (generate_aclnn_inner_srcs)
|
||||
add_custom_command(OUTPUT ${generate_aclnn_inner_srcs}
|
||||
COMMAND mkdir -p ${base_aclnn_binary_dir}/inner
|
||||
COMMAND OPS_PROTO_SEPARATE=1
|
||||
OPS_ACLNN_GEN=1
|
||||
OPS_PROJECT_NAME=aclnnInner
|
||||
${OP_BUILD_TOOL}
|
||||
$<TARGET_FILE:op_host_aclnnInner>
|
||||
${base_aclnn_binary_dir}/inner
|
||||
)
|
||||
endif ()
|
||||
|
||||
add_custom_target(opbuild_gen_inner
|
||||
DEPENDS ${generate_aclnn_inner_srcs} op_host_aclnnInner
|
||||
)
|
||||
|
||||
if (generate_exclude_proto_srcs)
|
||||
add_custom_command(OUTPUT ${generate_exclude_proto_srcs}
|
||||
COMMAND mkdir -p ${base_aclnn_binary_dir}/exc
|
||||
COMMAND OPS_PROTO_SEPARATE=1
|
||||
OPS_ACLNN_GEN=0
|
||||
OPS_PROJECT_NAME=aclnnExc
|
||||
${OP_BUILD_TOOL}
|
||||
$<TARGET_FILE:op_host_aclnnExc>
|
||||
${base_aclnn_binary_dir}/exc
|
||||
)
|
||||
endif ()
|
||||
|
||||
add_custom_target(opbuild_gen_exc
|
||||
DEPENDS ${generate_exclude_proto_srcs} op_host_aclnnExc
|
||||
)
|
||||
|
||||
add_custom_target(generate_transformer_adapt_py
|
||||
COMMAND ${HI_PYTHON} ${CMAKE_CURRENT_SOURCE_DIR}/cmake/scripts/util/ascendc_impl_build.py
|
||||
\"\"
|
||||
\"\"
|
||||
\"\"
|
||||
\"\"
|
||||
${ASCEND_IMPL_OUT_DIR}
|
||||
${ASCEND_AUTOGEN_DIR}
|
||||
--opsinfo-dir ${base_aclnn_binary_dir} ${base_aclnn_binary_dir}/inner ${base_aclnn_binary_dir}/exc
|
||||
)
|
||||
|
||||
add_dependencies(generate_transformer_adapt_py opbuild_gen_default opbuild_gen_inner opbuild_gen_exc)
|
||||
|
||||
foreach (_op_name ${OP_LIST})
|
||||
install(FILES ${ASCEND_IMPL_OUT_DIR}/dynamic/${_op_name}.py
|
||||
DESTINATION ${IMPL_DYNAMIC_INSTALL_DIR}
|
||||
OPTIONAL
|
||||
)
|
||||
endforeach ()
|
||||
|
||||
install(DIRECTORY ${OPS_ADV_UTILS_KERNEL_INC}/
|
||||
DESTINATION ${IMPL_INSTALL_DIR}/ascendc/common
|
||||
)
|
||||
|
||||
foreach (op_dir ${OP_DIR_LIST})
|
||||
get_filename_component(_op_name "${op_dir}" NAME)
|
||||
|
||||
if (EXISTS "${op_dir}/op_kernel")
|
||||
file(GLOB KERNEL_FILES
|
||||
${op_dir}/op_kernel/*.cpp
|
||||
${op_dir}/op_kernel/*.h
|
||||
)
|
||||
else()
|
||||
file(GLOB KERNEL_FILES
|
||||
${op_dir}/*.cpp
|
||||
${op_dir}/*.h
|
||||
)
|
||||
endif()
|
||||
|
||||
install(FILES ${KERNEL_FILES}
|
||||
DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name}
|
||||
OPTIONAL
|
||||
)
|
||||
|
||||
install(DIRECTORY ${op_dir}/arch32
|
||||
DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name}
|
||||
OPTIONAL
|
||||
)
|
||||
|
||||
install(DIRECTORY ${op_dir}/arch35
|
||||
DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name}
|
||||
OPTIONAL
|
||||
)
|
||||
|
||||
install(DIRECTORY ${op_dir}/arch38
|
||||
DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name}
|
||||
OPTIONAL
|
||||
)
|
||||
|
||||
install(DIRECTORY ${op_dir}/regbase/opkernel
|
||||
DESTINATION ${IMPL_INSTALL_DIR}/ascendc/${_op_name}/regbase
|
||||
OPTIONAL
|
||||
)
|
||||
endforeach ()
|
||||
|
||||
add_custom_target(prepare_build ALL)
|
||||
add_custom_target(generate_compile_cmd ALL)
|
||||
add_custom_target(generate_ops_info ALL)
|
||||
add_dependencies(prepare_build generate_transformer_adapt_py generate_compile_cmd)
|
||||
|
||||
foreach (compute_unit ${ASCEND_COMPUTE_UNIT})
|
||||
add_compile_cmd_target(
|
||||
COMPUTE_UNIT ${compute_unit}
|
||||
)
|
||||
|
||||
add_ops_info_target(
|
||||
COMPUTE_UNIT ${compute_unit}
|
||||
)
|
||||
endforeach ()
|
||||
|
||||
add_custom_target(ops_transformer_kernel ALL)
|
||||
add_custom_target(ops_transformer_config ALL)
|
||||
add_dependencies(ops_transformer_kernel ops_transformer_config)
|
||||
|
||||
foreach (compute_unit ${ASCEND_COMPUTE_UNIT})
|
||||
add_bin_compile_target(
|
||||
COMPUTE_UNIT
|
||||
${compute_unit}
|
||||
OP_INFO
|
||||
${OP_DIR_LIST}
|
||||
)
|
||||
endforeach ()
|
||||
endif ()
|
||||
30
csrc/aclnn_torch_adapter/NPUBridge.cpp
Normal file
30
csrc/aclnn_torch_adapter/NPUBridge.cpp
Normal file
@@ -0,0 +1,30 @@
|
||||
// Copyright (c) 2020, Huawei Technologies Co., Ltd
|
||||
// All rights reserved.
|
||||
//
|
||||
// This source code is licensed under the BSD-style license found in the
|
||||
// LICENSE file in the root directory of this source tree.
|
||||
|
||||
#include "NPUBridge.h"
|
||||
|
||||
namespace vllm_ascend
|
||||
{
|
||||
NPUStorageImpl *NPUBridge::GetNpuStorageImpl(c10::StorageImpl *storageImpl)
|
||||
{
|
||||
return static_cast<NPUStorageImpl *>(storageImpl);
|
||||
}
|
||||
|
||||
NPUStorageImpl *NPUBridge::GetNpuStorageImpl(c10::Storage &&storage)
|
||||
{
|
||||
return static_cast<NPUStorageImpl *>(storage.unsafeGetStorageImpl());
|
||||
}
|
||||
|
||||
NPUStorageImpl *NPUBridge::GetNpuStorageImpl(const at::Tensor &tensor)
|
||||
{
|
||||
return static_cast<NPUStorageImpl *>(tensor.storage().unsafeGetStorageImpl());
|
||||
}
|
||||
|
||||
NPUStorageDesc &NPUBridge::GetNpuStorageImplDesc(const at::Tensor &tensor)
|
||||
{
|
||||
return static_cast<NPUStorageImpl *>(tensor.storage().unsafeGetStorageImpl())->npu_desc_;
|
||||
}
|
||||
}
|
||||
29
csrc/aclnn_torch_adapter/NPUBridge.h
Normal file
29
csrc/aclnn_torch_adapter/NPUBridge.h
Normal file
@@ -0,0 +1,29 @@
|
||||
// Copyright (c) 2020, Huawei Technologies Co., Ltd
|
||||
// All rights reserved.
|
||||
//
|
||||
// This source code is licensed under the BSD-style license found in the
|
||||
// LICENSE file in the root directory of this source tree.
|
||||
|
||||
#pragma once
|
||||
#include <c10/core/StorageImpl.h>
|
||||
#include "NPUStorageImpl.h"
|
||||
|
||||
namespace vllm_ascend
|
||||
{
|
||||
|
||||
class NPUBridge
|
||||
{
|
||||
public:
|
||||
// at::tensor to NPUStorageImpl
|
||||
static NPUStorageImpl *GetNpuStorageImpl(const at::Tensor &tensor);
|
||||
|
||||
// c10::StorageImpl to NPUStorageImpl
|
||||
static NPUStorageImpl *GetNpuStorageImpl(c10::StorageImpl *storageImpl);
|
||||
|
||||
// c10::Storage to NPUStorageImpl
|
||||
static NPUStorageImpl *GetNpuStorageImpl(c10::Storage &&storage);
|
||||
|
||||
// tensor to NPUStorageDesc
|
||||
static NPUStorageDesc &GetNpuStorageImplDesc(const at::Tensor &tensor);
|
||||
};
|
||||
}
|
||||
52
csrc/aclnn_torch_adapter/NPUStorageImpl.cpp
Normal file
52
csrc/aclnn_torch_adapter/NPUStorageImpl.cpp
Normal file
@@ -0,0 +1,52 @@
|
||||
// Copyright (c) 2020, Huawei Technologies Co., Ltd
|
||||
// All rights reserved.
|
||||
//
|
||||
// This source code is licensed under the BSD-style license found in the
|
||||
// LICENSE file in the root directory of this source tree.
|
||||
|
||||
#include "NPUStorageImpl.h"
|
||||
|
||||
namespace vllm_ascend
|
||||
{
|
||||
|
||||
NPUStorageImpl::NPUStorageImpl(
|
||||
use_byte_size_t use_byte_size,
|
||||
size_t size_bytes,
|
||||
at::DataPtr data_ptr,
|
||||
at::Allocator *allocator,
|
||||
bool resizable) : c10::StorageImpl(use_byte_size,
|
||||
size_bytes,
|
||||
at::DataPtr(std::move(data_ptr)),
|
||||
allocator,
|
||||
resizable)
|
||||
{
|
||||
}
|
||||
|
||||
void NPUStorageImpl::release_resources()
|
||||
{
|
||||
StorageImpl::release_resources();
|
||||
}
|
||||
|
||||
c10::intrusive_ptr<c10::StorageImpl> make_npu_storage_impl(
|
||||
c10::StorageImpl::use_byte_size_t,
|
||||
c10::SymInt size_bytes,
|
||||
c10::DataPtr data_ptr,
|
||||
c10::Allocator *allocator,
|
||||
bool resizable)
|
||||
{
|
||||
if (data_ptr == nullptr)
|
||||
{
|
||||
data_ptr = allocator->allocate(size_bytes.as_int_unchecked());
|
||||
}
|
||||
// Correctly create NPUStorageImpl object.
|
||||
c10::intrusive_ptr<c10::StorageImpl> npu_storage_impl = c10::make_intrusive<NPUStorageImpl>(
|
||||
c10::StorageImpl::use_byte_size_t(),
|
||||
size_bytes.as_int_unchecked(),
|
||||
std::move(data_ptr),
|
||||
allocator,
|
||||
resizable);
|
||||
// There is no need to consider the NPUStorageDesc information, it will be carried out in the subsequent processing.
|
||||
return npu_storage_impl;
|
||||
}
|
||||
|
||||
}
|
||||
67
csrc/aclnn_torch_adapter/NPUStorageImpl.h
Normal file
67
csrc/aclnn_torch_adapter/NPUStorageImpl.h
Normal file
@@ -0,0 +1,67 @@
|
||||
// Copyright (c) 2020, Huawei Technologies Co., Ltd
|
||||
// All rights reserved.
|
||||
//
|
||||
// This source code is licensed under the BSD-style license found in the
|
||||
// LICENSE file in the root directory of this source tree.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/Tensor.h>
|
||||
#include <c10/core/StorageImpl.h>
|
||||
#include <c10/core/Allocator.h>
|
||||
#include <c10/core/ScalarType.h>
|
||||
#include <c10/util/typeid.h>
|
||||
#include <c10/util/order_preserving_flat_hash_map.h>
|
||||
|
||||
#include "acl/acl_rt.h"
|
||||
#include "acl/acl_base.h"
|
||||
|
||||
namespace vllm_ascend
|
||||
{
|
||||
|
||||
struct NPUStorageDesc
|
||||
{
|
||||
public:
|
||||
struct use_byte_size_t
|
||||
{
|
||||
};
|
||||
|
||||
c10::SmallVector<int64_t, 5> base_sizes_;
|
||||
c10::SmallVector<int64_t, 5> base_strides_;
|
||||
c10::SmallVector<int64_t, 5> storage_sizes_;
|
||||
int64_t base_offset_ = 0;
|
||||
use_byte_size_t base_dtype_ = {};
|
||||
aclFormat origin_format_ = ACL_FORMAT_UNDEFINED;
|
||||
aclFormat npu_format_ = ACL_FORMAT_ND;
|
||||
// used to make CANN GE tensor from storagImpl
|
||||
caffe2::TypeMeta data_type_ = caffe2::TypeMeta::Make<uint8_t>();
|
||||
};
|
||||
|
||||
struct NPUStorageImpl : public c10::StorageImpl
|
||||
{
|
||||
explicit NPUStorageImpl(
|
||||
use_byte_size_t use_byte_size,
|
||||
size_t size_bytes,
|
||||
at::DataPtr data_ptr,
|
||||
at::Allocator *allocator,
|
||||
bool resizable);
|
||||
~NPUStorageImpl() override = default;
|
||||
|
||||
void release_resources() override;
|
||||
|
||||
NPUStorageDesc npu_desc_;
|
||||
|
||||
NPUStorageDesc get_npu_desc() const
|
||||
{
|
||||
return npu_desc_;
|
||||
}
|
||||
};
|
||||
|
||||
c10::intrusive_ptr<c10::StorageImpl> make_npu_storage_impl(
|
||||
c10::StorageImpl::use_byte_size_t,
|
||||
c10::SymInt size_bytes,
|
||||
c10::DataPtr data_ptr,
|
||||
c10::Allocator *allocator,
|
||||
bool resizable);
|
||||
|
||||
}
|
||||
754
csrc/aclnn_torch_adapter/op_api_common.h
Normal file
754
csrc/aclnn_torch_adapter/op_api_common.h
Normal file
@@ -0,0 +1,754 @@
|
||||
// Copyright (c) 2023 Huawei Technologies Co., Ltd
|
||||
// All rights reserved.
|
||||
//
|
||||
// Licensed under the BSD 3-Clause License (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// https://opensource.org/licenses/BSD-3-Clause
|
||||
//
|
||||
// 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 OP_API_COMMON_ADAPTER
|
||||
#define OP_API_COMMON_ADAPTER
|
||||
|
||||
#include <fstream>
|
||||
#include <torch/types.h>
|
||||
#include <ATen/Tensor.h>
|
||||
#include <acl/acl_base.h>
|
||||
#include <c10/util/Exception.h>
|
||||
#include <dlfcn.h>
|
||||
#include <functional>
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
#include <torch_npu/csrc/framework/utils/CalcuOpUtil.h>
|
||||
#include <torch_npu/csrc/framework/utils/OpAdapter.h>
|
||||
#include "torch_npu/csrc/aten/NPUNativeFunctions.h"
|
||||
#include "torch_npu/csrc/core/npu/NPUStream.h"
|
||||
#include "torch_npu/csrc/framework/OpCommand.h"
|
||||
#include "torch_npu/csrc/framework/interface/EnvVariables.h"
|
||||
#include "torch_npu/csrc/framework/utils/CalcuOpUtil.h"
|
||||
#include "torch_npu/csrc/framework/utils/OpPreparation.h"
|
||||
#include "NPUBridge.h"
|
||||
#include "NPUStorageImpl.h"
|
||||
|
||||
#define NPU_NAME_SPACE at_npu::native
|
||||
using namespace at;
|
||||
|
||||
typedef struct aclOpExecutor aclOpExecutor;
|
||||
typedef struct aclTensor aclTensor;
|
||||
typedef struct aclScalar aclScalar;
|
||||
typedef struct aclIntArray aclIntArray;
|
||||
typedef struct aclFloatArray aclFloatArray;
|
||||
typedef struct aclBoolArray aclBoolArray;
|
||||
typedef struct aclTensorList aclTensorList;
|
||||
|
||||
typedef aclTensor *(*_aclCreateTensor)(
|
||||
const int64_t *view_dims, uint64_t view_dims_num, aclDataType data_type,
|
||||
const int64_t *stride, int64_t offset, aclFormat format,
|
||||
const int64_t *storage_dims, uint64_t storage_dims_num, void *tensor_data);
|
||||
typedef aclScalar *(*_aclCreateScalar)(void *value, aclDataType data_type);
|
||||
typedef aclIntArray *(*_aclCreateIntArray)(const int64_t *value, uint64_t size);
|
||||
typedef aclFloatArray *(*_aclCreateFloatArray)(const float *value,
|
||||
uint64_t size);
|
||||
typedef aclBoolArray *(*_aclCreateBoolArray)(const bool *value, uint64_t size);
|
||||
typedef aclTensorList *(*_aclCreateTensorList)(const aclTensor *const *value,
|
||||
uint64_t size);
|
||||
|
||||
typedef int (*_aclDestroyTensor)(const aclTensor *tensor);
|
||||
typedef int (*_aclDestroyScalar)(const aclScalar *scalar);
|
||||
typedef int (*_aclDestroyIntArray)(const aclIntArray *array);
|
||||
typedef int (*_aclDestroyFloatArray)(const aclFloatArray *array);
|
||||
typedef int (*_aclDestroyBoolArray)(const aclBoolArray *array);
|
||||
typedef int (*_aclDestroyTensorList)(const aclTensorList *array);
|
||||
|
||||
constexpr int kHashBufSize = 8192;
|
||||
constexpr int kHashBufMaxSize = kHashBufSize + 1024;
|
||||
extern thread_local char g_hashBuf[kHashBufSize];
|
||||
extern thread_local int g_hashOffset;
|
||||
|
||||
#ifdef MMCV_WITH_XLA
|
||||
#define DEVICE_TYPE at_npu::key::NativeDeviceType
|
||||
#else
|
||||
#define DEVICE_TYPE c10::DeviceType::PrivateUse1
|
||||
#endif
|
||||
|
||||
#define AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(_) \
|
||||
_(at::ScalarType::Byte, ACL_UINT8) \
|
||||
_(at::ScalarType::Char, ACL_INT8) \
|
||||
_(at::ScalarType::Short, ACL_INT16) \
|
||||
_(at::ScalarType::Int, ACL_INT32) \
|
||||
_(at::ScalarType::Long, ACL_INT64) \
|
||||
_(at::ScalarType::Half, ACL_FLOAT16) \
|
||||
_(at::ScalarType::Float, ACL_FLOAT) \
|
||||
_(at::ScalarType::Double, ACL_DOUBLE) \
|
||||
_(at::ScalarType::ComplexHalf, ACL_COMPLEX32) \
|
||||
_(at::ScalarType::ComplexFloat, ACL_COMPLEX64) \
|
||||
_(at::ScalarType::ComplexDouble, ACL_COMPLEX128) \
|
||||
_(at::ScalarType::Bool, ACL_BOOL) \
|
||||
_(at::ScalarType::QInt8, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::QUInt8, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::QInt32, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::BFloat16, ACL_BF16) \
|
||||
_(at::ScalarType::QUInt4x2, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::QUInt2x4, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Bits1x8, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Bits2x4, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Bits4x2, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Bits8, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Bits16, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Float8_e5m2, ACL_FLOAT8_E5M2) \
|
||||
_(at::ScalarType::Float8_e4m3fn, ACL_FLOAT8_E4M3FN) \
|
||||
_(at::ScalarType::Float8_e5m2fnuz, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Float8_e4m3fnuz, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::UInt16, ACL_UINT16) \
|
||||
_(at::ScalarType::UInt32, ACL_UINT32) \
|
||||
_(at::ScalarType::UInt64, ACL_UINT64) \
|
||||
_(at::ScalarType::UInt1, ACL_UINT1) \
|
||||
_(at::ScalarType::UInt2, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::UInt3, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::UInt4, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::UInt5, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::UInt6, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::UInt7, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Int1, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Int2, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Int3, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Int4, ACL_INT4) \
|
||||
_(at::ScalarType::Int5, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Int6, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Int7, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::Float8_e8m0fnu, ACL_FLOAT8_E8M0) \
|
||||
_(at::ScalarType::Float4_e2m1fn_x2, ACL_FLOAT4_E2M1) \
|
||||
_(at::ScalarType::Undefined, ACL_DT_UNDEFINED) \
|
||||
_(at::ScalarType::NumOptions, ACL_DT_UNDEFINED)
|
||||
|
||||
constexpr aclDataType kATenScalarTypeToAclDataTypeTable
|
||||
[static_cast<int64_t>(at::ScalarType::NumOptions) + 1] = {
|
||||
#define DEFINE_ENUM(_1, n) n,
|
||||
AT_ALL_SCALAR_TYPE_AND_ACL_DATATYPE_PAIR(DEFINE_ENUM)
|
||||
#undef DEFINE_ENUM
|
||||
};
|
||||
|
||||
#define GET_OP_API_FUNC(apiName) \
|
||||
reinterpret_cast<_##apiName>(GetOpApiFuncAddr(#apiName))
|
||||
|
||||
#define MEMCPY_TO_BUF(data_expression, size_expression) \
|
||||
if (g_hashOffset + (size_expression) > kHashBufSize) { \
|
||||
g_hashOffset = kHashBufMaxSize; \
|
||||
return; \
|
||||
} \
|
||||
memcpy(g_hashBuf + g_hashOffset, data_expression, size_expression); \
|
||||
g_hashOffset += size_expression;
|
||||
|
||||
bool IsOpInputBaseFormat(const at::Tensor &tensor)
|
||||
{
|
||||
if (!tensor.is_privateuseone()) {
|
||||
return true;
|
||||
}
|
||||
const auto format = vllm_ascend::NPUBridge::GetNpuStorageImplDesc(tensor).npu_format_;
|
||||
return (format == ACL_FORMAT_ND) || (format == ACL_FORMAT_NCHW) || (format == ACL_FORMAT_NHWC) ||
|
||||
(format == ACL_FORMAT_NCDHW);
|
||||
}
|
||||
|
||||
static std::vector<std::string> split_str(std::string s, const std::string &del)
|
||||
{
|
||||
int end = s.find(del);
|
||||
std::vector<std::string> path_list;
|
||||
while (end != -1) {
|
||||
path_list.push_back(s.substr(0, end));
|
||||
s.erase(s.begin(), s.begin() + end + 1);
|
||||
end = s.find(del);
|
||||
}
|
||||
path_list.push_back(s);
|
||||
return path_list;
|
||||
}
|
||||
|
||||
static bool is_file_exist(const std::string &path)
|
||||
{
|
||||
if (path.empty() || path.size() > PATH_MAX) {
|
||||
return false;
|
||||
}
|
||||
return (access(path.c_str(), F_OK) == 0) ? true : false;
|
||||
}
|
||||
|
||||
inline std::string real_path(const std::string &path)
|
||||
{
|
||||
if (path.empty() || path.size() > PATH_MAX) {
|
||||
return "";
|
||||
}
|
||||
char realPath[PATH_MAX] = {0};
|
||||
if (realpath(path.c_str(), realPath) == nullptr) {
|
||||
return "";
|
||||
}
|
||||
return std::string(realPath);
|
||||
}
|
||||
|
||||
inline std::vector<std::string> get_custom_lib_path()
|
||||
{
|
||||
char *ascend_custom_opppath = std::getenv("ASCEND_CUSTOM_OPP_PATH");
|
||||
std::vector<std::string> custom_lib_path_list;
|
||||
|
||||
if (ascend_custom_opppath == nullptr) {
|
||||
return std::vector<std::string>();
|
||||
}
|
||||
|
||||
std::string ascend_custom_opppath_str(ascend_custom_opppath);
|
||||
// split string with ":"
|
||||
custom_lib_path_list = split_str(ascend_custom_opppath_str, ":");
|
||||
if (custom_lib_path_list.empty()) {
|
||||
return std::vector<std::string>();
|
||||
}
|
||||
for (auto &it : custom_lib_path_list) {
|
||||
it = it + "/op_api/lib/";
|
||||
}
|
||||
|
||||
return custom_lib_path_list;
|
||||
}
|
||||
|
||||
inline std::vector<std::string> get_default_custom_lib_path()
|
||||
{
|
||||
char *ascend_opp_path = std::getenv("ASCEND_OPP_PATH");
|
||||
std::vector<std::string> default_vendors_list;
|
||||
|
||||
if (ascend_opp_path == nullptr) {
|
||||
return std::vector<std::string>();
|
||||
}
|
||||
|
||||
std::string vendors_path(ascend_opp_path);
|
||||
vendors_path = vendors_path + "/vendors";
|
||||
std::string vendors_config_file = real_path(vendors_path + "/config.ini");
|
||||
if (vendors_config_file.empty()) {
|
||||
return std::vector<std::string>();
|
||||
}
|
||||
|
||||
if (!is_file_exist(vendors_config_file)) {
|
||||
return std::vector<std::string>();
|
||||
}
|
||||
|
||||
std::ifstream ifs(vendors_config_file);
|
||||
std::string line;
|
||||
while (std::getline(ifs, line)) {
|
||||
if (line.find("load_priority=") == 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
std::string head = "load_priority=";
|
||||
line.erase(0, head.length());
|
||||
|
||||
// split string with ","
|
||||
default_vendors_list = split_str(line, ",");
|
||||
if (default_vendors_list.empty()) {
|
||||
return std::vector<std::string>();
|
||||
}
|
||||
for (auto &it : default_vendors_list) {
|
||||
it = real_path(vendors_path + "/" + it + "/op_api/lib/");
|
||||
}
|
||||
|
||||
return default_vendors_list;
|
||||
}
|
||||
|
||||
const std::vector<std::string> g_custom_lib_path = get_custom_lib_path();
|
||||
const std::vector<std::string> g_default_custom_lib_path = get_default_custom_lib_path();
|
||||
|
||||
inline const char *GetOpApiLibName(void) { return "libopapi.so"; }
|
||||
|
||||
inline const char *GetCustOpApiLibName(void) { return "libcust_opapi.so"; }
|
||||
|
||||
inline void *GetOpApiFuncAddrInLib(void *handler, const char *libName,
|
||||
const char *apiName) {
|
||||
auto funcAddr = dlsym(handler, apiName);
|
||||
return funcAddr;
|
||||
}
|
||||
|
||||
inline void *GetOpApiLibHandler(const char *libName) {
|
||||
auto handler = dlopen(libName, RTLD_LAZY);
|
||||
return handler;
|
||||
}
|
||||
|
||||
inline void *GetOpApiFuncAddr(const char *apiName)
|
||||
{
|
||||
if (!g_custom_lib_path.empty()) {
|
||||
for (auto &it : g_custom_lib_path) {
|
||||
auto cust_opapi_lib = real_path(it + "/" + GetCustOpApiLibName());
|
||||
if (cust_opapi_lib.empty()) {
|
||||
continue;
|
||||
}
|
||||
auto custOpApiHandler = GetOpApiLibHandler(cust_opapi_lib.c_str());
|
||||
if (custOpApiHandler != nullptr) {
|
||||
auto funcAddr =
|
||||
GetOpApiFuncAddrInLib(custOpApiHandler, GetCustOpApiLibName(), apiName);
|
||||
if (funcAddr != nullptr) {
|
||||
return funcAddr;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (!g_default_custom_lib_path.empty()) {
|
||||
for (auto &it : g_default_custom_lib_path) {
|
||||
auto default_cust_opapi_lib = real_path(it + "/" + GetCustOpApiLibName());
|
||||
if (default_cust_opapi_lib.empty()) {
|
||||
continue;
|
||||
}
|
||||
auto custOpApiHandler = GetOpApiLibHandler(default_cust_opapi_lib.c_str());
|
||||
if (custOpApiHandler != nullptr) {
|
||||
auto funcAddr =
|
||||
GetOpApiFuncAddrInLib(custOpApiHandler, GetCustOpApiLibName(), apiName);
|
||||
if (funcAddr != nullptr) {
|
||||
return funcAddr;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static auto opApiHandler = GetOpApiLibHandler(GetOpApiLibName());
|
||||
if (opApiHandler == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
return GetOpApiFuncAddrInLib(opApiHandler, GetOpApiLibName(), apiName);
|
||||
}
|
||||
|
||||
inline c10::Scalar ConvertTensorToScalar(const at::Tensor &tensor) {
|
||||
c10::Scalar expScalar;
|
||||
const at::Tensor *aclInput = &tensor;
|
||||
if (aclInput->scalar_type() == at::ScalarType::Double) {
|
||||
double value = *(double *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
} else if (aclInput->scalar_type() == at::ScalarType::Long) {
|
||||
int64_t value = *(int64_t *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
} else if (aclInput->scalar_type() == at::ScalarType::Float) {
|
||||
float value = *(float *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
} else if (aclInput->scalar_type() == at::ScalarType::Int) {
|
||||
int value = *(int *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
} else if (aclInput->scalar_type() == at::ScalarType::Half) {
|
||||
c10::Half value = *(c10::Half *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
} else if (aclInput->scalar_type() == at::ScalarType::Bool) {
|
||||
int8_t value = *(int8_t *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
} else if (aclInput->scalar_type() == at::ScalarType::ComplexDouble) {
|
||||
c10::complex<double> value = *(c10::complex<double> *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
} else if (aclInput->scalar_type() == at::ScalarType::ComplexFloat) {
|
||||
c10::complex<float> value = *(c10::complex<float> *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
} else if (aclInput->scalar_type() == at::ScalarType::BFloat16) {
|
||||
c10::BFloat16 value = *(c10::BFloat16 *)aclInput->data_ptr();
|
||||
c10::Scalar scalar(value);
|
||||
expScalar = scalar;
|
||||
}
|
||||
return expScalar;
|
||||
}
|
||||
|
||||
inline at::Tensor CopyTensorHostToDevice(const at::Tensor &cpu_tensor) {
|
||||
at::Tensor cpuPinMemTensor = cpu_tensor.pin_memory();
|
||||
int deviceIndex = 0;
|
||||
return cpuPinMemTensor.to(c10::Device(DEVICE_TYPE, deviceIndex),
|
||||
cpuPinMemTensor.scalar_type(), true, true);
|
||||
}
|
||||
|
||||
inline at::Tensor CopyScalarToDevice(const c10::Scalar &cpu_scalar,
|
||||
at::ScalarType scalar_data_type) {
|
||||
return CopyTensorHostToDevice(
|
||||
scalar_to_tensor(cpu_scalar).to(scalar_data_type));
|
||||
}
|
||||
|
||||
inline aclTensor *ConvertType(const at::Tensor &at_tensor) {
|
||||
static const auto aclCreateTensor = GET_OP_API_FUNC(aclCreateTensor);
|
||||
if (aclCreateTensor == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
if (!at_tensor.defined()) {
|
||||
return nullptr;
|
||||
}
|
||||
at::ScalarType scalar_data_type = at_tensor.scalar_type();
|
||||
aclDataType acl_data_type =
|
||||
kATenScalarTypeToAclDataTypeTable[static_cast<int64_t>(scalar_data_type)];
|
||||
TORCH_CHECK(
|
||||
acl_data_type != ACL_DT_UNDEFINED,
|
||||
std::string(c10::toString(scalar_data_type)) + " has not been supported")
|
||||
c10::SmallVector<int64_t, 5> storageDims;
|
||||
// if acl_data_type is ACL_STRING, storageDims is empty.
|
||||
auto itemsize = at_tensor.itemsize();
|
||||
TORCH_CHECK(itemsize != 0, "When ConvertType, tensor item size cannot be zero.");
|
||||
|
||||
const auto dimNum = at_tensor.sizes().size();
|
||||
aclFormat format = ACL_FORMAT_ND;
|
||||
if (!IsOpInputBaseFormat(at_tensor)) {
|
||||
format = vllm_ascend::NPUBridge::GetNpuStorageImpl(at_tensor)->npu_desc_.npu_format_;
|
||||
if (acl_data_type != ACL_STRING) {
|
||||
storageDims = vllm_ascend::NPUBridge::GetNpuStorageImpl(at_tensor)->npu_desc_.storage_sizes_;
|
||||
}
|
||||
} else {
|
||||
switch (dimNum) {
|
||||
case 3:
|
||||
format = ACL_FORMAT_NCL;
|
||||
break;
|
||||
case 4:
|
||||
format = ACL_FORMAT_NCHW;
|
||||
break;
|
||||
case 5:
|
||||
format = ACL_FORMAT_NCDHW;
|
||||
break;
|
||||
default:
|
||||
format = ACL_FORMAT_ND;
|
||||
}
|
||||
if (acl_data_type != ACL_STRING) {
|
||||
storageDims.push_back(at_tensor.storage().nbytes() / itemsize);
|
||||
}
|
||||
}
|
||||
|
||||
if (at_tensor.unsafeGetTensorImpl()->is_wrapped_number()) {
|
||||
c10::Scalar expScalar = ConvertTensorToScalar(at_tensor);
|
||||
at::Tensor aclInput = CopyScalarToDevice(expScalar, scalar_data_type);
|
||||
return aclCreateTensor(aclInput.sizes().data(), aclInput.sizes().size(),
|
||||
acl_data_type, aclInput.strides().data(),
|
||||
aclInput.storage_offset(), format,
|
||||
storageDims.data(), storageDims.size(),
|
||||
const_cast<void *>(aclInput.storage().data()));
|
||||
}
|
||||
|
||||
auto acl_tensor = aclCreateTensor(
|
||||
at_tensor.sizes().data(), at_tensor.sizes().size(), acl_data_type,
|
||||
at_tensor.strides().data(), at_tensor.storage_offset(), format,
|
||||
storageDims.data(), storageDims.size(),
|
||||
const_cast<void *>(at_tensor.storage().data()));
|
||||
return acl_tensor;
|
||||
}
|
||||
|
||||
inline aclScalar *ConvertType(const at::Scalar &at_scalar) {
|
||||
static const auto aclCreateScalar = GET_OP_API_FUNC(aclCreateScalar);
|
||||
if (aclCreateScalar == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
at::ScalarType scalar_data_type = at_scalar.type();
|
||||
aclDataType acl_data_type =
|
||||
kATenScalarTypeToAclDataTypeTable[static_cast<int64_t>(scalar_data_type)];
|
||||
TORCH_CHECK(
|
||||
acl_data_type != ACL_DT_UNDEFINED,
|
||||
std::string(c10::toString(scalar_data_type)) + " has not been supported")
|
||||
aclScalar *acl_scalar = nullptr;
|
||||
switch (scalar_data_type) {
|
||||
case at::ScalarType::Double: {
|
||||
double value = at_scalar.toDouble();
|
||||
acl_scalar = aclCreateScalar(&value, acl_data_type);
|
||||
break;
|
||||
}
|
||||
case at::ScalarType::Long: {
|
||||
int64_t value = at_scalar.toLong();
|
||||
acl_scalar = aclCreateScalar(&value, acl_data_type);
|
||||
break;
|
||||
}
|
||||
case at::ScalarType::Bool: {
|
||||
bool value = at_scalar.toBool();
|
||||
acl_scalar = aclCreateScalar(&value, acl_data_type);
|
||||
break;
|
||||
}
|
||||
case at::ScalarType::ComplexDouble: {
|
||||
auto value = at_scalar.toComplexDouble();
|
||||
acl_scalar = aclCreateScalar(&value, acl_data_type);
|
||||
break;
|
||||
}
|
||||
default:
|
||||
acl_scalar = nullptr;
|
||||
break;
|
||||
}
|
||||
return acl_scalar;
|
||||
}
|
||||
|
||||
inline aclIntArray *ConvertType(const at::IntArrayRef &at_array) {
|
||||
static const auto aclCreateIntArray = GET_OP_API_FUNC(aclCreateIntArray);
|
||||
if (aclCreateIntArray == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
auto array = aclCreateIntArray(at_array.data(), at_array.size());
|
||||
return array;
|
||||
}
|
||||
|
||||
template <std::size_t N>
|
||||
inline aclBoolArray *ConvertType(const std::array<bool, N> &value) {
|
||||
static const auto aclCreateBoolArray = GET_OP_API_FUNC(aclCreateBoolArray);
|
||||
if (aclCreateBoolArray == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
auto array = aclCreateBoolArray(value.data(), value.size());
|
||||
return array;
|
||||
}
|
||||
|
||||
inline aclBoolArray *ConvertType(const at::ArrayRef<bool> &value) {
|
||||
static const auto aclCreateBoolArray = GET_OP_API_FUNC(aclCreateBoolArray);
|
||||
if (aclCreateBoolArray == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
auto array = aclCreateBoolArray(value.data(), value.size());
|
||||
return array;
|
||||
}
|
||||
|
||||
inline aclTensorList *ConvertType(const at::TensorList &at_tensor_list) {
|
||||
static const auto aclCreateTensorList = GET_OP_API_FUNC(aclCreateTensorList);
|
||||
if (aclCreateTensorList == nullptr) {
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
std::vector<const aclTensor *> tensor_list(at_tensor_list.size());
|
||||
for (size_t i = 0; i < at_tensor_list.size(); i++) {
|
||||
tensor_list[i] = ConvertType(at_tensor_list[i]);
|
||||
}
|
||||
auto acl_tensor_list =
|
||||
aclCreateTensorList(tensor_list.data(), tensor_list.size());
|
||||
return acl_tensor_list;
|
||||
}
|
||||
|
||||
inline aclTensor *ConvertType(const c10::optional<at::Tensor> &opt_tensor) {
|
||||
if (opt_tensor.has_value() && opt_tensor.value().defined()) {
|
||||
return ConvertType(opt_tensor.value());
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
inline aclTensorList *ConvertType(
|
||||
const c10::optional<at::TensorList> &opt_tensor_list) {
|
||||
if (opt_tensor_list.has_value()) {
|
||||
return ConvertType(opt_tensor_list.value());
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
inline aclIntArray *ConvertType(
|
||||
const c10::optional<at::IntArrayRef> &opt_array) {
|
||||
if (opt_array.has_value()) {
|
||||
return ConvertType(opt_array.value());
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
inline aclScalar *ConvertType(const c10::optional<at::Scalar> &opt_scalar) {
|
||||
if (opt_scalar.has_value()) {
|
||||
return ConvertType(opt_scalar.value());
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
inline aclDataType ConvertType(const at::ScalarType scalarType) {
|
||||
return kATenScalarTypeToAclDataTypeTable[static_cast<int64_t>(scalarType)];
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
T ConvertType(T value) {
|
||||
return value;
|
||||
}
|
||||
|
||||
template <typename Tuple, size_t... I>
|
||||
auto ConvertToOpApiFunc(const Tuple ¶ms, void *opApiAddr,
|
||||
std::index_sequence<I...>) {
|
||||
typedef int (*OpApiFunc)(
|
||||
typename std::decay<decltype(std::get<I>(params))>::type...);
|
||||
auto func = reinterpret_cast<OpApiFunc>(opApiAddr);
|
||||
return func;
|
||||
}
|
||||
|
||||
template <typename Tuple>
|
||||
auto ConvertToOpApiFunc(const Tuple ¶ms, void *opApiAddr) {
|
||||
static constexpr auto size = std::tuple_size<Tuple>::value;
|
||||
return ConvertToOpApiFunc(params, opApiAddr,
|
||||
std::make_index_sequence<size>{});
|
||||
}
|
||||
|
||||
inline void Release(aclTensor *p) {
|
||||
static const auto aclDestroyTensor = GET_OP_API_FUNC(aclDestroyTensor);
|
||||
if (aclDestroyTensor == nullptr) {
|
||||
return;
|
||||
}
|
||||
aclDestroyTensor(p);
|
||||
}
|
||||
|
||||
inline void Release(aclScalar *p) {
|
||||
static const auto aclDestroyScalar = GET_OP_API_FUNC(aclDestroyScalar);
|
||||
if (aclDestroyScalar == nullptr) {
|
||||
return;
|
||||
}
|
||||
aclDestroyScalar(p);
|
||||
}
|
||||
|
||||
inline void Release(aclIntArray *p) {
|
||||
static const auto aclDestroyIntArray = GET_OP_API_FUNC(aclDestroyIntArray);
|
||||
if (aclDestroyIntArray == nullptr) {
|
||||
return;
|
||||
}
|
||||
|
||||
aclDestroyIntArray(p);
|
||||
}
|
||||
|
||||
inline void Release(aclBoolArray *p) {
|
||||
static const auto aclDestroyBoolArray = GET_OP_API_FUNC(aclDestroyBoolArray);
|
||||
if (aclDestroyBoolArray == nullptr) {
|
||||
return;
|
||||
}
|
||||
|
||||
aclDestroyBoolArray(p);
|
||||
}
|
||||
|
||||
inline void Release(aclTensorList *p) {
|
||||
static const auto aclDestroyTensorList =
|
||||
GET_OP_API_FUNC(aclDestroyTensorList);
|
||||
if (aclDestroyTensorList == nullptr) {
|
||||
return;
|
||||
}
|
||||
|
||||
aclDestroyTensorList(p);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void Release(T value) {
|
||||
(void)value;
|
||||
}
|
||||
|
||||
template <typename Tuple, size_t... I>
|
||||
void CallRelease(Tuple t, std::index_sequence<I...>) {
|
||||
(void)std::initializer_list<int>{(Release(std::get<I>(t)), 0)...};
|
||||
}
|
||||
|
||||
template <typename Tuple>
|
||||
void ReleaseConvertTypes(Tuple &t) {
|
||||
static constexpr auto size = std::tuple_size<Tuple>::value;
|
||||
CallRelease(t, std::make_index_sequence<size>{});
|
||||
}
|
||||
|
||||
template <typename... Ts>
|
||||
constexpr auto ConvertTypes(Ts &... args) {
|
||||
return std::make_tuple(ConvertType(args)...);
|
||||
}
|
||||
|
||||
template <typename Function, typename Tuple, size_t... I>
|
||||
auto call(Function f, Tuple t, std::index_sequence<I...>) {
|
||||
return f(std::get<I>(t)...);
|
||||
}
|
||||
|
||||
template <typename Function, typename Tuple>
|
||||
auto call(Function f, Tuple t) {
|
||||
static constexpr auto size = std::tuple_size<Tuple>::value;
|
||||
return call(f, t, std::make_index_sequence<size>{});
|
||||
}
|
||||
|
||||
template <std::size_t N>
|
||||
void AddParamToBuf(const std::array<bool, N> &value) {
|
||||
MEMCPY_TO_BUF(value.data(), value.size() * sizeof(bool));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void AddParamToBuf(const T &value) {
|
||||
MEMCPY_TO_BUF(&value, sizeof(T));
|
||||
}
|
||||
|
||||
void AddParamToBuf(const at::Tensor &);
|
||||
void AddParamToBuf(const at::Scalar &);
|
||||
void AddParamToBuf(const at::IntArrayRef &);
|
||||
void AddParamToBuf(const at::ArrayRef<bool> &);
|
||||
void AddParamToBuf(const at::TensorList &);
|
||||
void AddParamToBuf(const c10::optional<at::Tensor> &);
|
||||
void AddParamToBuf(const c10::optional<at::IntArrayRef> &);
|
||||
void AddParamToBuf(const c10::optional<at::Scalar> &);
|
||||
void AddParamToBuf(const at::ScalarType);
|
||||
void AddParamToBuf(const string &);
|
||||
void AddParamToBuf();
|
||||
|
||||
template <typename T, typename... Args>
|
||||
void AddParamToBuf(const T &arg, Args &... args) {
|
||||
AddParamToBuf(arg);
|
||||
AddParamToBuf(args...);
|
||||
}
|
||||
|
||||
uint64_t CalcHashId();
|
||||
typedef int (*InitHugeMemThreadLocal)(void *, bool);
|
||||
typedef void (*UnInitHugeMemThreadLocal)(void *, bool);
|
||||
typedef void (*ReleaseHugeMem)(void *, bool);
|
||||
|
||||
#define EXEC_NPU_CMD(aclnn_api, ...) \
|
||||
do { \
|
||||
static const auto getWorkspaceSizeFuncAddr = \
|
||||
GetOpApiFuncAddr(#aclnn_api "GetWorkspaceSize"); \
|
||||
static const auto opApiFuncAddr = GetOpApiFuncAddr(#aclnn_api); \
|
||||
static const auto initMemAddr = \
|
||||
GetOpApiFuncAddr("InitHugeMemThreadLocal"); \
|
||||
static const auto unInitMemAddr = \
|
||||
GetOpApiFuncAddr("UnInitHugeMemThreadLocal"); \
|
||||
static const auto releaseMemAddr = GetOpApiFuncAddr("ReleaseHugeMem"); \
|
||||
TORCH_CHECK( \
|
||||
getWorkspaceSizeFuncAddr != nullptr && opApiFuncAddr != nullptr, \
|
||||
#aclnn_api, " or ", #aclnn_api "GetWorkspaceSize", " not in ", \
|
||||
GetOpApiLibName(), ", or ", GetOpApiLibName(), "not found."); \
|
||||
auto acl_stream = c10_npu::getCurrentNPUStream().stream(false); \
|
||||
uint64_t workspace_size = 0; \
|
||||
uint64_t *workspace_size_addr = &workspace_size; \
|
||||
aclOpExecutor *executor = nullptr; \
|
||||
aclOpExecutor **executor_addr = &executor; \
|
||||
InitHugeMemThreadLocal initMemFunc = \
|
||||
reinterpret_cast<InitHugeMemThreadLocal>(initMemAddr); \
|
||||
UnInitHugeMemThreadLocal unInitMemFunc = \
|
||||
reinterpret_cast<UnInitHugeMemThreadLocal>(unInitMemAddr); \
|
||||
if (initMemFunc) { \
|
||||
initMemFunc(nullptr, false); \
|
||||
} \
|
||||
auto converted_params = \
|
||||
ConvertTypes(__VA_ARGS__, workspace_size_addr, executor_addr); \
|
||||
static auto getWorkspaceSizeFunc = \
|
||||
ConvertToOpApiFunc(converted_params, getWorkspaceSizeFuncAddr); \
|
||||
auto workspace_status = call(getWorkspaceSizeFunc, converted_params); \
|
||||
TORCH_CHECK(workspace_status == 0, \
|
||||
"call " #aclnn_api " failed, detail:", aclGetRecentErrMsg()); \
|
||||
void *workspace_addr = nullptr; \
|
||||
if (workspace_size != 0) { \
|
||||
at::TensorOptions options = \
|
||||
at::TensorOptions(torch_npu::utils::get_npu_device_type()); \
|
||||
auto workspace_tensor = \
|
||||
at::empty({workspace_size}, options.dtype(kByte)); \
|
||||
workspace_addr = const_cast<void *>(workspace_tensor.storage().data()); \
|
||||
} \
|
||||
auto acl_call = [converted_params, workspace_addr, workspace_size, \
|
||||
acl_stream, executor]() -> int { \
|
||||
typedef int (*OpApiFunc)(void *, uint64_t, aclOpExecutor *, \
|
||||
const aclrtStream); \
|
||||
OpApiFunc opApiFunc = reinterpret_cast<OpApiFunc>(opApiFuncAddr); \
|
||||
auto api_ret = \
|
||||
opApiFunc(workspace_addr, workspace_size, executor, acl_stream); \
|
||||
TORCH_CHECK(api_ret == 0, "call " #aclnn_api " failed, detail:", \
|
||||
aclGetRecentErrMsg()); \
|
||||
ReleaseConvertTypes(converted_params); \
|
||||
ReleaseHugeMem releaseMemFunc = \
|
||||
reinterpret_cast<ReleaseHugeMem>(releaseMemAddr); \
|
||||
if (releaseMemFunc) { \
|
||||
releaseMemFunc(nullptr, false); \
|
||||
} \
|
||||
return api_ret; \
|
||||
}; \
|
||||
at_npu::native::OpCommand cmd; \
|
||||
cmd.Name(#aclnn_api); \
|
||||
cmd.SetCustomHandler(acl_call); \
|
||||
cmd.Run(); \
|
||||
if (unInitMemFunc) { \
|
||||
unInitMemFunc(nullptr, false); \
|
||||
} \
|
||||
} while (false)
|
||||
|
||||
#endif
|
||||
30
csrc/attention/CMakeLists.txt
Normal file
30
csrc/attention/CMakeLists.txt
Normal file
@@ -0,0 +1,30 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
|
||||
set(OPTEST_NAME optest_${PKG_NAME})
|
||||
|
||||
file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
|
||||
foreach(SUB_DIR ${CURRENT_DIRS})
|
||||
if (DEFINED ASCEND_OP_NAME AND NOT "${ASCEND_OP_NAME}" STREQUAL "")
|
||||
if (NOT "${ASCEND_OP_NAME}" STREQUAL "all" AND NOT "${ASCEND_OP_NAME}" STREQUAL "ALL")
|
||||
if (NOT ${SUB_DIR} IN_LIST ASCEND_OP_NAME)
|
||||
continue()
|
||||
endif ()
|
||||
endif ()
|
||||
endif ()
|
||||
|
||||
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
|
||||
add_subdirectory(${SUB_DIR})
|
||||
else()
|
||||
if (EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/op_host/CMakeLists.txt")
|
||||
add_subdirectory(${SUB_DIR}/op_host)
|
||||
endif()
|
||||
endif()
|
||||
endforeach()
|
||||
514
csrc/attention/common/op_kernel/CopyInL1.h
Normal file
514
csrc/attention/common/op_kernel/CopyInL1.h
Normal file
@@ -0,0 +1,514 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file CopyInL1.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef COPYINL1_H
|
||||
#define COPYINL1_H
|
||||
|
||||
enum class KVLAYOUT
|
||||
{
|
||||
BNBD, // [blockNums, headNum, blockSize, headDim]
|
||||
BBH, // [blockNums, blockSize, headNum * headDim]
|
||||
NZ // [blockNums, headNum, d1, blockSize, d0], d1 = headDim / d0, d0 = 32 (block byte) / sizeof(KV_T)
|
||||
};
|
||||
|
||||
struct CopyParam{
|
||||
uint32_t width;
|
||||
uint32_t height;
|
||||
uint32_t orgWidth;
|
||||
};
|
||||
|
||||
struct PAShape{
|
||||
uint32_t blockNum;
|
||||
uint32_t blockSize;
|
||||
uint32_t headNum; // 一般为kv的head num
|
||||
uint32_t headDim; // mla下rope为64, 非rope为512
|
||||
uint32_t maxblockNumPerBatch; // block table 每一行的最大个数
|
||||
uint32_t actHeadDim; // 实际拷贝col大小,考虑到N切块 s*d, 对应d
|
||||
uint32_t copyRowNum;
|
||||
uint32_t copyRowNumAlign;
|
||||
uint32_t pageStride;
|
||||
};
|
||||
|
||||
struct Position{
|
||||
uint32_t bIdx;
|
||||
uint32_t n2Idx;
|
||||
uint32_t s2Offset;
|
||||
uint32_t dIdx; // N轴被切,对应D轴被切
|
||||
};
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void GmCopyInToL1(LocalTensor<L1Type>& L1Tensor, GlobalTensor<L1Type>& GmTensor, const CopyParam& mmCopyParam)
|
||||
{
|
||||
Nd2NzParams Gm2L1Nd2NzParams;
|
||||
Gm2L1Nd2NzParams.ndNum = 1; // ND矩阵的个数
|
||||
Gm2L1Nd2NzParams.nValue = mmCopyParam.height; // 单个ND矩阵的实际行数,单位为元素个数
|
||||
Gm2L1Nd2NzParams.dValue = mmCopyParam.width; // 单个ND矩阵的实际列数(vD),单位为元素个数
|
||||
Gm2L1Nd2NzParams.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移, 单位为元素个数
|
||||
Gm2L1Nd2NzParams.srcDValue = mmCopyParam.orgWidth; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
|
||||
Gm2L1Nd2NzParams.dstNzC0Stride = (Gm2L1Nd2NzParams.nValue + 15) >> 4 << 4; // 转换为NZ矩阵后,相邻Block起始地址之间的偏移, 单位为Block个数
|
||||
Gm2L1Nd2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后,ND之间相邻两行在NZ矩阵中起始地址之间的偏移, 单位为Block个数
|
||||
Gm2L1Nd2NzParams.dstNzMatrixStride = 0; // 两个NZ矩阵,起始地址之间的偏移, 单位为元素数量
|
||||
DataCopy(L1Tensor, GmTensor, Gm2L1Nd2NzParams);
|
||||
}
|
||||
|
||||
// 场景:key、value GM to L1
|
||||
// GM按ND格式存储
|
||||
// L1按NZ格式存储
|
||||
// GM的行、列、列的stride(D or ND)BNSD 和 BSH的区别
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void DataCopyGmNDToL1(LocalTensor<L1Type>& l1Tensor, GlobalTensor<L1Type>& gmTensor,
|
||||
uint32_t rowAct,
|
||||
uint32_t rowAlign,
|
||||
uint32_t col, // D
|
||||
uint32_t colStride) // D or N*D
|
||||
{
|
||||
Nd2NzParams nd2nzPara;
|
||||
nd2nzPara.ndNum = 1;
|
||||
nd2nzPara.nValue = rowAct; // 行数
|
||||
|
||||
nd2nzPara.dValue = col;
|
||||
nd2nzPara.srcDValue = colStride;
|
||||
nd2nzPara.dstNzC0Stride = rowAlign;
|
||||
nd2nzPara.dstNzNStride = 1;
|
||||
nd2nzPara.srcNdMatrixStride = 0;
|
||||
nd2nzPara.dstNzMatrixStride = 0;
|
||||
DataCopy(l1Tensor, gmTensor, nd2nzPara);
|
||||
}
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void DataCopyGmScaleNDToL1(LocalTensor<L1Type>& l1Tensor, GlobalTensor<L1Type>& gmTensor,
|
||||
uint32_t rowAct,
|
||||
uint32_t rowAlign,
|
||||
uint32_t col, // D
|
||||
uint32_t colStride) // D or N*D
|
||||
{
|
||||
Nd2NzParams nd2nzPara;
|
||||
nd2nzPara.ndNum = 1;
|
||||
nd2nzPara.nValue = rowAct;
|
||||
|
||||
nd2nzPara.dValue = col;
|
||||
nd2nzPara.srcDValue = colStride;
|
||||
nd2nzPara.dstNzC0Stride = rowAlign;
|
||||
nd2nzPara.dstNzNStride = 1;
|
||||
nd2nzPara.srcNdMatrixStride = 0;
|
||||
nd2nzPara.dstNzMatrixStride = nd2nzPara.nValue;
|
||||
|
||||
LocalTensor<bfloat16_t> l1TensorCast = l1Tensor.template ReinterpretCast<bfloat16_t>();
|
||||
GlobalTensor<bfloat16_t> gmTensorCast;
|
||||
gmTensorCast.SetGlobalBuffer(((__gm__ bfloat16_t*)(gmTensor.GetPhyAddr())));
|
||||
DataCopy(l1TensorCast, gmTensorCast, nd2nzPara);
|
||||
}
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void DataCopyGmScaleDNToL1(LocalTensor<L1Type>& l1Tensor, GlobalTensor<L1Type>& gmTensor,
|
||||
uint32_t rowAct,
|
||||
uint32_t rowAlign,
|
||||
uint32_t col,
|
||||
uint32_t colStride)
|
||||
{
|
||||
Dn2NzParams dn2nzPara;
|
||||
dn2nzPara.dnNum = 1;
|
||||
dn2nzPara.nValue = col / 2;
|
||||
dn2nzPara.dValue = rowAct;
|
||||
dn2nzPara.srcDValue = colStride / 2;
|
||||
dn2nzPara.dstNzC0Stride = dn2nzPara.nValue;
|
||||
dn2nzPara.dstNzNStride = 1;
|
||||
dn2nzPara.srcDnMatrixStride = 0;
|
||||
dn2nzPara.dstNzMatrixStride = dn2nzPara.nValue;
|
||||
|
||||
LocalTensor<bfloat16_t> l1TensorCast = l1Tensor.template ReinterpretCast<bfloat16_t>();
|
||||
GlobalTensor<bfloat16_t> gmTensorCast;
|
||||
gmTensorCast.SetGlobalBuffer(((__gm__ bfloat16_t*)(gmTensor.GetPhyAddr())));
|
||||
DataCopy(l1TensorCast, gmTensorCast, dn2nzPara);
|
||||
}
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void DataCopyGmNZToL1(LocalTensor<L1Type>& l1Tensor, GlobalTensor<L1Type>& gmTensor,
|
||||
uint32_t rowAct,
|
||||
uint32_t dstRowStride,
|
||||
uint32_t srcRowStride,
|
||||
uint32_t col)
|
||||
{
|
||||
uint32_t blockElementCnt = 32U / sizeof(L1Type);
|
||||
if constexpr (IsSameType<L1Type, int4b_t>::value) {
|
||||
blockElementCnt = 64U;
|
||||
}
|
||||
DataCopyParams intriParams;
|
||||
intriParams.blockCount = col / blockElementCnt;
|
||||
intriParams.blockLen = rowAct;
|
||||
intriParams.dstStride = dstRowStride;
|
||||
intriParams.srcStride = srcRowStride;
|
||||
DataCopy(l1Tensor, gmTensor, intriParams);
|
||||
}
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void GmCopyInToL1HasRopePANoContinue(LocalTensor<L1Type>& nopeTensor, LocalTensor<L1Type>& ropeTensor,
|
||||
GlobalTensor<L1Type>& nopeGmTensor, GlobalTensor<L1Type>& ropeGmTensor,
|
||||
GlobalTensor<int32_t>& blockTableGm, KVLAYOUT kvLayout,
|
||||
const PAShape &shape,
|
||||
const PAShape &ropeShape,
|
||||
const Position &startPos)
|
||||
{
|
||||
uint32_t copyFinishRowCnt = 0;
|
||||
uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; // 块表的基偏移量
|
||||
uint32_t curS2Idx = startPos.s2Offset;
|
||||
uint32_t blockElementCnt = 32U / sizeof(L1Type); // 每个块的元素数量
|
||||
// ropeshape的M方向与nopeshape保持一样, 此处只判断nopeshape的
|
||||
while(copyFinishRowCnt < shape.copyRowNum){
|
||||
uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引
|
||||
uint64_t remainRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数
|
||||
uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上获取的编号
|
||||
//计算可以拷贝行数
|
||||
uint32_t copyRowCnt = shape.blockSize - remainRowCnt; // 一次只能处理一个Block
|
||||
if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum){
|
||||
copyRowCnt = shape.copyRowNum - copyFinishRowCnt; // 一个block未拷满
|
||||
}
|
||||
uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; // PA的偏移
|
||||
if (shape.pageStride > 0) {
|
||||
offset = idInBlockTable * shape.pageStride;
|
||||
}
|
||||
uint64_t keyRopeOffset = idInBlockTable * ropeShape.blockSize * ropeShape.headNum * ropeShape.headDim;
|
||||
if (ropeShape.pageStride > 0) {
|
||||
keyRopeOffset = idInBlockTable * ropeShape.pageStride;
|
||||
}
|
||||
|
||||
if (kvLayout == KVLAYOUT::NZ) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize;
|
||||
keyRopeOffset += static_cast<uint64_t>(startPos.n2Idx * ropeShape.blockSize * ropeShape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * ropeShape.blockSize;
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = nopeTensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = nopeGmTensor[offset];
|
||||
DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim);
|
||||
|
||||
LocalTensor<L1Type> tmpRopeDstTensor = ropeTensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpRopeSrcTensor = ropeGmTensor[keyRopeOffset];
|
||||
DataCopyGmNZToL1(tmpRopeDstTensor, tmpRopeSrcTensor, copyRowCnt, (ropeShape.copyRowNumAlign - copyRowCnt), (ropeShape.blockSize - copyRowCnt), ropeShape.actHeadDim);
|
||||
} else {
|
||||
uint64_t dStride = shape.headDim;
|
||||
uint64_t dRopeStride = ropeShape.headDim;
|
||||
if (kvLayout == KVLAYOUT::BBH) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx;
|
||||
keyRopeOffset += static_cast<uint64_t>(startPos.n2Idx * ropeShape.headDim) + remainRowCnt * ropeShape.headDim * ropeShape.headNum;
|
||||
dStride = shape.headDim * shape.headNum;
|
||||
dRopeStride = ropeShape.headDim * ropeShape.headNum;
|
||||
} else{
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx;
|
||||
keyRopeOffset += static_cast<uint64_t>(startPos.n2Idx * ropeShape.headDim * ropeShape.blockSize) + remainRowCnt * ropeShape.headDim;
|
||||
}
|
||||
|
||||
uint32_t dValue = shape.actHeadDim;
|
||||
uint32_t srcDValue = dStride;
|
||||
uint32_t dRopeValue = ropeShape.actHeadDim;
|
||||
uint32_t srcRopeDValue = dRopeStride;
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = nopeTensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = nopeGmTensor[offset];
|
||||
DataCopyGmNDToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue);
|
||||
|
||||
LocalTensor<L1Type> tmpRopeDstTensor = ropeTensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpRopeSrcTensor = ropeGmTensor[keyRopeOffset];
|
||||
DataCopyGmNDToL1(tmpRopeDstTensor, tmpRopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dRopeValue, srcRopeDValue);
|
||||
}
|
||||
copyFinishRowCnt += copyRowCnt;
|
||||
curS2Idx += copyRowCnt;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void GmCopyInToL1HasRopePA(LocalTensor<L1Type>& nopeTensor, LocalTensor<L1Type>& ropeTensor,
|
||||
GlobalTensor<L1Type>& nopeGmTensor, GlobalTensor<L1Type>& ropeGmTensor,
|
||||
GlobalTensor<int32_t>& blockTableGm, KVLAYOUT kvLayout,
|
||||
const PAShape &shape,
|
||||
const PAShape &ropeShape,
|
||||
const Position &startPos)
|
||||
{
|
||||
uint32_t copyFinishRowCnt = 0;
|
||||
uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; // 块表的基偏移量
|
||||
uint32_t curS2Idx = startPos.s2Offset;
|
||||
uint32_t blockElementCnt = 32U / sizeof(L1Type); // 每个块的元素数量
|
||||
// ropeshape的M方向与nopeshape保持一样, 此处只判断nopeshape的
|
||||
while(copyFinishRowCnt < shape.copyRowNum){
|
||||
uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引
|
||||
uint64_t remainRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数
|
||||
uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上获取的编号
|
||||
//计算可以拷贝行数
|
||||
uint32_t copyRowCnt = shape.blockSize - remainRowCnt; // 一次只能处理一个Block
|
||||
if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum){
|
||||
copyRowCnt = shape.copyRowNum - copyFinishRowCnt; // 一个block未拷满
|
||||
}
|
||||
uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; // PA的偏移
|
||||
uint64_t keyRopeOffset = idInBlockTable * ropeShape.blockSize * ropeShape.headNum * ropeShape.headDim;
|
||||
if (kvLayout == KVLAYOUT::NZ) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize;
|
||||
keyRopeOffset += static_cast<uint64_t>(startPos.n2Idx * ropeShape.blockSize * ropeShape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * ropeShape.blockSize;
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = nopeTensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = nopeGmTensor[offset];
|
||||
DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim);
|
||||
|
||||
LocalTensor<L1Type> tmpRopeDstTensor = ropeTensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpRopeSrcTensor = ropeGmTensor[keyRopeOffset];
|
||||
DataCopyGmNZToL1(tmpRopeDstTensor, tmpRopeSrcTensor, copyRowCnt, (ropeShape.copyRowNumAlign - copyRowCnt), (ropeShape.blockSize - copyRowCnt), ropeShape.actHeadDim);
|
||||
} else {
|
||||
uint64_t dStride = shape.headDim;
|
||||
uint64_t dRopeStride = ropeShape.headDim;
|
||||
if (kvLayout == KVLAYOUT::BBH) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx;
|
||||
keyRopeOffset += static_cast<uint64_t>(startPos.n2Idx * ropeShape.headDim) + remainRowCnt * ropeShape.headDim * ropeShape.headNum;
|
||||
dStride = shape.headDim * shape.headNum;
|
||||
dRopeStride = ropeShape.headDim * ropeShape.headNum;
|
||||
} else{
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx;
|
||||
keyRopeOffset += static_cast<uint64_t>(startPos.n2Idx * ropeShape.headDim * ropeShape.blockSize) + remainRowCnt * ropeShape.headDim;
|
||||
}
|
||||
|
||||
uint32_t dValue = shape.actHeadDim;
|
||||
uint32_t srcDValue = dStride;
|
||||
uint32_t dRopeValue = ropeShape.actHeadDim;
|
||||
uint32_t srcRopeDValue = dRopeStride;
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = nopeTensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = nopeGmTensor[offset];
|
||||
DataCopyGmNDToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue);
|
||||
|
||||
LocalTensor<L1Type> tmpRopeDstTensor = ropeTensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpRopeSrcTensor = ropeGmTensor[keyRopeOffset];
|
||||
DataCopyGmNDToL1(tmpRopeDstTensor, tmpRopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dRopeValue, srcRopeDValue);
|
||||
}
|
||||
copyFinishRowCnt += copyRowCnt;
|
||||
curS2Idx += copyRowCnt;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void GmCopyInToL1PA(LocalTensor<L1Type>& l1Tensor, GlobalTensor<L1Type>& gmTensor,
|
||||
GlobalTensor<int32_t>& blockTableGm, KVLAYOUT kvLayout,
|
||||
const PAShape &shape, const Position &startPos)
|
||||
{
|
||||
uint32_t copyFinishRowCnt = 0;
|
||||
uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch; // 块表的基偏移量
|
||||
uint32_t curS2Idx = startPos.s2Offset;
|
||||
uint32_t blockElementCnt = 32U / sizeof(L1Type); // 每个块的元素数量
|
||||
while(copyFinishRowCnt < shape.copyRowNum){
|
||||
uint64_t blockIdOffset = curS2Idx / shape.blockSize; // 获取block table上的索引
|
||||
uint64_t remainRowCnt = curS2Idx % shape.blockSize; // 获取在单个块上超出的行数
|
||||
uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset); // 从block table上获取的编号
|
||||
//计算可以拷贝行数
|
||||
uint32_t copyRowCnt = shape.blockSize - remainRowCnt; // 一次只能处理一个Block
|
||||
if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum){
|
||||
copyRowCnt = shape.copyRowNum - copyFinishRowCnt; // 一个block未拷满
|
||||
}
|
||||
uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim; // PA的偏移
|
||||
if (kvLayout == KVLAYOUT::NZ) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize;
|
||||
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = gmTensor[offset];
|
||||
DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim);
|
||||
} else {
|
||||
uint64_t dStride = shape.headDim;
|
||||
if (kvLayout == KVLAYOUT::BBH) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx;
|
||||
dStride = shape.headDim * shape.headNum;
|
||||
} else {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx;
|
||||
}
|
||||
|
||||
uint32_t dValue = shape.actHeadDim;
|
||||
uint32_t srcDValue = dStride;
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = gmTensor[offset];
|
||||
DataCopyGmNDToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, shape.copyRowNumAlign, dValue, srcDValue);
|
||||
}
|
||||
copyFinishRowCnt += copyRowCnt;
|
||||
curS2Idx += copyRowCnt;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void GmScaleCopyInToL1PAForND(LocalTensor<L1Type>& l1Tensor, GlobalTensor<L1Type>& gmTensor,
|
||||
GlobalTensor<int32_t>& blockTableGm, KVLAYOUT kvLayout,
|
||||
const PAShape &shape, const Position &startPos)
|
||||
{
|
||||
uint32_t copyFinishRowCnt = 0;
|
||||
uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch;
|
||||
uint32_t curS2Idx = startPos.s2Offset;
|
||||
constexpr uint32_t blockElementCnt = 32U / sizeof(L1Type);
|
||||
while(copyFinishRowCnt < shape.copyRowNum) {
|
||||
uint64_t blockIdOffset = curS2Idx / shape.blockSize;
|
||||
uint64_t remainRowCnt = curS2Idx % shape.blockSize;
|
||||
uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset);
|
||||
uint32_t copyRowCnt = shape.blockSize - remainRowCnt;
|
||||
if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum) {
|
||||
copyRowCnt = shape.copyRowNum - copyFinishRowCnt;
|
||||
}
|
||||
uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim;
|
||||
if (kvLayout == KVLAYOUT::NZ) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize;
|
||||
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = gmTensor[offset];
|
||||
DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim);
|
||||
} else {
|
||||
uint64_t dStride = shape.headDim;
|
||||
if (kvLayout == KVLAYOUT::BBH) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx;
|
||||
dStride = shape.headDim * shape.headNum;
|
||||
} else {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx;
|
||||
}
|
||||
|
||||
uint32_t dValue = shape.actHeadDim;
|
||||
uint32_t srcDValue = dStride;
|
||||
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = gmTensor[offset * 2];
|
||||
DataCopyGmScaleNDToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, copyRowCnt, dValue, srcDValue);
|
||||
}
|
||||
copyFinishRowCnt += copyRowCnt;
|
||||
curS2Idx += copyRowCnt;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename L1Type>
|
||||
__aicore__ inline void GmScaleCopyInToL1PAForDN(LocalTensor<L1Type>& l1Tensor, GlobalTensor<L1Type>& gmTensor,
|
||||
GlobalTensor<int32_t>& blockTableGm, KVLAYOUT kvLayout,
|
||||
const PAShape &shape, const Position &startPos)
|
||||
{
|
||||
uint32_t copyFinishRowCnt = 0;
|
||||
uint64_t blockTableBaseOffset = startPos.bIdx * shape.maxblockNumPerBatch;
|
||||
uint32_t curS2Idx = startPos.s2Offset;
|
||||
constexpr uint32_t blockElementCnt = 32U / sizeof(L1Type);
|
||||
while(copyFinishRowCnt < shape.copyRowNum) {
|
||||
uint64_t blockIdOffset = curS2Idx / shape.blockSize;
|
||||
uint64_t remainRowCnt = curS2Idx % shape.blockSize;
|
||||
uint64_t idInBlockTable = blockTableGm.GetValue(blockTableBaseOffset + blockIdOffset);
|
||||
uint32_t copyRowCnt = shape.blockSize - remainRowCnt;
|
||||
if (copyFinishRowCnt + copyRowCnt > shape.copyRowNum) {
|
||||
copyRowCnt = shape.copyRowNum - copyFinishRowCnt;
|
||||
}
|
||||
uint64_t offset = idInBlockTable * shape.blockSize * shape.headNum * shape.headDim;
|
||||
if (kvLayout == KVLAYOUT::NZ) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.blockSize * shape.headDim) + remainRowCnt * blockElementCnt + startPos.dIdx * shape.blockSize;
|
||||
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = gmTensor[offset];
|
||||
DataCopyGmNZToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, (shape.copyRowNumAlign - copyRowCnt), (shape.blockSize - copyRowCnt), shape.actHeadDim);
|
||||
} else {
|
||||
uint64_t dStride = shape.headDim;
|
||||
if (kvLayout == KVLAYOUT::BBH) {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim) + remainRowCnt * shape.headDim * shape.headNum + startPos.dIdx;
|
||||
dStride = shape.headDim * shape.headNum;
|
||||
} else {
|
||||
offset += static_cast<uint64_t>(startPos.n2Idx * shape.headDim * shape.blockSize) + remainRowCnt * shape.headDim + startPos.dIdx;
|
||||
}
|
||||
|
||||
uint32_t dValue = shape.actHeadDim;
|
||||
uint32_t srcDValue = dStride;
|
||||
LocalTensor<L1Type> tmpNopeDstTensor = l1Tensor[copyFinishRowCnt * blockElementCnt];
|
||||
GlobalTensor<L1Type> tmpNopeSrcTensor = gmTensor[offset];
|
||||
|
||||
DataCopyGmScaleDNToL1(tmpNopeDstTensor, tmpNopeSrcTensor, copyRowCnt, copyRowCnt, dValue, srcDValue);
|
||||
}
|
||||
copyFinishRowCnt += copyRowCnt;
|
||||
curS2Idx += copyRowCnt;
|
||||
}
|
||||
}
|
||||
|
||||
template<typename INPUT_T>
|
||||
__aicore__ inline void CopyToL1Nd2Nz(const LocalTensor<INPUT_T> &l1Tensor, const GlobalTensor<INPUT_T> &gmTensor,
|
||||
uint32_t nValue, uint32_t dValue, uint32_t srcDValue)
|
||||
{
|
||||
Nd2NzParams gm2L1Nd2NzParams;
|
||||
gm2L1Nd2NzParams.ndNum = 1; // ND矩阵的个数
|
||||
gm2L1Nd2NzParams.nValue = nValue; // 单个ND矩阵的实际行数,单位为元素个数
|
||||
gm2L1Nd2NzParams.dValue = dValue; // 单个ND矩阵的实际列数,单位为元素个数
|
||||
gm2L1Nd2NzParams.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移, 单位为元素个数
|
||||
gm2L1Nd2NzParams.srcDValue = srcDValue; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
|
||||
#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__NPU_ARCH__ == 5102)
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value || IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value || IsSameType<INPUT_T, int8_t>::value) {
|
||||
gm2L1Nd2NzParams.dstNzC0Stride = (nValue + 31) >> 5 << 5;
|
||||
} else {
|
||||
gm2L1Nd2NzParams.dstNzC0Stride = (nValue + 15) >> 4 << 4;
|
||||
}
|
||||
#else
|
||||
gm2L1Nd2NzParams.dstNzC0Stride = (nValue + 15) >> 4 << 4; // NZ矩阵相邻Block起始地址之间的偏移, 单位为Block个数
|
||||
#endif
|
||||
gm2L1Nd2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后,ND之间相邻两行在NZ矩阵中起始地址之间的偏移, 单位为Block个数
|
||||
gm2L1Nd2NzParams.dstNzMatrixStride = 0; // 两个NZ矩阵,起始地址之间的偏移, 单位为元素数量
|
||||
DataCopy(l1Tensor, gmTensor, gm2L1Nd2NzParams);
|
||||
}
|
||||
|
||||
template<typename INPUT_T>
|
||||
__aicore__ inline void CopyScaleToL1Nd2Nz(const LocalTensor<INPUT_T> &l1Tensor, const GlobalTensor<INPUT_T> &gmTensor,
|
||||
uint32_t nValue, uint32_t dValue, uint32_t srcDValue)
|
||||
{
|
||||
Nd2NzParams gm2L1Nd2NzParams;
|
||||
gm2L1Nd2NzParams.ndNum = 1; // ND矩阵的个数
|
||||
gm2L1Nd2NzParams.nValue = nValue / 2; // 单个ND矩阵的实际行数,单位为元素个数
|
||||
gm2L1Nd2NzParams.dValue = dValue; // 单个ND矩阵的实际列数,单位为元素个数
|
||||
gm2L1Nd2NzParams.srcNdMatrixStride = 0; // 相邻ND矩阵起始地址之间的偏移, 单位为元素个数
|
||||
gm2L1Nd2NzParams.srcDValue = srcDValue; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
|
||||
gm2L1Nd2NzParams.dstNzC0Stride = nValue / 2; // NZ矩阵相邻Block起始地址之间的偏移, 单位为Block个数
|
||||
gm2L1Nd2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后,ND之间相邻两行在NZ矩阵中起始地址之间的偏移, 单位为Block个数
|
||||
gm2L1Nd2NzParams.dstNzMatrixStride = gm2L1Nd2NzParams.nValue; // 两个NZ矩阵,起始地址之间的偏移, 单位为元素数量
|
||||
|
||||
LocalTensor<bfloat16_t> l1TensorCast = l1Tensor.template ReinterpretCast<bfloat16_t>();
|
||||
GlobalTensor<bfloat16_t> gmTensorCast;
|
||||
gmTensorCast.SetGlobalBuffer(((__gm__ bfloat16_t*)(gmTensor.GetPhyAddr())));
|
||||
DataCopy(l1TensorCast, gmTensorCast, gm2L1Nd2NzParams);
|
||||
}
|
||||
|
||||
template<typename INPUT_T>
|
||||
__aicore__ inline void CopyScaleToL1Dn2Nz(const LocalTensor<INPUT_T> &l1Tensor, const GlobalTensor<INPUT_T> &gmTensor,
|
||||
uint32_t nValue, uint32_t dValue, uint32_t srcDValue)
|
||||
{
|
||||
Dn2NzParams gm2L1Dn2NzParams;
|
||||
gm2L1Dn2NzParams.dnNum = 1; // ND矩阵的个数
|
||||
gm2L1Dn2NzParams.nValue = nValue / 2; // 单个DN矩阵的实际列数,单位为元素个数
|
||||
gm2L1Dn2NzParams.dValue = dValue; // 单个DN矩阵的实际行数,单位为元素个数
|
||||
gm2L1Dn2NzParams.srcDnMatrixStride = 0; // 相邻Dn矩阵起始地址之间的偏移, 单位为元素个数
|
||||
gm2L1Dn2NzParams.srcDValue = srcDValue / 2; // 同一个Dn矩阵中相邻行起始地址之间的偏移, 单位为元素个数
|
||||
gm2L1Dn2NzParams.dstNzC0Stride = nValue / 2;
|
||||
gm2L1Dn2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后,ND之间相邻两行在NZ矩阵中起始地址之间的偏移, 单位为Block个数
|
||||
gm2L1Dn2NzParams.dstNzMatrixStride = gm2L1Dn2NzParams.nValue; // 两个NZ矩阵,起始地址之间的偏移, 单位为元素数量
|
||||
|
||||
LocalTensor<bfloat16_t> l1TensorCast = l1Tensor.template ReinterpretCast<bfloat16_t>();
|
||||
GlobalTensor<bfloat16_t> gmTensorCast;
|
||||
gmTensorCast.SetGlobalBuffer(((__gm__ bfloat16_t*)(gmTensor.GetPhyAddr())));
|
||||
DataCopy(l1TensorCast, gmTensorCast, gm2L1Dn2NzParams);
|
||||
}
|
||||
|
||||
template<typename INPUT_T>
|
||||
__aicore__ inline void CopyToL1Nd2NzGS1Merge(const LocalTensor<INPUT_T> &l1Tensor, const GlobalTensor<INPUT_T> &gmTensor,
|
||||
uint32_t ndNum, uint32_t nValue, uint32_t dValue, uint32_t srcNdMatrixStride, uint32_t srcDValue, uint32_t dstNzC0Stride) // BSNGD 合轴拷贝
|
||||
{
|
||||
Nd2NzParams gm2L1Nd2NzParams;
|
||||
gm2L1Nd2NzParams.ndNum = ndNum; // ND矩阵的个数
|
||||
gm2L1Nd2NzParams.nValue = nValue; // 单个ND矩阵的实际行数,单位为元素个数
|
||||
gm2L1Nd2NzParams.dValue = dValue; // 单个ND矩阵的实际列数,单位为元素个数
|
||||
gm2L1Nd2NzParams.srcNdMatrixStride = srcNdMatrixStride; // 相邻ND矩阵起始地址之间的偏移, 单位为元素个数
|
||||
gm2L1Nd2NzParams.srcDValue = srcDValue; // 同一个ND矩阵中相邻行起始地址之间的偏移, 单位为元素个数
|
||||
#if (__CCE_AICORE__ == 310) || (defined __DAV_310R6__) || (__NPU_ARCH__ == 5102)
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value || IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value || IsSameType<INPUT_T, int8_t>::value) {
|
||||
gm2L1Nd2NzParams.dstNzC0Stride = (dstNzC0Stride + 31) >> 5 << 5; // NZ矩阵相邻Block起始地址之间的偏移,单位为Block个数,32对齐
|
||||
} else {
|
||||
gm2L1Nd2NzParams.dstNzC0Stride = (dstNzC0Stride + 15) >> 4 << 4; // NZ矩阵相邻Block起始地址之间的偏移,单位为Block个数,16对齐
|
||||
}
|
||||
#else
|
||||
gm2L1Nd2NzParams.dstNzC0Stride = (dstNzC0Stride + 15) >> 4 << 4; // NZ矩阵相邻Block起始地址之间的偏移,单位为Block个数,16对齐
|
||||
#endif
|
||||
gm2L1Nd2NzParams.dstNzNStride = 1; // 转换为NZ矩阵后,ND之间相邻两行在NZ矩阵中起始地址之间的偏移, 单位为Block个数
|
||||
gm2L1Nd2NzParams.dstNzMatrixStride = nValue * 32 / sizeof(INPUT_T); // 两个NZ矩阵,起始地址之间的偏移, 单位为元素数量
|
||||
DataCopy(l1Tensor, gmTensor, gm2L1Nd2NzParams);
|
||||
}
|
||||
#endif
|
||||
56
csrc/attention/common/op_kernel/FixpipeOut.h
Normal file
56
csrc/attention/common/op_kernel/FixpipeOut.h
Normal file
@@ -0,0 +1,56 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file FixpipeOut.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef FIXPIPEOUT_H
|
||||
#define FIXPIPEOUT_H
|
||||
|
||||
constexpr FixpipeConfig PFA_CFG_ROW_MAJOR_UB = {CO2Layout::ROW_MAJOR, true}; // ROW_MAJOR: 使能NZ2ND,输出数据格式为ND格式; true: 用于用户指定目的地址的位置是否是UB
|
||||
constexpr FixpipeConfig PFA_CFG_ROW_MAJOR_GM = {CO2Layout::ROW_MAJOR, false}; // ROW_MAJOR: 使能NZ2ND,输出数据格式为ND格式; true: 用于用户指定目的地址的位置是否是UB
|
||||
constexpr FixpipeConfig FA_CFG_NZ_UB = {CO2Layout::NZ, true}; // 不使能NZ2ND,输出数据格式为NZ格式; true: 用于用户指定目的地址的位置是否是UB
|
||||
|
||||
struct fixpipeOutParams {
|
||||
uint32_t fixpOutMSize;
|
||||
uint32_t fixpOutNSize;
|
||||
};
|
||||
|
||||
template<typename mmOutputType, typename computeType, typename l0cType>
|
||||
__aicore__ inline void FixpipeMmCopyOutToUB(LocalTensor<mmOutputType>& mmResUb, LocalTensor<l0cType>& L0CTensor, const fixpipeOutParams& fixpOutParam)
|
||||
{
|
||||
FixpipeParamsC310<CO2Layout::ROW_MAJOR> L0C2UbFixpParams; // L0C->UB
|
||||
L0C2UbFixpParams.nSize = (fixpOutParam.fixpOutNSize + 7) >> 3 << 3; // L0C上的bmm1结果矩阵N方向的size大小;同mmadParams.n;8个元素(32B)对齐
|
||||
L0C2UbFixpParams.mSize = (fixpOutParam.fixpOutMSize + 1) >> 1 << 1; // 有效数据不足16行,只需输出部分行即可;L0C上的bmm1结果矩阵M方向的size大小必须是偶数
|
||||
L0C2UbFixpParams.srcStride = ((L0C2UbFixpParams.mSize + 15) >> 4) << 4; // L0C上matmul结果相邻连续数据片断间隔(前面一个数据块的头与后面数据块的头的间隔),单位为16 *sizeof(T) //源NZ矩阵中相邻Z排布的起始地址偏移
|
||||
L0C2UbFixpParams.dstStride = (L0C2UbFixpParams.nSize + 15) >> 4 << 4; // mmResUb上两行之间的间隔,单位:element。 // 128:根据比对dump文件得到,ND方案(S1 * S2)时脏数据用mask剔除
|
||||
L0C2UbFixpParams.dualDstCtl = 1; // 双目标模式,按M维度拆分, M / 2 * N写入每个UB,M必须为2的倍数
|
||||
L0C2UbFixpParams.params.ndNum = 1;
|
||||
L0C2UbFixpParams.params.srcNdStride = 0;
|
||||
L0C2UbFixpParams.params.dstNdStride = 0;
|
||||
Fixpipe<mmOutputType, computeType, PFA_CFG_ROW_MAJOR_UB>(mmResUb, L0CTensor, L0C2UbFixpParams); // 将matmul结果从L0C搬运到UB
|
||||
}
|
||||
|
||||
template<typename mmOutputType, typename computeType, typename l0cType>
|
||||
__aicore__ inline void FixpipeMmCopyOutToGm(GlobalTensor<mmOutputType>& mmResGm,LocalTensor<l0cType>& L0CTensor, const fixpipeOutParams& fixpOutParam)
|
||||
{
|
||||
FixpipeParamsC310<CO2Layout::ROW_MAJOR> L0C2GmFixpParams; // L0C->Gm
|
||||
L0C2GmFixpParams.nSize = (fixpOutParam.fixpOutNSize + 7) >> 3 << 3; // L0C上的bmm1结果矩阵N方向的size大小;同mmadParams.n;8个元素(32B)对齐;分档计算且vector1中通过mask筛选出实际有效值
|
||||
L0C2GmFixpParams.mSize = (fixpOutParam.fixpOutMSize + 1) >> 1 << 1; // 有效数据不足16行,只需输出部分行即可;L0C上的bmm1结果矩阵M方向的size大小;同mmadParams.m
|
||||
L0C2GmFixpParams.srcStride = ((L0C2GmFixpParams.mSize + 15) >> 4) << 4; // L0C上bmm1结果相邻连续数据片断间隔(前面一个数据块的头与后面数据块的头的间隔)
|
||||
L0C2GmFixpParams.dstStride = (L0C2GmFixpParams.nSize + 15) >> 4 << 4; // mmResGm上两行之间的间隔
|
||||
L0C2GmFixpParams.dualDstCtl = 1;
|
||||
L0C2GmFixpParams.params.ndNum = 1;
|
||||
L0C2GmFixpParams.params.srcNdStride = 0;
|
||||
L0C2GmFixpParams.params.dstNdStride = 0;
|
||||
Fixpipe<mmOutputType, computeType, PFA_CFG_ROW_MAJOR_GM>(mmResGm, L0CTensor, L0C2GmFixpParams); // 将matmul结果从L0C搬运到Gm
|
||||
}
|
||||
#endif
|
||||
@@ -0,0 +1,139 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_basic_block_aligned128_no_update_sfa.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef VF_BASIC_BLOCK_ALIGNED128_NO_UPDATE_SFA_H
|
||||
#define VF_BASIC_BLOCK_ALIGNED128_NO_UPDATE_SFA_H
|
||||
|
||||
#include "vf_basic_block_utils.h"
|
||||
|
||||
using namespace regbaseutil;
|
||||
|
||||
namespace FaVectorApi {
|
||||
// no update, originN == 128
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__simd_vf__ void ProcessVec1NoUpdateImpl128VF(
|
||||
__ubuf__ T2 * expUb, __ubuf__ T * expSumUb, __ubuf__ T * maxUb, __ubuf__ T * maxUbStart,
|
||||
__ubuf__ T * srcUb, const uint32_t blockStride, const uint32_t repeatStride,
|
||||
const uint16_t m, const T scale, const T minValue)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x_unroll;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_tmp;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_brc;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_sum;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_even;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_odd;
|
||||
|
||||
// bfloat16_t
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_even_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_odd_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_bf16;
|
||||
// half
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_even_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_odd_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_fp16;
|
||||
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_max;
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum;
|
||||
|
||||
AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask<T, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_all_b16 =
|
||||
AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x_unroll, srcUb + floatRepSize + i * s2BaseSize);
|
||||
|
||||
AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_all); // Muls(scale)
|
||||
AscendC::MicroAPI::Muls(vreg_input_x_unroll, vreg_input_x_unroll, scale, preg_all);
|
||||
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_all);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + floatRepSize + i * s2BaseSize, vreg_input_x_unroll, preg_all);
|
||||
AscendC::MicroAPI::Max(vreg_max_tmp, vreg_input_x, vreg_input_x_unroll, preg_all);
|
||||
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::MAX, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_input_max, vreg_max_tmp, preg_all);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)maxUb), vreg_input_max, ureg_max, 1);
|
||||
}
|
||||
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)maxUb), ureg_max, 0);
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
// maxUb is [S1, 1], BRC_B32 is reading one fp32 element and broadcast it to all 64 vreg element
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(
|
||||
vreg_max_brc, maxUbStart + i);
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_DINTLV_B32>(
|
||||
vreg_input_x, vreg_input_x_unroll, srcUb + i * s2BaseSize);
|
||||
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp_even, vreg_input_x, vreg_max_brc, preg_all);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp_odd, vreg_input_x_unroll, vreg_max_brc, preg_all);
|
||||
|
||||
// x_sum = sum(x_exp, axis=-1, keepdims=True)
|
||||
AscendC::MicroAPI::Add(vreg_exp_sum, vreg_exp_even, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::SUM, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_exp_sum, vreg_exp_sum, preg_all);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)expSumUb), vreg_exp_sum, ureg_exp_sum, 1);
|
||||
|
||||
if constexpr (IsSameType<T2, bfloat16_t>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_even_bf16, vreg_exp_even, preg_all);
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitOne>(vreg_exp_odd_bf16, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Or((RegTensor<uint16_t>&)vreg_exp_bf16, (RegTensor<uint16_t>&)vreg_exp_even_bf16,
|
||||
(RegTensor<uint16_t>&)vreg_exp_odd_bf16, preg_all_b16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_exp_bf16, blockStride, repeatStride, preg_all_b16);
|
||||
} else if constexpr (IsSameType<T2, half>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_even_fp16, vreg_exp_even, preg_all);
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitOne>(vreg_exp_odd_fp16, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Or((RegTensor<uint16_t>&)vreg_exp_fp16, (RegTensor<uint16_t>&)vreg_exp_even_fp16,
|
||||
(RegTensor<uint16_t>&)vreg_exp_odd_fp16, preg_all_b16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_exp_fp16, blockStride, repeatStride, preg_all_b16);
|
||||
}
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)expSumUb), ureg_exp_sum, 0);
|
||||
}
|
||||
|
||||
// no update, originN == 128
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__aicore__ inline void ProcessVec1NoUpdateImpl128(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor,
|
||||
const LocalTensor<T>& expSumTensor, const LocalTensor<T>& maxTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
// 写的时候固定用65或者33的stride去写,因为正向目前使能settail之后mm2的s1方向必须算满128或者64行
|
||||
// stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1)
|
||||
const uint32_t blockStride = s1BaseSize >> 1 | 0x1;
|
||||
const uint32_t repeatStride = 1;
|
||||
__ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ T * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr();
|
||||
__ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr();
|
||||
__ubuf__ T * maxUbStart = (__ubuf__ T*)maxTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr();
|
||||
|
||||
ProcessVec1NoUpdateImpl128VF<T, T2, s1BaseSize, s2BaseSize>(
|
||||
expUb, expSumUb, maxUb, maxUbStart, srcUb, blockStride, repeatStride, m, scale, minValue);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
#endif // VF_BASIC_BLOCK_ALIGNED128_NO_UPDATE_SFA_H
|
||||
@@ -0,0 +1,143 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_basic_block_aligned128_update_sfa.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef VF_BASIC_BLOCK_ALIGNED128_UPDATE_SFA_H
|
||||
#define VF_BASIC_BLOCK_ALIGNED128_UPDATE_SFA_H
|
||||
|
||||
#include "vf_basic_block_utils.h"
|
||||
|
||||
using namespace regbaseutil;
|
||||
|
||||
namespace FaVectorApi {
|
||||
// update, originN == 128
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 128, uint32_t s2BaseSize = 128>
|
||||
__simd_vf__ void ProcessVec1UpdateImpl128VF(
|
||||
__ubuf__ T2 * expUb, __ubuf__ T * srcUb, __ubuf__ T * inMaxUb,
|
||||
__ubuf__ T * tmpExpSumUb, __ubuf__ T * tmpMaxUb, __ubuf__ T * tmpMaxUb2, const uint32_t blockStride,
|
||||
const uint32_t repeatStride, const uint16_t m, const T scale, const T minValue)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x_unroll;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_tmp;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_in_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_new;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_brc;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_cur_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_sum;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_in_exp_sum;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_even;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_odd;
|
||||
|
||||
// bfloat16_t
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_even_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_odd_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_bf16;
|
||||
// half
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_even_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_odd_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_fp16;
|
||||
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_max;
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum;
|
||||
AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_all_b16 =
|
||||
AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
// x_max = max(src, axis=-1, keepdims=True); x_max = Max(x_max, inMax)
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x_unroll, srcUb + floatRepSize + i * s2BaseSize);
|
||||
|
||||
AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_all); // Muls(scale)
|
||||
AscendC::MicroAPI::Muls(vreg_input_x_unroll, vreg_input_x_unroll, scale, preg_all);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_all);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + floatRepSize + i * s2BaseSize, vreg_input_x_unroll, preg_all);
|
||||
AscendC::MicroAPI::Max(vreg_max_tmp, vreg_input_x, vreg_input_x_unroll, preg_all);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::MAX, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_max_tmp, vreg_max_tmp, preg_all);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpMaxUb), vreg_max_tmp, ureg_max, 1);
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpMaxUb), ureg_max, 0);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_in_max, inMaxUb);
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
AscendC::MicroAPI::LoadAlign(vreg_cur_max, tmpMaxUb2); // 获取新的max[s1, 1]
|
||||
AscendC::MicroAPI::Max(vreg_max_new, vreg_cur_max, vreg_in_max, preg_all); // 计算新、旧max的最大值
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)tmpMaxUb2, vreg_max_new, preg_all);
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_max_brc, tmpMaxUb2 + i);
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_DINTLV_B32>(
|
||||
vreg_input_x, vreg_input_x_unroll, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp_even, vreg_input_x, vreg_max_brc, preg_all);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp_odd, vreg_input_x_unroll, vreg_max_brc, preg_all);
|
||||
|
||||
// x_sum = sum(x_exp, axis=-1, keepdims=True)
|
||||
AscendC::MicroAPI::Add(vreg_exp_sum, vreg_exp_even, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::SUM, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_exp_sum, vreg_exp_sum, preg_all);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpExpSumUb), vreg_exp_sum, ureg_exp_sum, 1);
|
||||
|
||||
if constexpr (IsSameType<T2, bfloat16_t>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_even_bf16, vreg_exp_even, preg_all);
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitOne>(vreg_exp_odd_bf16, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Or((RegTensor<uint16_t>&)vreg_exp_bf16, (RegTensor<uint16_t>&)vreg_exp_even_bf16,
|
||||
(RegTensor<uint16_t>&)vreg_exp_odd_bf16, preg_all_b16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_exp_bf16, blockStride, repeatStride, preg_all_b16);
|
||||
} else if constexpr (IsSameType<T2, half>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_even_fp16, vreg_exp_even, preg_all);
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitOne>(vreg_exp_odd_fp16, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Or((RegTensor<uint16_t>&)vreg_exp_fp16, (RegTensor<uint16_t>&)vreg_exp_even_fp16,
|
||||
(RegTensor<uint16_t>&)vreg_exp_odd_fp16, preg_all_b16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_exp_fp16, blockStride, repeatStride, preg_all_b16);
|
||||
}
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpExpSumUb), ureg_exp_sum, 0);
|
||||
}
|
||||
|
||||
// update, originN == 128
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__aicore__ inline void ProcessVec1UpdateImpl128(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
// 写的时候固定用65或者33的stride去写,因为正向目前使能settail之后mm2的s1方向必须算满128或者64行
|
||||
// stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1)
|
||||
const uint32_t blockStride = s1BaseSize >> 1 | 0x1;
|
||||
const uint32_t repeatStride = 1;
|
||||
|
||||
__ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr();
|
||||
__ubuf__ T * inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr();
|
||||
__ubuf__ T * tmpExpSumUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr();
|
||||
__ubuf__ T * tmpMaxUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64;
|
||||
__ubuf__ T * tmpMaxUb2 = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64;
|
||||
|
||||
ProcessVec1UpdateImpl128VF <T, T2, s1BaseSize, s2BaseSize>(
|
||||
expUb, srcUb, inMaxUb, tmpExpSumUb, tmpMaxUb, tmpMaxUb2, blockStride, repeatStride, m, scale, minValue);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
#endif // VF_BASIC_BLOCK_ALIGNED128_UPDATE_SFA_H
|
||||
@@ -0,0 +1,149 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_basic_block_unaligned128_no_update_sfa.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef VF_BASIC_BLOCK_UNALIGNED128_NO_UPDATE_SFA_H
|
||||
#define VF_BASIC_BLOCK_UNALIGNED128_NO_UPDATE_SFA_H
|
||||
|
||||
#include "vf_basic_block_utils.h"
|
||||
|
||||
using namespace regbaseutil;
|
||||
|
||||
namespace FaVectorApi {
|
||||
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__simd_vf__ void ProcessVec1NoUpdateGeneralImpl128VF(
|
||||
__ubuf__ T2 * expUb, __ubuf__ T * expSumUb, __ubuf__ T * maxUb, __ubuf__ T * maxUbStart,
|
||||
__ubuf__ T * srcUb, const uint32_t blockStride, const uint32_t repeatStride,
|
||||
const uint16_t m, const T scale, const T minValue, uint32_t pltOriTailN, uint32_t pltTailN)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_min;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x_unroll;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x_unroll_new;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_tmp;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_brc;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_sum;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_even;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_odd;
|
||||
|
||||
// bfloat16_t
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_even_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_odd_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_bf16;
|
||||
// half
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_even_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_odd_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_fp16;
|
||||
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_max;
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum;
|
||||
|
||||
AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_all_b16 =
|
||||
AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_all_b8 = AscendC::MicroAPI::CreateMask<T2, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_tail_n = AscendC::MicroAPI::UpdateMask<float>(pltTailN);
|
||||
AscendC::MicroAPI::MaskReg preg_ori_tail_n = AscendC::MicroAPI::UpdateMask<float>(pltOriTailN);
|
||||
AscendC::MicroAPI::MaskReg preg_reduce_n =
|
||||
AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::VL8>();
|
||||
|
||||
AscendC::MicroAPI::Duplicate(vreg_min, minValue);
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x_unroll, srcUb + floatRepSize + i * s2BaseSize);
|
||||
AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_all); // Muls(scale)
|
||||
AscendC::MicroAPI::Muls(vreg_input_x_unroll, vreg_input_x_unroll, scale, preg_ori_tail_n);
|
||||
AscendC::MicroAPI::Select(vreg_input_x_unroll_new, vreg_input_x_unroll, vreg_min, preg_ori_tail_n);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_all);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + floatRepSize + i * s2BaseSize, vreg_input_x_unroll_new, preg_tail_n);
|
||||
|
||||
AscendC::MicroAPI::Max(vreg_max_tmp, vreg_input_x, vreg_input_x_unroll_new, preg_all);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::MAX, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_input_max, vreg_max_tmp, preg_all);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)maxUb), vreg_input_max, ureg_max, 1);
|
||||
}
|
||||
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)maxUb), ureg_max, 0);
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_max_brc, maxUbStart + i);
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_DINTLV_B32>(
|
||||
vreg_input_x, vreg_input_x_unroll, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp_even, vreg_input_x, vreg_max_brc, preg_all);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp_odd, vreg_input_x_unroll, vreg_max_brc, preg_all);
|
||||
|
||||
// x_sum = sum(x_exp, axis=-1, keepdims=True)
|
||||
AscendC::MicroAPI::Add(vreg_exp_sum, vreg_exp_even, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::SUM, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_exp_sum, vreg_exp_sum, preg_all);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)expSumUb), vreg_exp_sum, ureg_exp_sum, 1);
|
||||
|
||||
if constexpr (IsSameType<T2, bfloat16_t>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_even_bf16, vreg_exp_even, preg_all);
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitOne>(vreg_exp_odd_bf16, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Or((RegTensor<uint16_t>&)vreg_exp_bf16, (RegTensor<uint16_t>&)vreg_exp_even_bf16,
|
||||
(RegTensor<uint16_t>&)vreg_exp_odd_bf16, preg_all_b16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_exp_bf16, blockStride, repeatStride, preg_all_b16);
|
||||
} else if constexpr (IsSameType<T2, half>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_even_fp16, vreg_exp_even, preg_all);
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitOne>(vreg_exp_odd_fp16, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Or((RegTensor<uint16_t>&)vreg_exp_fp16, (RegTensor<uint16_t>&)vreg_exp_even_fp16,
|
||||
(RegTensor<uint16_t>&)vreg_exp_odd_fp16, preg_all_b16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_exp_fp16, blockStride, repeatStride, preg_all_b16);
|
||||
}
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)expSumUb), ureg_exp_sum, 0);
|
||||
}
|
||||
|
||||
// no update, 64 < originN <= 128
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__aicore__ inline void ProcessVec1NoUpdateGeneralImpl128(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor,
|
||||
const LocalTensor<T>& expSumTensor, const LocalTensor<T>& maxTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
// 写的时候固定用65或者33的stride去写,因为正向目前使能settail之后mm2的s1方向必须算满128或者64行
|
||||
// stride, high 16bits: blockStride (65*16*2/32),单位block, low 16bits: repeatStride (1)
|
||||
const uint32_t blockStride = s1BaseSize >> 1 | 0x1;
|
||||
const uint32_t repeatStride = 1;
|
||||
__ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ T * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr();
|
||||
__ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr();
|
||||
__ubuf__ T * maxUbStart = (__ubuf__ T*)maxTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr();
|
||||
|
||||
const uint32_t oriTailN = originN - floatRepSize;
|
||||
const uint32_t tailN = s2BaseSize - floatRepSize;
|
||||
uint32_t pltOriTailN = oriTailN;
|
||||
uint32_t pltTailN = tailN;
|
||||
|
||||
ProcessVec1NoUpdateGeneralImpl128VF<T, T2, s1BaseSize, s2BaseSize>(
|
||||
expUb, expSumUb, maxUb, maxUbStart, srcUb, blockStride, repeatStride, m, scale, minValue,
|
||||
pltOriTailN, pltTailN);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
#endif // VF_BASIC_BLOCK_UNALIGNED128_NO_UPDATE_SFA_H
|
||||
@@ -0,0 +1,159 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_basic_block_unaligned128_update_sfa.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef VF_BASIC_BLOCK_UNALIGNED128_UPDATE_SFA_H
|
||||
#define VF_BASIC_BLOCK_UNALIGNED128_UPDATE_SFA_H
|
||||
|
||||
#include "vf_basic_block_utils.h"
|
||||
|
||||
using namespace regbaseutil;
|
||||
|
||||
namespace FaVectorApi {
|
||||
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__simd_vf__ void ProcessVec1UpdateGeneralImpl128VF(
|
||||
__ubuf__ T2 * expUb, __ubuf__ T * srcUb, __ubuf__ T * inMaxUb,
|
||||
__ubuf__ T * tmpExpSumUb, __ubuf__ T * tmpMaxUb, __ubuf__ T * tmpMaxUb2, const uint32_t blockStride,
|
||||
const uint32_t repeatStride, const uint16_t m, const T scale, const T minValue, uint32_t pltOriTailN,
|
||||
uint32_t pltTailN, uint32_t pltN)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_min;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x_unroll;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x_unroll_new;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_tmp;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_cur_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_new;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_sum;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_in_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_brc;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_even;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_odd;
|
||||
|
||||
// bfloat16_t
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_even_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_odd_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_bf16;
|
||||
// half
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_even_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_odd_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_fp16;
|
||||
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_max;
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum;
|
||||
|
||||
AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_all_b16 = AscendC::MicroAPI::CreateMask<uint16_t,
|
||||
AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_n_b16 = AscendC::MicroAPI::UpdateMask<uint16_t>(pltN);
|
||||
AscendC::MicroAPI::MaskReg preg_tail_n = AscendC::MicroAPI::UpdateMask<T>(pltTailN);
|
||||
AscendC::MicroAPI::MaskReg preg_ori_tail_n = AscendC::MicroAPI::UpdateMask<T>(pltOriTailN);
|
||||
|
||||
AscendC::MicroAPI::Duplicate(vreg_min, minValue);
|
||||
// x_max = max(src, axis=-1, keepdims=True); x_max = Max(x_max, inMax)
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x_unroll, srcUb + floatRepSize + i * s2BaseSize);
|
||||
AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_all); // Muls(scale)
|
||||
AscendC::MicroAPI::Muls(vreg_input_x_unroll, vreg_input_x_unroll, scale, preg_ori_tail_n);
|
||||
AscendC::MicroAPI::Select(vreg_input_x_unroll_new, vreg_input_x_unroll, vreg_min, preg_ori_tail_n);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_all);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + floatRepSize + i * s2BaseSize, vreg_input_x_unroll_new, preg_tail_n);
|
||||
AscendC::MicroAPI::Max(vreg_max_tmp, vreg_input_x, vreg_input_x_unroll_new, preg_all);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::MAX, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_cur_max, vreg_max_tmp, preg_all);
|
||||
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpMaxUb), vreg_cur_max, ureg_max, 1);
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpMaxUb), ureg_max, 0);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_in_max, inMaxUb);
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
AscendC::MicroAPI::LoadAlign(vreg_cur_max, tmpMaxUb2); // 获取新的max[s1, 1]
|
||||
AscendC::MicroAPI::Max(vreg_max_new, vreg_cur_max, vreg_in_max, preg_all); // 计算新、旧max的最大值
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)tmpMaxUb2, vreg_max_new, preg_all);
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(
|
||||
vreg_max_brc, tmpMaxUb2 + i);
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_DINTLV_B32>(
|
||||
vreg_input_x, vreg_input_x_unroll, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp_even, vreg_input_x, vreg_max_brc, preg_all);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp_odd, vreg_input_x_unroll, vreg_max_brc, preg_all);
|
||||
|
||||
// x_sum = sum(x_exp, axis=-1, keepdims=True)
|
||||
AscendC::MicroAPI::Add(vreg_exp_sum, vreg_exp_even, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::SUM, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_exp_sum, vreg_exp_sum, preg_all);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpExpSumUb), vreg_exp_sum, ureg_exp_sum, 1);
|
||||
|
||||
if constexpr (IsSameType<T2, bfloat16_t>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_even_bf16, vreg_exp_even, preg_all);
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitOne>(vreg_exp_odd_bf16, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Or((RegTensor<uint16_t>&)vreg_exp_bf16, (RegTensor<uint16_t>&)vreg_exp_even_bf16,
|
||||
(RegTensor<uint16_t>&)vreg_exp_odd_bf16, preg_all_b16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_exp_bf16, blockStride, repeatStride, preg_n_b16);
|
||||
} else if constexpr (IsSameType<T2, half>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_even_fp16, vreg_exp_even, preg_all);
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitOne>(vreg_exp_odd_fp16, vreg_exp_odd, preg_all);
|
||||
AscendC::MicroAPI::Or((RegTensor<uint16_t>&)vreg_exp_fp16, (RegTensor<uint16_t>&)vreg_exp_even_fp16,
|
||||
(RegTensor<uint16_t>&)vreg_exp_odd_fp16, preg_all_b16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_exp_fp16, blockStride, repeatStride, preg_n_b16);
|
||||
}
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpExpSumUb), ureg_exp_sum, 0);
|
||||
}
|
||||
|
||||
|
||||
// update, 64 < originN <= 128
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__aicore__ inline void ProcessVec1UpdateGeneralImpl128(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
// 写的时候固定用65或者33的stride去写,因为正向目前使能settail之后mm2的s1方向必须算满128或者64行
|
||||
// stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1)
|
||||
const uint32_t blockStride = s1BaseSize >> 1 | 0x1;
|
||||
const uint32_t repeatStride = 1;
|
||||
const uint32_t oriTailN = originN - floatRepSize;
|
||||
const uint32_t tailN = s2BaseSize - floatRepSize;
|
||||
uint32_t pltOriTailN = oriTailN;
|
||||
uint32_t pltTailN = tailN;
|
||||
uint32_t pltN = s2BaseSize;
|
||||
|
||||
__ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr();
|
||||
__ubuf__ T * inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr();
|
||||
__ubuf__ T * tmpExpSumUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr();
|
||||
__ubuf__ T * tmpMaxUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64;
|
||||
__ubuf__ T * tmpMaxUb2 = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64;
|
||||
|
||||
ProcessVec1UpdateGeneralImpl128VF<T, T2, s1BaseSize, s2BaseSize>(
|
||||
expUb, srcUb, inMaxUb, tmpExpSumUb, tmpMaxUb, tmpMaxUb2, blockStride, repeatStride,
|
||||
m, scale, minValue, pltOriTailN, pltTailN, pltN);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
#endif // VF_BASIC_BLOCK_UNALIGNED128_UPDATE_SFA_H
|
||||
@@ -0,0 +1,129 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_basic_block_unaligned64_no_update_sfa.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef VF_BASIC_BLOCK_UNALIGNED64_NO_UPDATE_SFA_H
|
||||
#define VF_BASIC_BLOCK_UNALIGNED64_NO_UPDATE_SFA_H
|
||||
|
||||
#include "vf_basic_block_utils.h"
|
||||
|
||||
using namespace regbaseutil;
|
||||
|
||||
namespace FaVectorApi {
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__simd_vf__ void ProcessVec1NoUpdateImpl64VF(
|
||||
__ubuf__ T2 * expUb, __ubuf__ T * expSumUb, __ubuf__ T * maxUb, __ubuf__ T * maxUbStart,
|
||||
__ubuf__ T * srcUb, const uint32_t blockStride, const uint32_t repeatStride,
|
||||
const uint16_t m, const T scale, const T minValue, uint32_t pltOriginalN, uint32_t pltSrcN)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_min;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_brc;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_sum;
|
||||
|
||||
// bfloat16_t
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_dst_even_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_dst_odd_bf16;
|
||||
// half
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_dst_even_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_dst_odd_fp16;
|
||||
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_max;
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum;
|
||||
|
||||
AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_all_b16 =
|
||||
AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_src_n = AscendC::MicroAPI::UpdateMask<float>(pltSrcN);
|
||||
AscendC::MicroAPI::MaskReg preg_src_n_b16 =
|
||||
AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::H>();
|
||||
AscendC::MicroAPI::MaskReg preg_ori_src_n = AscendC::MicroAPI::UpdateMask<T>(pltOriginalN);
|
||||
|
||||
// x_max = max(src, axis=-1, keepdims=True)
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_ori_src_n); // Muls(scale)
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_src_n);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::MAX, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_input_max, vreg_input_x, preg_ori_src_n);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)maxUb), vreg_input_max, ureg_max, 1);
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)maxUb), ureg_max, 0);
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(
|
||||
vreg_max_brc, maxUbStart + i);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp, vreg_input_x, vreg_max_brc, preg_ori_src_n);
|
||||
|
||||
// x_sum = sum(x_exp, axis=-1, keepdims=True)
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::SUM, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_exp_sum, vreg_exp, preg_ori_src_n);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)expSumUb), vreg_exp_sum, ureg_exp_sum, 1);
|
||||
|
||||
if constexpr (IsSameType<T2, bfloat16_t>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_bf16, vreg_exp, preg_all_b16);
|
||||
AscendC::MicroAPI::DeInterleave(vreg_dst_even_bf16, vreg_dst_odd_bf16,
|
||||
vreg_exp_bf16, vreg_exp_bf16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_dst_even_bf16, blockStride, repeatStride, preg_src_n_b16);
|
||||
} else if constexpr (IsSameType<T2, half>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_fp16, vreg_exp, preg_all_b16);
|
||||
AscendC::MicroAPI::DeInterleave(vreg_dst_even_fp16, vreg_dst_odd_fp16,
|
||||
vreg_exp_fp16, vreg_exp_fp16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_dst_even_fp16, blockStride, repeatStride, preg_src_n_b16);
|
||||
}
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)expSumUb), ureg_exp_sum, 0);
|
||||
}
|
||||
|
||||
// no update, originN <= 64
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__aicore__ inline void ProcessVec1NoUpdateImpl64(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor,
|
||||
const LocalTensor<T>& expSumTensor, const LocalTensor<T>& maxTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
__ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ T * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr();
|
||||
__ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr();
|
||||
__ubuf__ T * maxUbStart = (__ubuf__ T*)maxTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr();
|
||||
|
||||
// 写的时候固定用65或者33的stride去写,因为正向目前使能settail之后mm2的s1方向必须算满128或者64行
|
||||
// stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1)
|
||||
const uint32_t blockStride = s1BaseSize >> 1 | 0x1;
|
||||
const uint32_t repeatStride = 1;
|
||||
uint32_t pltOriginalN = originN;
|
||||
uint32_t pltSrcN = s2BaseSize;
|
||||
|
||||
ProcessVec1NoUpdateImpl64VF<T, T2, s1BaseSize, s2BaseSize>(
|
||||
expUb, expSumUb, maxUb, maxUbStart, srcUb, blockStride, repeatStride, m, scale, minValue,
|
||||
pltOriginalN, pltSrcN);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
#endif // VF_BASIC_BLOCK_UNALIGNED64_NO_UPDATE_SFA_H
|
||||
@@ -0,0 +1,141 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_basic_block_aligned64_update_sfa.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef VF_BASIC_BLOCK_ALIGNED64_UPDATE_SFA_H
|
||||
#define VF_BASIC_BLOCK_ALIGNED64_UPDATE_SFA_H
|
||||
|
||||
#include "vf_basic_block_utils.h"
|
||||
|
||||
using namespace regbaseutil;
|
||||
|
||||
namespace FaVectorApi {
|
||||
// update, originN <= 64
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 128, uint32_t s2BaseSize = 128>
|
||||
__simd_vf__ void ProcessVec1UpdateImpl64VF(
|
||||
__ubuf__ T2 * expUb, __ubuf__ T * srcUb, __ubuf__ T * inMaxUb,
|
||||
__ubuf__ T * tmpExpSumUb, __ubuf__ T * tmpMaxUb, __ubuf__ T * tmpMaxUb2, const uint32_t blockStride,
|
||||
const uint32_t repeatStride, const uint16_t m, const T scale, const T minValue, uint32_t pltOriginalN,
|
||||
uint32_t pltSrcN)
|
||||
{
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_input_x;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_tmp;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_in_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_new;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_max_brc;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_cur_max;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp;
|
||||
AscendC::MicroAPI::RegTensor<float> vreg_exp_sum;
|
||||
|
||||
// bfloat16_t
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_exp_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_dst_even_bf16;
|
||||
AscendC::MicroAPI::RegTensor<bfloat16_t> vreg_dst_odd_bf16;
|
||||
// half
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_exp_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_dst_even_fp16;
|
||||
AscendC::MicroAPI::RegTensor<half> vreg_dst_odd_fp16;
|
||||
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_max;
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg_exp_sum;
|
||||
|
||||
AscendC::MicroAPI::MaskReg preg_all = AscendC::MicroAPI::CreateMask<float, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_all_b16 =
|
||||
AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::ALL>();
|
||||
AscendC::MicroAPI::MaskReg preg_ori_src_n = AscendC::MicroAPI::UpdateMask<T>(pltOriginalN);
|
||||
AscendC::MicroAPI::MaskReg preg_src_n = AscendC::MicroAPI::UpdateMask<T>(pltSrcN);
|
||||
AscendC::MicroAPI::MaskReg preg_src_n_b16 =
|
||||
AscendC::MicroAPI::CreateMask<uint16_t, AscendC::MicroAPI::MaskPattern::H>();
|
||||
|
||||
// x_max = max(src, axis=-1, keepdims=True)
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::Muls(vreg_input_x, vreg_input_x, scale, preg_ori_src_n);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)srcUb + i * s2BaseSize, vreg_input_x, preg_src_n);
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::MAX, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_cur_max, vreg_input_x, preg_ori_src_n);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpMaxUb), vreg_cur_max, ureg_max, 1);
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpMaxUb), ureg_max, 0);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_in_max, inMaxUb);
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
AscendC::MicroAPI::LoadAlign(vreg_cur_max, tmpMaxUb2);
|
||||
AscendC::MicroAPI::Max(vreg_max_new, vreg_cur_max, vreg_in_max, preg_all); // 计算新、旧的最大值
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)tmpMaxUb2, vreg_max_new, preg_all);
|
||||
|
||||
AscendC::MicroAPI::LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
AscendC::MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(
|
||||
vreg_max_brc, tmpMaxUb2 + i);
|
||||
AscendC::MicroAPI::LoadAlign(vreg_input_x, srcUb + i * s2BaseSize);
|
||||
AscendC::MicroAPI::ExpSub(vreg_exp, vreg_input_x, vreg_max_brc, preg_ori_src_n);
|
||||
|
||||
// x_sum = sum(x_exp, axis=-1, keepdims=True)
|
||||
AscendC::MicroAPI::Reduce<MicroAPI::ReduceType::SUM, float, float, MicroAPI::MaskMergeMode::ZEROING>(
|
||||
vreg_exp_sum, vreg_exp, preg_ori_src_n);
|
||||
AscendC::MicroAPI::StoreUnAlign<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpExpSumUb), vreg_exp_sum, ureg_exp_sum, 1);
|
||||
|
||||
if constexpr (IsSameType<T2, bfloat16_t>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_bf16, vreg_exp, preg_all_b16);
|
||||
AscendC::MicroAPI::DeInterleave(vreg_dst_even_bf16, vreg_dst_odd_bf16,
|
||||
vreg_exp_bf16, vreg_exp_bf16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_dst_even_bf16, blockStride, repeatStride, preg_src_n_b16);
|
||||
} else if constexpr (IsSameType<T2, half>::value) {
|
||||
AscendC::MicroAPI::Cast<T2, T, castTraitZero>(vreg_exp_fp16, vreg_exp, preg_all_b16);
|
||||
AscendC::MicroAPI::DeInterleave(vreg_dst_even_fp16, vreg_dst_odd_fp16,
|
||||
vreg_exp_fp16, vreg_exp_fp16);
|
||||
AscendC::MicroAPI::StoreAlign<T2, MicroAPI::DataCopyMode::DATA_BLOCK_COPY,
|
||||
MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T2 *&)expUb), vreg_dst_even_fp16, blockStride, repeatStride, preg_src_n_b16);
|
||||
}
|
||||
}
|
||||
AscendC::MicroAPI::StoreUnAlignPost<float, MicroAPI::PostLiteral::POST_MODE_UPDATE>(
|
||||
((__ubuf__ T *&)tmpExpSumUb), ureg_exp_sum, 0);
|
||||
}
|
||||
|
||||
|
||||
// update, originN <= 64
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128>
|
||||
__aicore__ inline void ProcessVec1UpdateImpl64(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
// 写的时候固定用65或者33的stride去写,因为正向目前使能settail之后mm2的s1方向必须算满128或者64行
|
||||
// stride, high 16bits: blockStride (m*16*2/32), low 16bits: repeatStride (1)
|
||||
const uint32_t blockStride = s1BaseSize >> 1 | 0x1;
|
||||
const uint32_t repeatStride = 1;
|
||||
uint32_t pltOriginalN = originN;
|
||||
uint32_t pltSrcN = s2BaseSize;
|
||||
|
||||
__ubuf__ T2 * expUb = (__ubuf__ T2*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr();
|
||||
__ubuf__ T * inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr();
|
||||
__ubuf__ T * tmpExpSumUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr();
|
||||
__ubuf__ T * tmpMaxUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64;
|
||||
__ubuf__ T * tmpMaxUb2 = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64;
|
||||
|
||||
ProcessVec1UpdateImpl64VF <T, T2, s1BaseSize, s2BaseSize>(
|
||||
expUb, srcUb, inMaxUb, tmpExpSumUb, tmpMaxUb, tmpMaxUb2, blockStride, repeatStride, m, scale, minValue,
|
||||
pltOriginalN, pltSrcN);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
#endif // VF_BASIC_BLOCK_ALIGNED64_UPDATE_SFA_H
|
||||
112
csrc/attention/common/op_kernel/arch35/vf/vf_basic_block_utils.h
Normal file
112
csrc/attention/common/op_kernel/arch35/vf/vf_basic_block_utils.h
Normal file
@@ -0,0 +1,112 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_basic_block_utils.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef VF_BASIC_BLOCK_UTILS_H
|
||||
#define VF_BASIC_BLOCK_UTILS_H
|
||||
|
||||
#if ASC_DEVKIT_MAJOR >= 9
|
||||
#include "kernel_basic_intf.h"
|
||||
#else
|
||||
#include "kernel_operator.h"
|
||||
#endif
|
||||
|
||||
namespace FaVectorApi {
|
||||
constexpr uint32_t floatRepSize = 64;
|
||||
constexpr uint32_t halfRepSize = 128;
|
||||
constexpr uint32_t blockBytesU8 = 32;
|
||||
constexpr float fp8e4m3MaxValue = 448.0f;
|
||||
constexpr float int8MaxValue = 127.0f;
|
||||
constexpr float hifp8MaxValue = 32768.0f;
|
||||
constexpr float floatEps = 2.220446049250313e-16;
|
||||
/* **************************************************************************************************
|
||||
* Muls + Select(optional) + SoftmaxFlashV2 + Cast(fp32->fp16/bf16) + ND2NZ
|
||||
* ************************************************************************************************* */
|
||||
using namespace MicroAPI;
|
||||
|
||||
constexpr static AscendC::MicroAPI::CastTrait castTraitZero = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_ROUND,
|
||||
};
|
||||
|
||||
constexpr static AscendC::MicroAPI::CastTrait castTraitOne = {
|
||||
AscendC::MicroAPI::RegLayout::ONE,
|
||||
AscendC::MicroAPI::SatMode::SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_ROUND,
|
||||
};
|
||||
|
||||
constexpr static AscendC::MicroAPI::CastTrait castTraitTwo = {
|
||||
AscendC::MicroAPI::RegLayout::TWO,
|
||||
AscendC::MicroAPI::SatMode::SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_ROUND,
|
||||
};
|
||||
|
||||
constexpr static AscendC::MicroAPI::CastTrait castTraitThree = {
|
||||
AscendC::MicroAPI::RegLayout::THREE,
|
||||
AscendC::MicroAPI::SatMode::SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_ROUND,
|
||||
};
|
||||
|
||||
constexpr static AscendC::MicroAPI::CastTrait castTraitRintZero = {
|
||||
AscendC::MicroAPI::RegLayout::ZERO,
|
||||
AscendC::MicroAPI::SatMode::SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_RINT,
|
||||
};
|
||||
|
||||
constexpr static AscendC::MicroAPI::CastTrait castTraitRintOne = {
|
||||
AscendC::MicroAPI::RegLayout::ONE,
|
||||
AscendC::MicroAPI::SatMode::SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_RINT,
|
||||
};
|
||||
|
||||
constexpr static AscendC::MicroAPI::CastTrait castTraitRintTwo = {
|
||||
AscendC::MicroAPI::RegLayout::TWO,
|
||||
AscendC::MicroAPI::SatMode::SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_RINT,
|
||||
};
|
||||
|
||||
constexpr static AscendC::MicroAPI::CastTrait castTraitRintThree = {
|
||||
AscendC::MicroAPI::RegLayout::THREE,
|
||||
AscendC::MicroAPI::SatMode::SAT,
|
||||
AscendC::MicroAPI::MaskMergeMode::ZEROING,
|
||||
AscendC::RoundMode::CAST_RINT,
|
||||
};
|
||||
|
||||
#define USE_MLA_FULLQUANT_V1_P(vreg_exp, vreg_rowmax_p, MaskReg) \
|
||||
do { \
|
||||
Muls(vreg_exp, vreg_exp, fp8e4m3MaxValue, MaskReg); \
|
||||
Div(vreg_exp, vreg_exp, vreg_rowmax_p, MaskReg); \
|
||||
} while (0)
|
||||
|
||||
#define USE_MLA_FULLQUANT_V1_P_INT8(vreg_exp, vreg_rowmax_p, MaskReg) \
|
||||
do { \
|
||||
Muls(vreg_exp, vreg_exp, int8MaxValue, MaskReg); \
|
||||
Div(vreg_exp, vreg_exp, vreg_rowmax_p, MaskReg); \
|
||||
} while (0)
|
||||
|
||||
#define USE_MLA_FULLQUANT_V1_P_HIFP8(vreg_exp, vreg_rowmax_p, MaskReg) \
|
||||
do { \
|
||||
Muls(vreg_exp, vreg_exp, hifp8MaxValue, MaskReg); \
|
||||
Div(vreg_exp, vreg_exp, vreg_rowmax_p, MaskReg); \
|
||||
} while (0)
|
||||
} // namespace
|
||||
|
||||
#endif // VF_BASIC_BLOCK_UTILS_H
|
||||
727
csrc/attention/common/op_kernel/arch35/vf/vf_flashupdate_new.h
Normal file
727
csrc/attention/common/op_kernel/arch35/vf/vf_flashupdate_new.h
Normal file
@@ -0,0 +1,727 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_flashupdate_new.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef MY_FLASH_UPDATE_NEW_INTERFACE_H
|
||||
#define MY_FLASH_UPDATE_NEW_INTERFACE_H
|
||||
|
||||
#include "kernel_tensor.h"
|
||||
|
||||
namespace FaVectorApi {
|
||||
// bf16->fp32
|
||||
static constexpr MicroAPI::CastTrait castTraitFp16_32_update = {MicroAPI::RegLayout::ZERO, MicroAPI::SatMode::UNKNOWN,
|
||||
MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
constexpr uint16_t REDUCE_SIZE = 1;
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t srcD, uint16_t reduceSize, bool isUpdatePre, bool isMlaFullQuant>
|
||||
__simd_vf__ inline void FlashUpdateBasicVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, __ubuf__ float * preUb,
|
||||
__ubuf__ float * expMaxUb, __ubuf__ float * rowMaxUb, const uint16_t m, const uint16_t d,
|
||||
const float deScaleV, const float deScaleVPre)
|
||||
{
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
constexpr uint16_t dLoops = srcD / floatRepSize;
|
||||
RegTensor<float> vreg_exp_max;
|
||||
RegTensor<float> vreg_row_max;
|
||||
RegTensor<float> vreg_input_pre;
|
||||
RegTensor<float> vreg_input_cur;
|
||||
RegTensor<float> vreg_mul;
|
||||
RegTensor<float> vreg_add;
|
||||
|
||||
MaskReg preg_all = CreateMask<float, MaskPattern::ALL>();
|
||||
|
||||
// dstTensor = preTensor * expMaxTensor + curTensor
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_exp_max, expMaxUb + i * reduceSize); // [m,8]
|
||||
if constexpr (isMlaFullQuant) {
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_row_max, rowMaxUb + i * reduceSize);
|
||||
}
|
||||
|
||||
for (uint16_t j = 0; j < dLoops; ++j) {
|
||||
LoadAlign(vreg_input_pre, preUb + i * d + j * floatRepSize);
|
||||
LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize);
|
||||
if constexpr (isMlaFullQuant) {
|
||||
Mul(vreg_input_cur, vreg_input_cur, vreg_row_max, preg_all);
|
||||
}
|
||||
Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_all);
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value ||
|
||||
IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value ||
|
||||
IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all);
|
||||
if constexpr (isUpdatePre) {
|
||||
Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all);
|
||||
}
|
||||
}
|
||||
Add(vreg_add, vreg_mul, vreg_input_cur, preg_all);
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_add, preg_all);
|
||||
}
|
||||
}
|
||||
}
|
||||
/* **************************************************************************************************
|
||||
* FlashUpdate, fp32
|
||||
* ************************************************************************************************* */
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t srcD, uint16_t reduceSize, bool isUpdatePre, bool isMlaFullQuant>
|
||||
__aicore__ inline void FlashUpdateBasic(const LocalTensor<T>& dstTensor, const LocalTensor<T>& curTensor,
|
||||
const LocalTensor<T>& preTensor, const LocalTensor<T>& expMaxTensor, const LocalTensor<T>& rowMaxTensor,
|
||||
const uint16_t m, const uint16_t d, const float deScaleV, const float deScaleVPre)
|
||||
{
|
||||
__ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr();
|
||||
__ubuf__ float * preUb = (__ubuf__ T*)preTensor.GetPhyAddr();
|
||||
__ubuf__ float * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr();
|
||||
__ubuf__ float * rowMaxUb = (__ubuf__ T*)rowMaxTensor.GetPhyAddr();
|
||||
|
||||
FlashUpdateBasicVF<T, INPUT_T, OUTPUT_T, srcD, reduceSize, isUpdatePre, isMlaFullQuant>(
|
||||
dstUb, curUb, preUb, expMaxUb, rowMaxUb, m, d, deScaleV, deScaleVPre);
|
||||
}
|
||||
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t reduceSize, bool isUpdatePre>
|
||||
__simd_vf__ inline void FlashUpdateGeneralVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, __ubuf__ float * preUb,
|
||||
__ubuf__ float * expMaxUb, const uint16_t m, const uint16_t d,
|
||||
const float deScaleV, const float deScaleVPre, const uint32_t pltTailD, const uint16_t hasTail)
|
||||
{
|
||||
RegTensor<float> vreg_exp_max;
|
||||
RegTensor<float> vreg_input_pre;
|
||||
RegTensor<float> vreg_input_cur;
|
||||
RegTensor<float> vreg_mul;
|
||||
RegTensor<float> vreg_add;
|
||||
|
||||
MaskReg preg_all = CreateMask<float, MaskPattern::ALL>();
|
||||
uint32_t tmpTailD = pltTailD;
|
||||
MaskReg preg_tail_d = UpdateMask<float>(tmpTailD);
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
const uint16_t dLoops = d / floatRepSize;
|
||||
|
||||
// dstTensor = preTensor * expMaxTensor + curTensor
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_exp_max, expMaxUb + i * reduceSize); // [m,8]
|
||||
|
||||
for (uint16_t j = 0; j < dLoops; ++j) {
|
||||
LoadAlign(vreg_input_pre, preUb + i * d + j * floatRepSize);
|
||||
LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize);
|
||||
|
||||
Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_all);
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value ||
|
||||
IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value ||
|
||||
IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all);
|
||||
if constexpr (isUpdatePre) {
|
||||
Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all);
|
||||
}
|
||||
}
|
||||
Add(vreg_add, vreg_mul, vreg_input_cur, preg_all);
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_add, preg_all);
|
||||
}
|
||||
for (uint16_t t = 0; t < hasTail; ++t) {
|
||||
LoadAlign(vreg_input_pre, preUb + i * d + dLoops * floatRepSize);
|
||||
LoadAlign(vreg_input_cur, curUb + i * d + dLoops * floatRepSize);
|
||||
|
||||
Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_tail_d);
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value ||
|
||||
IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value ||
|
||||
IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all);
|
||||
if constexpr (isUpdatePre) {
|
||||
Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all);
|
||||
}
|
||||
}
|
||||
Add(vreg_add, vreg_mul, vreg_input_cur, preg_tail_d);
|
||||
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)dstUb + i * d + dLoops * floatRepSize, vreg_add, preg_tail_d);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t reduceSize, bool isUpdatePre>
|
||||
__aicore__ inline void FlashUpdateGeneral(const LocalTensor<T>& dstTensor, const LocalTensor<T>& curTensor,
|
||||
const LocalTensor<T>& preTensor, const LocalTensor<T>& expMaxTensor, const uint16_t m, const uint16_t d,
|
||||
const float deScaleV, const float deScaleVPre)
|
||||
{
|
||||
__ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr();
|
||||
__ubuf__ float * preUb = (__ubuf__ T*)preTensor.GetPhyAddr();
|
||||
__ubuf__ float * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr();
|
||||
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
const uint16_t tailD = d % floatRepSize;
|
||||
uint32_t pltTailD = static_cast<uint32_t>(tailD);
|
||||
|
||||
uint16_t hasTail = 0;
|
||||
if (tailD > 0) {
|
||||
hasTail = 1;
|
||||
}
|
||||
|
||||
FlashUpdateGeneralVF<T, INPUT_T, OUTPUT_T, reduceSize, isUpdatePre>(
|
||||
dstUb, curUb, preUb, expMaxUb, m, d, deScaleV, deScaleVPre, pltTailD, hasTail);
|
||||
}
|
||||
|
||||
/*
|
||||
* @ingroup FlashUpdate
|
||||
* @brief compute, dstTensor = preTensor * expMaxTensor + curTensor
|
||||
* @param [out] dstTensor, output LocalTensor
|
||||
* @param [in] curTensor, input LocalTensor
|
||||
* @param [in] preTensor, input LocalTensor
|
||||
* @param [in] expMaxTensor, input LocalTensor
|
||||
* @param [in] m, input rows
|
||||
* @param [in] d, input columns, should be 32 bytes aligned
|
||||
*/
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t srcD, bool isUpdatePre, bool isMlaFullQuant>
|
||||
__aicore__ inline void FlashUpdateNew(const LocalTensor<T>& dstTensor, const LocalTensor<T>& curTensor,
|
||||
const LocalTensor<T>& preTensor, const LocalTensor<T>& expMaxTensor, const LocalTensor<T>& rowMaxTensor, const uint16_t m, const uint16_t d,
|
||||
const float deScaleV, const float deScaleVPre)
|
||||
{
|
||||
static_assert(IsSameType<T, float>::value, "VF FlashUpdate, T must be float");
|
||||
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
if constexpr(srcD % floatRepSize == 0) {
|
||||
FlashUpdateBasic<T, INPUT_T, OUTPUT_T, srcD, REDUCE_SIZE, isUpdatePre, isMlaFullQuant>(dstTensor, curTensor, preTensor, expMaxTensor, rowMaxTensor,
|
||||
m, d, deScaleV, deScaleVPre);
|
||||
} else {
|
||||
|
||||
FlashUpdateGeneral<T, INPUT_T, OUTPUT_T, REDUCE_SIZE, isUpdatePre>(dstTensor, curTensor, preTensor, expMaxTensor, m, d,
|
||||
deScaleV, deScaleVPre);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t srcD, uint16_t reduceSize, bool isUpdatePre, bool isMlaFullQuant>
|
||||
__simd_vf__ inline void FlashUpdateLastBasicVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, __ubuf__ float * preUb,
|
||||
__ubuf__ float * expMaxUb, __ubuf__ float * expSumUb, __ubuf__ float * rowMaxUb, const uint16_t m, const uint16_t d,
|
||||
const float deScaleV, const float deScaleVPre)
|
||||
{
|
||||
RegTensor<float> vreg_exp_max;
|
||||
RegTensor<float> vreg_row_max;
|
||||
RegTensor<float> vreg_input_pre;
|
||||
RegTensor<float> vreg_input_cur;
|
||||
RegTensor<float> vreg_mul;
|
||||
RegTensor<float> vreg_add;
|
||||
RegTensor<float> vreg_div;
|
||||
RegTensor<half> vreg_cast;
|
||||
RegTensor<float> vreg_exp_sum;
|
||||
|
||||
MaskReg preg_all = CreateMask<float, MaskPattern::ALL>();
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
constexpr uint16_t dLoops = srcD / floatRepSize;
|
||||
constexpr float fp8e4m3MaxValueRec = 1 / 448.0f;
|
||||
constexpr float int8MaxValueRec = 1 / 127.0f;
|
||||
constexpr float hifp8MaxValueRec = 1 / 32768.0f;
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_exp_max, expMaxUb + i * reduceSize);
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_exp_sum, expSumUb + i * reduceSize);
|
||||
if constexpr (isMlaFullQuant) {
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_row_max, rowMaxUb + i * reduceSize);
|
||||
}
|
||||
for (uint16_t j = 0; j < dLoops; ++j) {
|
||||
LoadAlign(vreg_input_pre, preUb + i * d + j * floatRepSize);
|
||||
LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize);
|
||||
if constexpr (isMlaFullQuant) {
|
||||
Mul(vreg_input_cur, vreg_input_cur, vreg_row_max, preg_all);
|
||||
}
|
||||
Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_all);
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value ||
|
||||
IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value ||
|
||||
IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all);
|
||||
if constexpr (isUpdatePre) {
|
||||
Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all);
|
||||
}
|
||||
}
|
||||
Add(vreg_add, vreg_mul, vreg_input_cur, preg_all);
|
||||
Div(vreg_div, vreg_add, vreg_exp_sum, preg_all);
|
||||
if constexpr (isMlaFullQuant) {
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e4m3fn_t>::value) {
|
||||
Muls(vreg_div, vreg_div, fp8e4m3MaxValueRec, preg_all);
|
||||
} else if constexpr (IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_div, vreg_div, int8MaxValueRec, preg_all);
|
||||
} else {
|
||||
Muls(vreg_div, vreg_div, hifp8MaxValueRec, preg_all);
|
||||
}
|
||||
}
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_div, preg_all);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t srcD, uint16_t reduceSize, bool isUpdatePre, bool isMlaFullQuant>
|
||||
__aicore__ inline void FlashUpdateLastBasic(const LocalTensor<T>& dstTensor,
|
||||
const LocalTensor<T>& curTensor, const LocalTensor<T>& preTensor,
|
||||
const LocalTensor<T>& expMaxTensor, const LocalTensor<T>& rowMaxTensor, const LocalTensor<T>& expSumTensor,
|
||||
const uint16_t m, const uint16_t d, const float deScaleV, const float deScaleVPre)
|
||||
{
|
||||
__ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr();
|
||||
__ubuf__ float * preUb = (__ubuf__ T*)preTensor.GetPhyAddr();
|
||||
__ubuf__ float * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr();
|
||||
__ubuf__ float * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr();
|
||||
__ubuf__ float * rowMaxUb = (__ubuf__ T*)rowMaxTensor.GetPhyAddr();
|
||||
|
||||
FlashUpdateLastBasicVF<T, INPUT_T, OUTPUT_T, srcD, reduceSize, isUpdatePre, isMlaFullQuant>(
|
||||
dstUb, curUb, preUb, expMaxUb, expSumUb, rowMaxUb, m, d, deScaleV, deScaleVPre);
|
||||
}
|
||||
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t reduceSize, bool isUpdatePre>
|
||||
__simd_vf__ inline void FlashUpdateLastGeneralVF(__ubuf__ float * dstUb, __ubuf__ float * curUb,
|
||||
__ubuf__ float * preUb, __ubuf__ float * expMaxUb, __ubuf__ float * expSumUb, const uint16_t m, const uint16_t d,
|
||||
const float deScaleV, const float deScaleVPre, const uint32_t pltTailD, const uint16_t hasTail)
|
||||
{
|
||||
RegTensor<float> vreg_exp_max;
|
||||
RegTensor<float> vreg_input_pre;
|
||||
RegTensor<float> vreg_input_cur;
|
||||
RegTensor<float> vreg_mul;
|
||||
RegTensor<float> vreg_add;
|
||||
RegTensor<float> vreg_div;
|
||||
RegTensor<half> vreg_cast;
|
||||
RegTensor<float> vreg_exp_sum;
|
||||
|
||||
MaskReg preg_all = CreateMask<float, MaskPattern::ALL>();
|
||||
uint32_t tmpTailD = pltTailD;
|
||||
MaskReg preg_tail_d = UpdateMask<float>(tmpTailD);
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
uint16_t dLoops = d / floatRepSize;
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_exp_max, expMaxUb + i * reduceSize);
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_exp_sum, expSumUb + i * reduceSize);
|
||||
for (uint16_t j = 0; j < dLoops; ++j) {
|
||||
LoadAlign(vreg_input_pre, preUb + i * d + j * floatRepSize);
|
||||
LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize);
|
||||
|
||||
Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_all);
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value ||
|
||||
IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value ||
|
||||
IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all);
|
||||
if constexpr (isUpdatePre) {
|
||||
Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all);
|
||||
}
|
||||
}
|
||||
Add(vreg_add, vreg_mul, vreg_input_cur, preg_all);
|
||||
Div(vreg_div, vreg_add, vreg_exp_sum, preg_all);
|
||||
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_div, preg_all);
|
||||
}
|
||||
|
||||
for (uint16_t t = 0; t < hasTail; ++t) {
|
||||
LoadAlign(vreg_input_pre, preUb + i * d + dLoops * floatRepSize);
|
||||
LoadAlign(vreg_input_cur, curUb + i * d + dLoops * floatRepSize);
|
||||
Mul(vreg_mul, vreg_exp_max, vreg_input_pre, preg_tail_d);
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value ||
|
||||
IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value ||
|
||||
IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all);
|
||||
if constexpr (isUpdatePre) {
|
||||
Muls(vreg_mul, vreg_mul, deScaleVPre, preg_all);
|
||||
}
|
||||
}
|
||||
Add(vreg_add, vreg_mul, vreg_input_cur, preg_tail_d);
|
||||
Div(vreg_div, vreg_add, vreg_exp_sum, preg_tail_d);
|
||||
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)dstUb + i * d + dLoops * floatRepSize, vreg_div, preg_tail_d);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t reduceSize, bool isUpdatePre>
|
||||
__aicore__ inline void FlashUpdateLastGeneral(const LocalTensor<T>& dstTensor,
|
||||
const LocalTensor<T>& curTensor, const LocalTensor<T>& preTensor,
|
||||
const LocalTensor<T>& expMaxTensor, const LocalTensor<T>& expSumTensor,
|
||||
const uint16_t m, const uint16_t d, const float deScaleV, const float deScaleVPre)
|
||||
{
|
||||
__ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr();
|
||||
__ubuf__ float * preUb = (__ubuf__ T*)preTensor.GetPhyAddr();
|
||||
__ubuf__ float * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr();
|
||||
__ubuf__ float * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr();
|
||||
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
uint16_t tailD = d % floatRepSize;
|
||||
uint32_t pltTailD = tailD;
|
||||
|
||||
uint16_t hasTail = 0;
|
||||
if (tailD > 0) {
|
||||
hasTail = 1;
|
||||
}
|
||||
|
||||
FlashUpdateLastGeneralVF<T, INPUT_T, OUTPUT_T, reduceSize, isUpdatePre>(
|
||||
dstUb, curUb, preUb, expMaxUb, expSumUb, m, d, deScaleV, deScaleVPre, pltTailD, hasTail);
|
||||
}
|
||||
|
||||
/*
|
||||
* @ingroup FlashUpdateLast
|
||||
* @brief compute, dstTensor = (preTensor * expMaxTensor + curTensor) / expSumTensor
|
||||
* @param [out] dstTensor, output LocalTensor
|
||||
* @param [in] curTensor, input LocalTensor
|
||||
* @param [in] preTensor, input LocalTensor
|
||||
* @param [in] expMaxTensor, input LocalTensor
|
||||
* @param [in] expSumTensor, input LocalTensor
|
||||
* @param [in] m, input rows
|
||||
* @param [in] d, input columns, 32 bytes align
|
||||
*/
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint16_t srcD, bool isUpdatePre, bool isMlaFullQuant>
|
||||
__aicore__ inline void FlashUpdateLastNew(const LocalTensor<T>& dstTensor,
|
||||
const LocalTensor<T>& curTensor, const LocalTensor<T>& preTensor,
|
||||
const LocalTensor<T>& expMaxTensor, const LocalTensor<T>& rowMaxTensor, const LocalTensor<T>& expSumTensor,
|
||||
uint16_t m, uint16_t d, const float deScaleV, const float deScaleVPre)
|
||||
{
|
||||
static_assert(IsSameType<T, float>::value, "VF FlashUpdateLast, T must be float");
|
||||
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
if constexpr(srcD % floatRepSize == 0) {
|
||||
FlashUpdateLastBasic<T, INPUT_T, OUTPUT_T, srcD, REDUCE_SIZE, isUpdatePre, isMlaFullQuant>(
|
||||
dstTensor, curTensor, preTensor, expMaxTensor, rowMaxTensor, expSumTensor, m, d, deScaleV, deScaleVPre);
|
||||
} else {
|
||||
FlashUpdateLastGeneral<T, INPUT_T, OUTPUT_T, REDUCE_SIZE, isUpdatePre>(
|
||||
dstTensor, curTensor, preTensor, expMaxTensor, expSumTensor, m, d, deScaleV, deScaleVPre);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint32_t srcD, bool isMlaFullQuant>
|
||||
__simd_vf__ inline void LastDivNewVF(__ubuf__ float * dstUb, __ubuf__ float * curUb, __ubuf__ float * expSumUb,
|
||||
const uint16_t m, const uint16_t d, const float deScaleV)
|
||||
{
|
||||
RegTensor<float> vreg_input_cur;
|
||||
RegTensor<float> vreg_div;
|
||||
RegTensor<float> vreg_exp_sum;
|
||||
MaskReg preg_all = CreateMask<float, MaskPattern::ALL>();
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
const uint16_t dLoops = d >> 6;
|
||||
constexpr float fp8e4m3MaxValueRec = 1 / 448.0f;
|
||||
constexpr float int8MaxValueRec = 1 / 127.0f;
|
||||
constexpr float hifp8MaxValueRec = 1 / 32768.0f;
|
||||
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
uint32_t sreg_init = d;
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_exp_sum, expSumUb + i * REDUCE_SIZE);
|
||||
for (uint16_t j = 0; j < dLoops; ++j) {
|
||||
MaskReg preg_update = UpdateMask<float>(sreg_init);
|
||||
|
||||
LoadAlign(vreg_input_cur, curUb + i * d + j * floatRepSize);
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e5m2_t>::value ||
|
||||
IsSameType<INPUT_T, fp8_e4m3fn_t>::value ||
|
||||
IsSameType<INPUT_T, hifloat8_t>::value ||
|
||||
IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_input_cur, vreg_input_cur, deScaleV, preg_all);
|
||||
}
|
||||
Div(vreg_div, vreg_input_cur, vreg_exp_sum, preg_update);
|
||||
if constexpr (isMlaFullQuant) {
|
||||
if constexpr (IsSameType<INPUT_T, fp8_e4m3fn_t>::value) {
|
||||
Muls(vreg_div, vreg_div, fp8e4m3MaxValueRec, preg_all);
|
||||
} else if constexpr (IsSameType<INPUT_T, int8_t>::value) {
|
||||
Muls(vreg_div, vreg_div, int8MaxValueRec, preg_all);
|
||||
} else {
|
||||
Muls(vreg_div, vreg_div, hifp8MaxValueRec, preg_all);
|
||||
}
|
||||
}
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_div, preg_update);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// dstTensor = curTensor / expSumTensor, curTensor: [64,128], expSumTensor: [64,8]
|
||||
template <typename T, typename INPUT_T, typename OUTPUT_T, uint32_t srcD, bool isMlaFullQuant>
|
||||
__aicore__ inline void LastDivNew(const LocalTensor<T>& dstTensor, const LocalTensor<T>& curTensor,
|
||||
const LocalTensor<T>& expSumTensor, const uint16_t m, const uint16_t d, const float deScaleV)
|
||||
{
|
||||
__ubuf__ float * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ float * curUb = (__ubuf__ T*)curTensor.GetPhyAddr();
|
||||
__ubuf__ float * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr();
|
||||
|
||||
LastDivNewVF<T, INPUT_T, OUTPUT_T, srcD, isMlaFullQuant>(dstUb, curUb, expSumUb, m, d, deScaleV);
|
||||
}
|
||||
|
||||
template <typename T, uint32_t srcD>
|
||||
__simd_vf__ inline void InvalidLineUpdateVF(__ubuf__ T * dstUb, __ubuf__ T * srcUb, __ubuf__ T * maxUb,
|
||||
const uint16_t m, const uint16_t d, const T minValue, const T invalidValue)
|
||||
{
|
||||
RegTensor<float> vreg_invalid_value;
|
||||
RegTensor<float> vreg_max;
|
||||
RegTensor<float> vreg_input;
|
||||
RegTensor<float> vreg_input_brc;
|
||||
|
||||
MaskReg preg_all = CreateMask<float, MaskPattern::ALL>();
|
||||
MaskReg preg_compare;
|
||||
const uint16_t dLoops = d >> 6;
|
||||
|
||||
Duplicate(vreg_invalid_value, invalidValue);
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
LoadAlign<T, MicroAPI::LoadDist::DIST_BRC_B32>(vreg_max, maxUb + i);
|
||||
Compares<T, CMPMODE::EQ>(preg_compare, vreg_max, minValue, preg_all);
|
||||
for (uint16_t j = 0; j < dLoops; ++j) {
|
||||
LoadAlign(vreg_input, srcUb + i * d + j * floatRepSize);
|
||||
Select(vreg_input_brc, vreg_invalid_value, vreg_input, preg_compare);
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)dstUb + i * d + j * floatRepSize, vreg_input_brc, preg_all);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, uint32_t srcD>
|
||||
__aicore__ inline void InvalidLineUpdate(const LocalTensor<T>& dstTensor, const LocalTensor<T>& srcTensor,
|
||||
const LocalTensor<T>& maxTensor, const uint16_t m, const uint16_t d, const T minValue, const T invalidValue)
|
||||
{
|
||||
__ubuf__ T * dstUb = (__ubuf__ T*)dstTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcUb = (__ubuf__ T*)srcTensor.GetPhyAddr();
|
||||
__ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr();
|
||||
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
uint16_t dLoops = d >> 6;
|
||||
|
||||
InvalidLineUpdateVF<T, srcD>(dstUb, srcUb, maxUb, m, d, minValue, invalidValue);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ inline void ComputeLseOutputVF(__ubuf__ T *srcSumUb, __ubuf__ T *srcMaxUb, __ubuf__ T *dstUb, const uint32_t dealCount)
|
||||
{
|
||||
MicroAPI::RegTensor<T> vregSum;
|
||||
MicroAPI::RegTensor<T> vregMax;
|
||||
MicroAPI::RegTensor<T> vregRes;
|
||||
MicroAPI::RegTensor<T> vregResFinal;
|
||||
MicroAPI::RegTensor<float> vregMinValue;
|
||||
MicroAPI::RegTensor<float> vregInfValue;
|
||||
MicroAPI::MaskReg pregCompare;
|
||||
constexpr uint32_t dealRows = 8;
|
||||
constexpr uint32_t floatRepSize = 64; // 64: 一个寄存器存64个float
|
||||
constexpr float infValue = 3e+99; // 3e+99 for float inf
|
||||
constexpr uint32_t tmpMin = 0xFF7FFFFF;
|
||||
float minValue = *((float*)&tmpMin);
|
||||
uint16_t updateLoops = dealCount / dealRows;
|
||||
uint16_t tailLSize = dealCount % dealRows * 8;
|
||||
uint32_t pltTail = static_cast<uint32_t>(tailLSize);
|
||||
|
||||
MicroAPI::MaskReg pregAll = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregTail = MicroAPI::UpdateMask<T>(pltTail);
|
||||
MicroAPI::Duplicate<float, float>(vregMinValue, minValue);
|
||||
MicroAPI::Duplicate<float, float>(vregInfValue, infValue);
|
||||
|
||||
for (uint16_t i = 0; i < updateLoops; ++i) {
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_E2B_B32>(vregSum, srcSumUb + (i * dealRows));
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_E2B_B32>(vregMax, srcMaxUb + (i * dealRows));
|
||||
|
||||
MicroAPI::Log<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregSum, pregAll);
|
||||
MicroAPI::Add<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregRes, vregMax, pregAll);
|
||||
|
||||
MicroAPI::Compare<float, CMPMODE::EQ>(pregCompare, vregMax, vregMinValue, pregAll);
|
||||
MicroAPI::Select<T>(vregResFinal, vregInfValue, vregRes, pregCompare);
|
||||
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(dstUb + (i * floatRepSize), vregResFinal, pregAll);
|
||||
}
|
||||
|
||||
if (tailLSize != 0) {
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_E2B_B32>(vregSum, srcSumUb + dealRows * updateLoops);
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_E2B_B32>(vregMax, srcMaxUb + dealRows * updateLoops);
|
||||
|
||||
MicroAPI::Log<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregSum, pregTail);
|
||||
MicroAPI::Add<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregRes, vregMax, pregTail);
|
||||
|
||||
MicroAPI::Compare<float, CMPMODE::EQ>(pregCompare, vregMax, vregMinValue, pregTail);
|
||||
MicroAPI::Select<T>(vregResFinal, vregInfValue, vregRes, pregCompare);
|
||||
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(dstUb + floatRepSize * updateLoops, vregResFinal, pregTail);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void ComputeLseOutputVF(const LocalTensor<T>& dstTensor, const LocalTensor<T>& softmaxSumTensor,
|
||||
const LocalTensor<T>& softmaxMaxTensor, uint32_t dealCount)
|
||||
{
|
||||
__ubuf__ T * srcSumUb = (__ubuf__ T *)softmaxSumTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcMaxUb = (__ubuf__ T *)softmaxMaxTensor.GetPhyAddr();
|
||||
__ubuf__ T * dstUb = (__ubuf__ T *)dstTensor.GetPhyAddr();
|
||||
|
||||
ComputeLseOutputVF<T>(srcSumUb, srcMaxUb, dstUb, dealCount);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ inline void SinkSubExpAddVF(__ubuf__ T *srcSumUb, __ubuf__ T *srcMaxUb, const T sinkValue, const uint32_t dealCount)
|
||||
{
|
||||
MicroAPI::RegTensor<T> vregSum;
|
||||
MicroAPI::RegTensor<T> vregMax;
|
||||
MicroAPI::RegTensor<T> vregRes;
|
||||
MicroAPI::RegTensor<T> vregSink;
|
||||
|
||||
constexpr uint32_t floatRepSize = 64;
|
||||
|
||||
uint16_t updateLoops = dealCount / floatRepSize;
|
||||
uint16_t tailSize = dealCount % floatRepSize;
|
||||
uint32_t pltTail = static_cast<uint32_t>(tailSize);
|
||||
|
||||
//mask
|
||||
MicroAPI::MaskReg pregAll = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregTail = MicroAPI::UpdateMask<T>(pltTail);
|
||||
|
||||
Duplicate(vregSink, sinkValue);
|
||||
|
||||
for (uint16_t i = 0; i < updateLoops; ++i) {
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregSum, srcSumUb + (i * floatRepSize));
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregMax, srcMaxUb + (i * floatRepSize));
|
||||
|
||||
MicroAPI::Sub<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregSink, vregMax, pregAll);
|
||||
MicroAPI::Exp<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregRes, pregAll);
|
||||
MicroAPI::Add<T, MicroAPI::MaskMergeMode::ZEROING>(vregSum, vregSum, vregRes, pregAll);
|
||||
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(srcSumUb + (i * floatRepSize), vregSum, pregAll);
|
||||
}
|
||||
|
||||
for (uint16_t i = 0; i < tailSize; i = i + tailSize) {
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregSum, srcSumUb + (updateLoops * floatRepSize));
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregMax, srcMaxUb + (updateLoops * floatRepSize));
|
||||
|
||||
MicroAPI::Sub<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregSink, vregMax, pregTail);
|
||||
MicroAPI::Exp<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregRes, pregTail);
|
||||
MicroAPI::Add<T, MicroAPI::MaskMergeMode::ZEROING>(vregSum, vregSum, vregRes, pregTail);
|
||||
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(srcSumUb + (updateLoops * floatRepSize), vregSum, pregTail);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void SinkSubExpAddVF(const LocalTensor<T>& softmaxSumTensor, const LocalTensor<T>& softmaxMaxTensor,
|
||||
const T sinkValue, uint32_t dealCount)
|
||||
{
|
||||
__ubuf__ T * srcSumUb = (__ubuf__ T *)softmaxSumTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcMaxUb = (__ubuf__ T *)softmaxMaxTensor.GetPhyAddr();
|
||||
|
||||
SinkSubExpAddVF<T>(srcSumUb, srcMaxUb, sinkValue, dealCount);
|
||||
}
|
||||
|
||||
template <typename T, typename SINK_T>
|
||||
__simd_vf__ inline void SinkSubExpAddGSFusedVF(__ubuf__ T *srcSumUb, __ubuf__ T *srcMaxUb, __ubuf__ uint16_t *sinkUb, const uint32_t dealCount)
|
||||
{
|
||||
MicroAPI::RegTensor<T> vregSum;
|
||||
MicroAPI::RegTensor<T> vregMax;
|
||||
MicroAPI::RegTensor<T> vregRes;
|
||||
MicroAPI::RegTensor<SINK_T> vregSink;
|
||||
MicroAPI::RegTensor<T> vregSinkCast;
|
||||
|
||||
constexpr uint32_t floatRepSize = 64;
|
||||
|
||||
uint16_t updateLoops = dealCount / floatRepSize;
|
||||
uint16_t tailSize = dealCount % floatRepSize;
|
||||
uint32_t pltTail = static_cast<uint32_t>(tailSize);
|
||||
|
||||
//mask
|
||||
MicroAPI::MaskReg pregAll = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg pregTail = MicroAPI::UpdateMask<T>(pltTail);
|
||||
MicroAPI::MaskReg pregSinkAll = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
|
||||
MicroAPI::LoadAlign<uint16_t, MicroAPI::LoadDist::DIST_UNPACK_B16>((MicroAPI::RegTensor<uint16_t>&)vregSink, sinkUb);
|
||||
MicroAPI::Cast<T, SINK_T, castTraitFp16_32_update>(vregSinkCast, vregSink, pregSinkAll);
|
||||
|
||||
for (uint16_t i = 0; i < updateLoops; ++i) {
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregSum, srcSumUb + (i * floatRepSize));
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregMax, srcMaxUb + (i * floatRepSize));
|
||||
|
||||
MicroAPI::Sub<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregSinkCast, vregMax, pregAll);
|
||||
MicroAPI::Exp<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregRes, pregAll);
|
||||
MicroAPI::Add<T, MicroAPI::MaskMergeMode::ZEROING>(vregSum, vregSum, vregRes, pregAll);
|
||||
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(srcSumUb + (i * floatRepSize), vregSum, pregAll);
|
||||
}
|
||||
|
||||
if (tailSize != 0) {
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregSum, srcSumUb + (updateLoops * floatRepSize));
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregMax, srcMaxUb + (updateLoops * floatRepSize));
|
||||
|
||||
MicroAPI::Sub<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregSinkCast, vregMax, pregTail);
|
||||
MicroAPI::Exp<T, MicroAPI::MaskMergeMode::ZEROING>(vregRes, vregRes, pregTail);
|
||||
MicroAPI::Add<T, MicroAPI::MaskMergeMode::ZEROING>(vregSum, vregSum, vregRes, pregTail);
|
||||
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(srcSumUb + (updateLoops * floatRepSize), vregSum, pregTail);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename SINK_T>
|
||||
__aicore__ inline void SinkSubExpAddGSFusedVF(const LocalTensor<SINK_T>& dstTensor, const LocalTensor<T>& softmaxSumTensor,
|
||||
const LocalTensor<T>& softmaxMaxTensor, uint32_t dealCount)
|
||||
{
|
||||
__ubuf__ T * srcSumUb = (__ubuf__ T *)softmaxSumTensor.GetPhyAddr();
|
||||
__ubuf__ T * srcMaxUb = (__ubuf__ T *)softmaxMaxTensor.GetPhyAddr();
|
||||
__ubuf__ uint16_t * dstUb = (__ubuf__ uint16_t *)dstTensor.GetPhyAddr();
|
||||
|
||||
SinkSubExpAddGSFusedVF<T, SINK_T>(srcSumUb, srcMaxUb, dstUb, dealCount);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ inline void RowInvalidUpdateVF(__ubuf__ T *finalUb, __ubuf__ float *maxUb, const uint16_t m,
|
||||
const uint16_t d, int64_t dSize, const uint32_t pltTailD, const uint16_t hasTail)
|
||||
{
|
||||
constexpr uint16_t floatRepSize = 64; // 64: 一个寄存器可以存储64个float类型数据
|
||||
const uint16_t dLoops = d / floatRepSize;
|
||||
|
||||
|
||||
constexpr uint32_t tmpZero = 0x00000000; // zero value of fp16 and fp32
|
||||
const T zeroValue = *((T*)&tmpZero);
|
||||
constexpr uint32_t tmpMin = 0xFF7FFFFF; // min value of float
|
||||
const float minValue = *((float*)&tmpMin);
|
||||
MicroAPI::RegTensor<float> vregMinValue;
|
||||
MicroAPI::RegTensor<T> vregZeroValue;
|
||||
MicroAPI::RegTensor<float> vregMax;
|
||||
MicroAPI::RegTensor<T> vregFinal;
|
||||
MicroAPI::RegTensor<T> vregFinalNew;
|
||||
|
||||
MicroAPI::MaskReg pregAll = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
uint32_t tmpTailD = pltTailD;
|
||||
MicroAPI::MaskReg pregTailD = MicroAPI::UpdateMask<T>(tmpTailD);
|
||||
MicroAPI::MaskReg pregCompare;
|
||||
|
||||
MicroAPI::Duplicate<float, float>(vregMinValue, minValue);
|
||||
MicroAPI::Duplicate<T, T>(vregZeroValue, zeroValue);
|
||||
for (uint16_t i = 0; i < m; ++i) {
|
||||
MicroAPI::LoadAlign<float, MicroAPI::LoadDist::DIST_BRC_B32>(vregMax, maxUb + i);
|
||||
MicroAPI::Compare<float, CMPMODE::EQ>(pregCompare, vregMax, vregMinValue, pregAll);
|
||||
for (uint16_t j = 0; j < dLoops; ++j) {
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregFinal, finalUb + i * dSize + j * floatRepSize);
|
||||
MicroAPI::Select<T>(vregFinalNew, vregZeroValue, vregFinal, pregCompare);
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(finalUb + i * dSize + j * floatRepSize,
|
||||
vregFinalNew, pregAll);
|
||||
}
|
||||
for (uint16_t t = 0; t < hasTail; ++t) {
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregFinal, finalUb + i * dSize + dLoops * floatRepSize);
|
||||
MicroAPI::Select<T>(vregFinalNew, vregZeroValue, vregFinal, pregCompare);
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(finalUb + i * dSize + dLoops * floatRepSize,
|
||||
vregFinalNew, pregTailD);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void RowInvalidUpdateVF(const LocalTensor<T>& finalTensor, const LocalTensor<float>& maxTensor,
|
||||
const uint16_t m, const uint16_t d, int64_t dSize)
|
||||
{
|
||||
__ubuf__ T * finalUb = (__ubuf__ T*)finalTensor.GetPhyAddr();
|
||||
__ubuf__ float * maxUb = (__ubuf__ float*)maxTensor.GetPhyAddr();
|
||||
|
||||
constexpr uint16_t floatRepSize = 64;
|
||||
const uint16_t tailD = d % floatRepSize;
|
||||
uint32_t pltTailD = static_cast<uint32_t>(tailD);
|
||||
uint16_t hasTail = 0;
|
||||
if (tailD > 0) {
|
||||
hasTail = 1;
|
||||
}
|
||||
|
||||
RowInvalidUpdateVF<T>(finalUb, maxUb, m, d, dSize, pltTailD, hasTail);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
#endif // MY_FLASH_UPDATE_INTERFACE_H
|
||||
@@ -0,0 +1,164 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_mul_sel_softmaxflashv2_cast_nz_sfa.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef MUL_SEL_SOFTMAX_FLASH_V2_CAST_NZ_SFA_INTERFACE_H
|
||||
#define MUL_SEL_SOFTMAX_FLASH_V2_CAST_NZ_SFA_INTERFACE_H
|
||||
|
||||
#include "vf_basic_block_aligned128_no_update_sfa.h"
|
||||
#include "vf_basic_block_aligned128_update_sfa.h"
|
||||
#include "vf_basic_block_unaligned64_update_sfa.h"
|
||||
#include "vf_basic_block_unaligned64_no_update_sfa.h"
|
||||
#include "vf_basic_block_unaligned128_no_update_sfa.h"
|
||||
#include "vf_basic_block_unaligned128_update_sfa.h"
|
||||
|
||||
using namespace regbaseutil;
|
||||
|
||||
namespace FaVectorApi {
|
||||
/* **************************************************************************************************
|
||||
* Muls + Select(optional) + SoftmaxFlashV2 + Cast(fp32->fp16/bf16) + ND2NZ
|
||||
* ************************************************************************************************* */
|
||||
using AscendC::LocalTensor;
|
||||
|
||||
enum class OriginNRange {
|
||||
EQ_128_SFA = 0, // originN == 128, better performance than GT_64_AND_LTE_128 (s2BaseSize=128)
|
||||
GT_0_AND_LTE_64_SFA, // 0 < originN <= 64 (s2BaseSize <= 64 or tail s2)
|
||||
GT_64_AND_LTE_128_SFA, // 64 < originN <= 128, support for non-alignment (s2BaseSize=128)
|
||||
N_INVALID_SFA
|
||||
};
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128,
|
||||
OriginNRange oriNRange = OriginNRange::EQ_128_SFA>
|
||||
__aicore__ inline void ProcessVec1NoUpdate(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor,
|
||||
const LocalTensor<T>& expSumTensor, const LocalTensor<T>& maxTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
if constexpr (oriNRange == OriginNRange::EQ_128_SFA) {
|
||||
ProcessVec1NoUpdateImpl128<T, T2, s1BaseSize, s2BaseSize>(
|
||||
dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue);
|
||||
} else if constexpr (oriNRange == OriginNRange::GT_0_AND_LTE_64_SFA) {
|
||||
ProcessVec1NoUpdateImpl64<T, T2, s1BaseSize, s2BaseSize>(
|
||||
dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue);
|
||||
} else if constexpr (oriNRange == OriginNRange::GT_64_AND_LTE_128_SFA) {
|
||||
ProcessVec1NoUpdateGeneralImpl128<T, T2, s1BaseSize, s2BaseSize>(
|
||||
dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename T2, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128,
|
||||
OriginNRange oriNRange = OriginNRange::EQ_128_SFA>
|
||||
__aicore__ inline void ProcessVec1Update(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor,
|
||||
const LocalTensor<T>& expSumTensor, const LocalTensor<T>& maxTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
if constexpr (oriNRange == OriginNRange::EQ_128_SFA) {
|
||||
ProcessVec1UpdateImpl128<T, T2, s1BaseSize, s2BaseSize>(
|
||||
dstTensor, srcTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue);
|
||||
} else if constexpr (oriNRange == OriginNRange::GT_0_AND_LTE_64_SFA) {
|
||||
ProcessVec1UpdateImpl64<T, T2, s1BaseSize, s2BaseSize>(
|
||||
dstTensor, srcTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue);
|
||||
} else if constexpr (oriNRange == OriginNRange::GT_64_AND_LTE_128_SFA) {
|
||||
ProcessVec1UpdateGeneralImpl128<T, T2, s1BaseSize, s2BaseSize>(
|
||||
dstTensor, srcTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T, typename T2, bool isUpdate = false, uint32_t s1BaseSize = 64, uint32_t s2BaseSize = 128,
|
||||
OriginNRange oriNRange = OriginNRange::EQ_128_SFA>
|
||||
__aicore__ inline void ProcessVec1Vf(
|
||||
const LocalTensor<T2>& dstTensor, const LocalTensor<T>& srcTensor,
|
||||
const LocalTensor<T>& expSumTensor, const LocalTensor<T>& maxTensor, const LocalTensor<T>& inMaxTensor,
|
||||
const LocalTensor<T>& sharedTmpBuffer, const uint16_t m, const uint32_t originN, const T scale, const T minValue)
|
||||
{
|
||||
static_assert(IsSameType<T, float>::value, "VF mul_sel_softmaxflashv2_cast_nz, T must be float");
|
||||
static_assert((IsSameType<T2, half>::value || IsSameType<T2, bfloat16_t>::value),
|
||||
"VF mul_sel_softmaxflashv2_cast_nz, T2 must be half or bfloat16");
|
||||
|
||||
if constexpr (!isUpdate) {
|
||||
ProcessVec1NoUpdate<T, T2, s1BaseSize, s2BaseSize, oriNRange>(
|
||||
dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue);
|
||||
} else {
|
||||
ProcessVec1Update<T, T2, s1BaseSize, s2BaseSize, oriNRange>(
|
||||
dstTensor, srcTensor, expSumTensor, maxTensor, inMaxTensor, sharedTmpBuffer, m, originN, scale, minValue);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ inline void UpdateExpSumAndExpMaxVF(__ubuf__ T * maxUb, __ubuf__ T * inMaxUb, __ubuf__ T * expMaxUb,
|
||||
__ubuf__ T * expSumUb, __ubuf__ T * inExpSumUb, __ubuf__ T * tmpExpSumUb, __ubuf__ T * tmpMaxUb, const uint32_t m)
|
||||
{
|
||||
RegTensor<float> vreg_input_x;
|
||||
RegTensor<float> vreg_input_x_unroll;
|
||||
RegTensor<float> vreg_max;
|
||||
RegTensor<float> vreg_in_max;
|
||||
RegTensor<float> vreg_exp_sum;
|
||||
RegTensor<float> vreg_in_exp_sum;
|
||||
RegTensor<float> vreg_exp_max;
|
||||
RegTensor<float> vreg_exp_sum_brc;
|
||||
RegTensor<float> vreg_exp_sum_update;
|
||||
MaskReg preg_all = CreateMask<float, MaskPattern::ALL>();
|
||||
// 注意:当m大于64的时候需要开启循环
|
||||
LoadAlign(vreg_max, tmpMaxUb);
|
||||
LoadAlign(vreg_in_max, inMaxUb);
|
||||
FusedExpSub(vreg_exp_max, vreg_in_max, vreg_max, preg_all);
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)expMaxUb, vreg_exp_max, preg_all);
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)maxUb, vreg_max, preg_all);
|
||||
LoadAlign(vreg_in_exp_sum, inExpSumUb);
|
||||
|
||||
// x_sum = exp_max * insum + x_sum
|
||||
LoadAlign(vreg_exp_sum_brc, tmpExpSumUb);
|
||||
Mul(vreg_exp_sum_update, vreg_exp_max, vreg_in_exp_sum, preg_all);
|
||||
Add(vreg_exp_sum_update, vreg_exp_sum_update, vreg_exp_sum_brc, preg_all);
|
||||
StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(
|
||||
(__ubuf__ T *&)expSumUb, vreg_exp_sum_update, preg_all);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void SFAUpdateExpSumAndExpMax(
|
||||
const LocalTensor<T>& expSumTensor, const LocalTensor<T>& maxTensor,
|
||||
const LocalTensor<T>& expMaxTensor, const LocalTensor<T>& inExpSumTensor,
|
||||
const LocalTensor<T>& inMaxTensor, const LocalTensor<T>& sharedTmpBuffer, const uint32_t m)
|
||||
{
|
||||
__ubuf__ T * maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr();
|
||||
__ubuf__ T * inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr();
|
||||
|
||||
__ubuf__ T * expMaxUb = (__ubuf__ T*)expMaxTensor.GetPhyAddr();
|
||||
__ubuf__ T * expSumUb = (__ubuf__ T*)expSumTensor.GetPhyAddr();
|
||||
__ubuf__ T * inExpSumUb = (__ubuf__ T*)inExpSumTensor.GetPhyAddr();
|
||||
|
||||
__ubuf__ T * tmpExpSumUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr();
|
||||
__ubuf__ T * tmpMaxUb = (__ubuf__ T*)sharedTmpBuffer.GetPhyAddr() + 64;
|
||||
|
||||
UpdateExpSumAndExpMaxVF<T>(maxUb, inMaxUb, expMaxUb, expSumUb, inExpSumUb, tmpExpSumUb, tmpMaxUb, m);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ inline void DuplicateSumWithR0VF(__ubuf__ T * sumUb, const T R0, uint32_t m) {
|
||||
AscendC::MicroAPI::RegTensor<T> vreg_sum;
|
||||
AscendC::MicroAPI::MaskReg preg_m = AscendC::MicroAPI::UpdateMask<T>(m);
|
||||
AscendC::MicroAPI::UnalignRegForStore ureg;
|
||||
AscendC::MicroAPI::Duplicate<T, MicroAPI::MaskMergeMode::ZEROING, T>(vreg_sum, R0, preg_m);
|
||||
AscendC::MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM_B32>(sumUb, vreg_sum, preg_m);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DuplicateSumWithR0(const LocalTensor<T>& sumTensor, const T R0, uint32_t m)
|
||||
{
|
||||
__ubuf__ T * sumUb = (__ubuf__ T*)sumTensor.GetPhyAddr();
|
||||
DuplicateSumWithR0VF<T>(sumUb, R0, m);
|
||||
}
|
||||
} // namespace
|
||||
#endif // MUL_SEL_SOFTMAX_FLASH_V2_CAST_NZ_SFA_INTERFACE_H
|
||||
292
csrc/attention/common/op_kernel/buffer.h
Normal file
292
csrc/attention/common/op_kernel/buffer.h
Normal file
@@ -0,0 +1,292 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file buffer.h
|
||||
* \brief同步管理
|
||||
*/
|
||||
#ifndef BUFFER_H
|
||||
#define BUFFER_H
|
||||
#include<type_traits>
|
||||
#include"lib/matmul_intf.h"
|
||||
#if ASC_DEVKIT_MAJOR >= 9
|
||||
#include "kernel_basic_intf.h"
|
||||
#else
|
||||
#include "kernel_operator.h"
|
||||
#endif
|
||||
using namespace AscendC;
|
||||
namespace fa_base_matmul {
|
||||
__BLOCK_LOCAL__ __inline__ uint32_t idCounterNum;
|
||||
#define MAKE_ID ((++idCounterNum) % 11)
|
||||
|
||||
// 核间同步中,AIC(flagId 0-10)对应AIV0(flagId 0-10),对应AIV1(flagId 16-26)
|
||||
#define AIV0_AIV1_OFFSET 16
|
||||
|
||||
enum class BufferType {
|
||||
L1 = 0,
|
||||
L0A = 1,
|
||||
L0B = 2,
|
||||
L0C = 3,
|
||||
UB = 4,
|
||||
GM = 5,
|
||||
C2 = 6,
|
||||
};
|
||||
|
||||
enum class SyncType {
|
||||
NO_SYNC,
|
||||
INNER_CORE_SYNC,
|
||||
CROSS_CORE_SYNC_FORWARD,
|
||||
CROSS_CORE_SYNC_BOTH,
|
||||
CROSS_CORE_SYNC_BACKWARD,
|
||||
};
|
||||
|
||||
constexpr uint32_t INVALID_CROSS_CORE_EVENT_ID = 16;
|
||||
static constexpr uint64_t CROSS_CORE_SYNC_MODE = 4;
|
||||
|
||||
template<BufferType Type>
|
||||
struct BufferInfo{
|
||||
// Cons 消费者,Prod 生产者
|
||||
__aicore__ const static constexpr HardEvent ConsWaitProdStatus() {
|
||||
if constexpr (Type == BufferType::L1) {
|
||||
return HardEvent::MTE2_MTE1;
|
||||
} else if constexpr (Type == BufferType::L0A) {
|
||||
return HardEvent::MTE1_M;
|
||||
} else if constexpr (Type == BufferType::L0B) {
|
||||
return HardEvent::MTE1_M;
|
||||
} else if constexpr (Type == BufferType::L0C) {
|
||||
return HardEvent::M_FIX;
|
||||
} else if constexpr (Type == BufferType::C2) {
|
||||
return HardEvent::MTE1_M;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ const static constexpr HardEvent ProdWaitConsStatus() {
|
||||
if constexpr (Type == BufferType::L1) {
|
||||
return HardEvent::MTE1_MTE2;
|
||||
} else if constexpr (Type == BufferType::L0A) {
|
||||
return HardEvent::M_MTE1;
|
||||
} else if constexpr (Type == BufferType::L0B) {
|
||||
return HardEvent::M_MTE1;
|
||||
} else if constexpr (Type == BufferType::L0C) {
|
||||
return HardEvent::FIX_M;
|
||||
} else if constexpr (Type == BufferType::C2) {
|
||||
return HardEvent::M_MTE1;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ const static constexpr TPosition GetTPosition() {
|
||||
if constexpr (Type == BufferType::L1) {
|
||||
return TPosition::A1;
|
||||
} else if constexpr (Type == BufferType::L0A) {
|
||||
return TPosition::A2;
|
||||
} else if constexpr (Type == BufferType::L0B) {
|
||||
return TPosition::B2;
|
||||
} else if constexpr (Type == BufferType::L0C) {
|
||||
return TPosition::CO1;
|
||||
} else if constexpr (Type == BufferType::UB) {
|
||||
return TPosition::VECIN;
|
||||
} else if constexpr (Type == BufferType::GM) {
|
||||
return TPosition::GM;
|
||||
} else if constexpr (Type == BufferType::C2) {
|
||||
return TPosition::C2;
|
||||
}
|
||||
}
|
||||
|
||||
static constexpr HardEvent EventP2C = ConsWaitProdStatus(); // 生产者到消费者方向的HardEvent:消费者等生产者提供/生产者通知消费者已生成
|
||||
static constexpr HardEvent EventC2P = ProdWaitConsStatus(); // 消费者到生产者方向的HardEvent:生产者等消费者消耗/消费者通知生产者已消耗’
|
||||
static constexpr TPosition Position = GetTPosition();
|
||||
};
|
||||
|
||||
// buffer绑定生产者、消费者关系
|
||||
// L1 buffer的生产者为MTE2或者MTE3,消费者为MTE1
|
||||
// L0A buffer的生产者为MTE1,消费者为M
|
||||
// L0B buffer的生产者为MTE1,消费者为M
|
||||
// L0C buffer的生产者为M,消费者为FIX
|
||||
template<BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
class Buffer {
|
||||
using TensorType = std::conditional_t<bufferType == BufferType::GM, GlobalTensor<uint8_t>, LocalTensor<uint8_t>>;
|
||||
|
||||
template <typename T>
|
||||
using TargetTensorType = std::conditional_t<bufferType == BufferType::GM, GlobalTensor<T>, LocalTensor<T>>;
|
||||
public:
|
||||
__aicore__ inline Buffer() {}
|
||||
__aicore__ inline Buffer(TensorType tensor, uint32_t size) {
|
||||
tensor_ = tensor;
|
||||
size_ = size;
|
||||
if constexpr (syncType == SyncType::CROSS_CORE_SYNC_FORWARD) {
|
||||
id0_ = MAKE_ID;
|
||||
id1_ = INVALID_CROSS_CORE_EVENT_ID;
|
||||
} else if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BACKWARD) {
|
||||
id0_ = INVALID_CROSS_CORE_EVENT_ID;
|
||||
id1_ = MAKE_ID;
|
||||
} else if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BOTH) {
|
||||
id0_ = MAKE_ID;
|
||||
id1_ = MAKE_ID;
|
||||
} else {
|
||||
id0_ = INVALID_CROSS_CORE_EVENT_ID;
|
||||
id1_ = INVALID_CROSS_CORE_EVENT_ID;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void Init() {
|
||||
if ASCEND_IS_AIC {
|
||||
if constexpr (syncType == SyncType::INNER_CORE_SYNC) {
|
||||
p2cEventId_ = GetTPipePtr()->AllocEventID<BufferInfo<bufferType>::EventP2C>(); // 确保只能被调用一次
|
||||
c2pEventId_ = GetTPipePtr()->AllocEventID<BufferInfo<bufferType>::EventC2P>();
|
||||
SetFlag<BufferInfo<bufferType>::EventC2P>(c2pEventId_);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void UnInit() {
|
||||
if ASCEND_IS_AIC {
|
||||
if constexpr (syncType == SyncType::INNER_CORE_SYNC) {
|
||||
WaitFlag<BufferInfo<bufferType>::EventC2P>(c2pEventId_);
|
||||
GetTPipePtr()->ReleaseEventID<BufferInfo<bufferType>::EventP2C>(p2cEventId_); // 确保只能被调用一次
|
||||
GetTPipePtr()->ReleaseEventID<BufferInfo<bufferType>::EventC2P>(c2pEventId_);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<HardEvent EventType>
|
||||
__aicore__ inline void Wait() {
|
||||
if ASCEND_IS_AIC {
|
||||
if constexpr (syncType == SyncType::INNER_CORE_SYNC) {
|
||||
if constexpr (EventType == BufferInfo<bufferType>::EventP2C) {
|
||||
WaitFlag<BufferInfo<bufferType>::EventP2C>(p2cEventId_); // 消费者等待生产者完成生产
|
||||
} else {
|
||||
WaitFlag<BufferInfo<bufferType>::EventC2P>(c2pEventId_); // 生产者等待消费者完成消费
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<HardEvent EventType>
|
||||
__aicore__ inline void Set() {
|
||||
if ASCEND_IS_AIC {
|
||||
if constexpr (syncType == SyncType::INNER_CORE_SYNC) {
|
||||
if constexpr (EventType == BufferInfo<bufferType>::EventP2C) {
|
||||
SetFlag<BufferInfo<bufferType>::EventP2C>(p2cEventId_); // 生产者通知消费者已完成生产
|
||||
} else {
|
||||
SetFlag<BufferInfo<bufferType>::EventC2P>(c2pEventId_); // 消费者通知生产者已完成消费
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void SetEventID() {
|
||||
if ASCEND_IS_AIC {
|
||||
p2cEventId_ = GetTPipePtr()->AllocEventID<BufferInfo<bufferType>::EventP2C>(); // 确保只能被调用一次
|
||||
c2pEventId_ = GetTPipePtr()->AllocEventID<BufferInfo<bufferType>::EventC2P>();
|
||||
}
|
||||
}
|
||||
|
||||
template<HardEvent EventType>
|
||||
__aicore__ inline TEventID GetEventID() {
|
||||
if ASCEND_IS_AIC {
|
||||
if constexpr (EventType == BufferInfo<bufferType>::EventP2C) {
|
||||
return p2cEventId_; // 生产者通知消费者已完成生产
|
||||
} else {
|
||||
return c2pEventId_; // 消费者通知生产者已完成消费
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<bool isReuse = false>
|
||||
__aicore__ inline void WaitCrossCore() {
|
||||
if constexpr (bufferType == BufferType::GM && syncType == SyncType::CROSS_CORE_SYNC_BACKWARD) {
|
||||
// AIC属于消费者,AIV属于生产者,且一个AIC对应两个AIV
|
||||
if ASCEND_IS_AIC {
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE2>(id1_);
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE2>(id1_ + AIV0_AIV1_OFFSET);
|
||||
} else {
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE2>(id0_);
|
||||
}
|
||||
} else if constexpr (bufferType == BufferType::UB || bufferType == BufferType::GM) {
|
||||
// AIC属于生产者,AIV属于消费者,且一个AIC对应两个AIV
|
||||
if ASCEND_IS_AIC {
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_FIX>(id1_);
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_FIX>(id1_ + AIV0_AIV1_OFFSET);
|
||||
} else {
|
||||
if constexpr (isReuse) {
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE3>(id0_);
|
||||
} else {
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_V>(id0_);
|
||||
}
|
||||
}
|
||||
} else if constexpr (bufferType == BufferType::L1) {
|
||||
// AIC属于消费者,AIV属于生产者,且一个AIC对应两个AIV
|
||||
if ASCEND_IS_AIC {
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE1>(id0_);
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE1>(id0_ + AIV0_AIV1_OFFSET);
|
||||
} else {
|
||||
if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BOTH) {
|
||||
CrossCoreWaitFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE3>(id1_);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<bool isReuse = false>
|
||||
__aicore__ inline void SetCrossCore() {
|
||||
if constexpr (bufferType == BufferType::GM && syncType == SyncType::CROSS_CORE_SYNC_BACKWARD) {
|
||||
// AIC属于消费者,AIV属于生产者,且一个AIC对应两个AIV
|
||||
if ASCEND_IS_AIC {
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_FIX>(id0_);
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_FIX>(id0_ + AIV0_AIV1_OFFSET);
|
||||
} else {
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE3>(id1_);
|
||||
}
|
||||
} else if constexpr (bufferType == BufferType::UB || bufferType == BufferType::GM) {
|
||||
// AIC属于生产者,AIV属于消费者,且一个AIC对应两个AIV
|
||||
if ASCEND_IS_AIC {
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_FIX>(id0_);
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_FIX>(id0_ + AIV0_AIV1_OFFSET);
|
||||
} else {
|
||||
if constexpr (isReuse) {
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE3>(id1_);
|
||||
} else {
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_V>(id1_);
|
||||
}
|
||||
}
|
||||
} else if constexpr (bufferType == BufferType::L1) {
|
||||
// AIC属于消费者,AIV属于生产者,且一个AIC对应两个AIV
|
||||
if ASCEND_IS_AIC {
|
||||
if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BOTH) {
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE1>(id1_);
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE1>(id1_ + AIV0_AIV1_OFFSET);
|
||||
}
|
||||
} else {
|
||||
CrossCoreSetFlag<CROSS_CORE_SYNC_MODE, PIPE_MTE3>(id0_);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
__aicore__ inline TargetTensorType<T> GetTensor() {
|
||||
return tensor_.template ReinterpretCast<T>();
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
__aicore__ inline TargetTensorType<T> GetTensor(uint64_t startindex) {
|
||||
TargetTensorType<T> tmpTensor = tensor_.template ReinterpretCast<T>();
|
||||
return tmpTensor[startindex];
|
||||
}
|
||||
|
||||
private:
|
||||
TensorType tensor_;
|
||||
uint32_t size_;
|
||||
TEventID p2cEventId_;
|
||||
TEventID c2pEventId_;
|
||||
uint32_t id0_; // 用作正向同步:生产者通知消费者,或者消费者等待生产者;
|
||||
uint32_t id1_; // 用作反向同步:消费者通知生产者,或者生产者等待消费者;
|
||||
};
|
||||
}
|
||||
#endif
|
||||
61
csrc/attention/common/op_kernel/buffer_manager.h
Normal file
61
csrc/attention/common/op_kernel/buffer_manager.h
Normal file
@@ -0,0 +1,61 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file buffer_manager.h
|
||||
* \brief buffer内存管理
|
||||
*/
|
||||
#ifndef BUFFER_MANAGER_H
|
||||
#define BUFFER_MANAGER_H
|
||||
|
||||
#if (__NPU_ARCH__ == 5102)
|
||||
#include "buffer_mix_core.h"
|
||||
#else
|
||||
#include "buffer.h"
|
||||
#endif
|
||||
|
||||
// L1 TPosition::A1
|
||||
// L0A TPosition::A2
|
||||
// L0B TPosition::B2
|
||||
// L0C TPosition::CO1
|
||||
// UB TPosition::VECIN
|
||||
namespace fa_base_matmul {
|
||||
template<BufferType bufferType>
|
||||
class BufferManager {
|
||||
using TensorType = std::conditional_t<bufferType == BufferType::GM, GlobalTensor<uint8_t>, LocalTensor<uint8_t>>;
|
||||
public:
|
||||
__aicore__ inline void Init(TPipe *pipe, uint32_t size) {
|
||||
static_assert(bufferType != BufferType::GM, "GM should use workspace.");
|
||||
TBuf<BufferInfo<bufferType>::Position> tbuf;
|
||||
pipe->InitBuffer(tbuf, size);
|
||||
mem_ = tbuf.template Get<uint8_t>();
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(__gm__ uint8_t* workspace) {
|
||||
static_assert(bufferType == BufferType::GM, "BufferType should be GM.");
|
||||
mem_.SetGlobalBuffer((__gm__ uint8_t*)workspace);
|
||||
}
|
||||
|
||||
template<SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
__aicore__ inline Buffer<bufferType, syncType> AllocBuffer(uint32_t size) {
|
||||
TensorType temp = mem_[offset_];
|
||||
offset_ += size;
|
||||
return Buffer<bufferType, syncType>(temp, size);
|
||||
}
|
||||
|
||||
template<SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
__aicore__ inline void FreeBuffer(Buffer<bufferType, syncType> &buffer){
|
||||
}
|
||||
private:
|
||||
uint32_t offset_ = 0;
|
||||
TensorType mem_;
|
||||
};
|
||||
}
|
||||
#endif
|
||||
216
csrc/attention/common/op_kernel/buffer_mix_core.h
Normal file
216
csrc/attention/common/op_kernel/buffer_mix_core.h
Normal file
@@ -0,0 +1,216 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file buffer_mix_core.h
|
||||
* \brief同步管理
|
||||
*/
|
||||
#ifndef BUFFER_MIX_CORE_H
|
||||
#define BUFFER_MIX_CORE_H
|
||||
#include <type_traits>
|
||||
#include "lib/matmul_intf.h"
|
||||
#if ASC_DEVKIT_MAJOR >= 9
|
||||
#include "kernel_basic_intf.h"
|
||||
#else
|
||||
#include "kernel_operator.h"
|
||||
#endif
|
||||
using namespace AscendC;
|
||||
namespace fa_base_matmul {
|
||||
__BLOCK_LOCAL__ __inline__ uint32_t idCounterNum;
|
||||
#define MAKE_ID ((++idCounterNum) % 11)
|
||||
|
||||
// 核间同步中,AIC(flagId 0-10)对应AIV0(flagId 0-10),对应AIV1(flagId 16-26)
|
||||
#define AIV0_AIV1_OFFSET 16
|
||||
|
||||
enum class BufferType {
|
||||
L1 = 0,
|
||||
L0A = 1,
|
||||
L0B = 2,
|
||||
L0C = 3,
|
||||
UB = 4,
|
||||
GM = 5,
|
||||
};
|
||||
|
||||
enum class SyncType {
|
||||
NO_SYNC,
|
||||
INNER_CORE_SYNC,
|
||||
CROSS_CORE_SYNC_FORWARD,
|
||||
CROSS_CORE_SYNC_BOTH,
|
||||
};
|
||||
|
||||
constexpr uint32_t INVALID_CROSS_CORE_EVENT_ID = 16;
|
||||
static constexpr uint64_t CROSS_CORE_SYNC_MODE = 4;
|
||||
|
||||
template <BufferType Type>
|
||||
struct BufferInfo {
|
||||
// Cons 消费者,Prod 生产者
|
||||
__aicore__ const static constexpr HardEvent ConsWaitProdStatus()
|
||||
{
|
||||
if constexpr (Type == BufferType::L1) {
|
||||
return HardEvent::MTE2_MTE1;
|
||||
} else if constexpr (Type == BufferType::L0A) {
|
||||
return HardEvent::MTE1_M;
|
||||
} else if constexpr (Type == BufferType::L0B) {
|
||||
return HardEvent::MTE1_M;
|
||||
} else if constexpr (Type == BufferType::L0C) {
|
||||
return HardEvent::M_FIX;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ const static constexpr HardEvent ProdWaitConsStatus()
|
||||
{
|
||||
if constexpr (Type == BufferType::L1) {
|
||||
return HardEvent::MTE1_MTE2;
|
||||
} else if constexpr (Type == BufferType::L0A) {
|
||||
return HardEvent::M_MTE1;
|
||||
} else if constexpr (Type == BufferType::L0B) {
|
||||
return HardEvent::M_MTE1;
|
||||
} else if constexpr (Type == BufferType::L0C) {
|
||||
return HardEvent::FIX_M;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ const static constexpr TPosition GetTPosition()
|
||||
{
|
||||
if constexpr (Type == BufferType::L1) {
|
||||
return TPosition::A1;
|
||||
} else if constexpr (Type == BufferType::L0A) {
|
||||
return TPosition::A2;
|
||||
} else if constexpr (Type == BufferType::L0B) {
|
||||
return TPosition::B2;
|
||||
} else if constexpr (Type == BufferType::L0C) {
|
||||
return TPosition::CO1;
|
||||
} else if constexpr (Type == BufferType::UB) {
|
||||
return TPosition::VECIN;
|
||||
} else if constexpr (Type == BufferType::GM) {
|
||||
return TPosition::GM;
|
||||
}
|
||||
}
|
||||
|
||||
static constexpr HardEvent EventP2C =
|
||||
ConsWaitProdStatus(); // 生产者到消费者方向的HardEvent:消费者等生产者提供/生产者通知消费者已生成
|
||||
static constexpr HardEvent EventC2P =
|
||||
ProdWaitConsStatus(); // 消费者到生产者方向的HardEvent:生产者等消费者消耗/消费者通知生产者已消耗’
|
||||
static constexpr TPosition Position = GetTPosition();
|
||||
};
|
||||
|
||||
// buffer绑定生产者、消费者关系
|
||||
// L1 buffer的生产者为MTE2或者MTE3,消费者为MTE1
|
||||
// L0A buffer的生产者为MTE1,消费者为M
|
||||
// L0B buffer的生产者为MTE1,消费者为M
|
||||
// L0C buffer的生产者为M,消费者为FIX
|
||||
template <BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
class Buffer {
|
||||
using TensorType = std::conditional_t<bufferType == BufferType::GM, GlobalTensor<uint8_t>, LocalTensor<uint8_t>>;
|
||||
|
||||
template <typename T>
|
||||
using TargetTensorType = std::conditional_t<bufferType == BufferType::GM, GlobalTensor<T>, LocalTensor<T>>;
|
||||
|
||||
public:
|
||||
__aicore__ inline Buffer()
|
||||
{
|
||||
}
|
||||
__aicore__ inline Buffer(TensorType tensor, uint32_t size)
|
||||
{
|
||||
tensor_ = tensor;
|
||||
size_ = size;
|
||||
if constexpr (syncType == SyncType::CROSS_CORE_SYNC_FORWARD) {
|
||||
id0_ = MAKE_ID;
|
||||
id1_ = INVALID_CROSS_CORE_EVENT_ID;
|
||||
} else if constexpr (syncType == SyncType::CROSS_CORE_SYNC_BOTH) {
|
||||
id0_ = MAKE_ID;
|
||||
id1_ = MAKE_ID;
|
||||
} else {
|
||||
id0_ = INVALID_CROSS_CORE_EVENT_ID;
|
||||
id1_ = INVALID_CROSS_CORE_EVENT_ID;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void Init()
|
||||
{
|
||||
if constexpr (syncType == SyncType::INNER_CORE_SYNC) {
|
||||
p2cEventId_ = GetTPipePtr()->AllocEventID<BufferInfo<bufferType>::EventP2C>(); // 确保只能被调用一次
|
||||
c2pEventId_ = GetTPipePtr()->AllocEventID<BufferInfo<bufferType>::EventC2P>();
|
||||
SetFlag<BufferInfo<bufferType>::EventC2P>(c2pEventId_);
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void UnInit()
|
||||
{
|
||||
if constexpr (syncType == SyncType::INNER_CORE_SYNC) {
|
||||
WaitFlag<BufferInfo<bufferType>::EventC2P>(c2pEventId_);
|
||||
GetTPipePtr()->ReleaseEventID<BufferInfo<bufferType>::EventP2C>(p2cEventId_); // 确保只能被调用一次
|
||||
GetTPipePtr()->ReleaseEventID<BufferInfo<bufferType>::EventC2P>(c2pEventId_);
|
||||
}
|
||||
}
|
||||
|
||||
template <HardEvent EventType>
|
||||
__aicore__ inline void Wait()
|
||||
{
|
||||
if constexpr (syncType == SyncType::INNER_CORE_SYNC) {
|
||||
if constexpr (EventType == BufferInfo<bufferType>::EventP2C) {
|
||||
WaitFlag<BufferInfo<bufferType>::EventP2C>(p2cEventId_); // 消费者等待生产者完成生产
|
||||
} else {
|
||||
WaitFlag<BufferInfo<bufferType>::EventC2P>(c2pEventId_); // 生产者等待消费者完成消费
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <HardEvent EventType>
|
||||
__aicore__ inline void Set()
|
||||
{
|
||||
if constexpr (syncType == SyncType::INNER_CORE_SYNC) {
|
||||
if constexpr (EventType == BufferInfo<bufferType>::EventP2C) {
|
||||
SetFlag<BufferInfo<bufferType>::EventP2C>(p2cEventId_); // 生产者通知消费者已完成生产
|
||||
} else {
|
||||
SetFlag<BufferInfo<bufferType>::EventC2P>(c2pEventId_); // 消费者通知生产者已完成消费
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void SetEventID()
|
||||
{
|
||||
p2cEventId_ = GetTPipePtr()->AllocEventID<BufferInfo<bufferType>::EventP2C>(); // 确保只能被调用一次
|
||||
c2pEventId_ = GetTPipePtr()->AllocEventID<BufferInfo<bufferType>::EventC2P>();
|
||||
}
|
||||
|
||||
template <HardEvent EventType>
|
||||
__aicore__ inline TEventID GetEventID()
|
||||
{
|
||||
if constexpr (EventType == BufferInfo<bufferType>::EventP2C) {
|
||||
return p2cEventId_; // 生产者通知消费者已完成生产
|
||||
} else {
|
||||
return c2pEventId_; // 消费者通知生产者已完成消费
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline TargetTensorType<T> GetTensor()
|
||||
{
|
||||
return tensor_.template ReinterpretCast<T>();
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline TargetTensorType<T> GetTensor(uint64_t startindex)
|
||||
{
|
||||
TargetTensorType<T> tmpTensor = tensor_.template ReinterpretCast<T>();
|
||||
return tmpTensor[startindex];
|
||||
}
|
||||
|
||||
private:
|
||||
TensorType tensor_;
|
||||
uint32_t size_;
|
||||
TEventID p2cEventId_;
|
||||
TEventID c2pEventId_;
|
||||
uint32_t id0_; // 用作正向同步:生产者通知消费者,或者消费者等待生产者;
|
||||
uint32_t id1_; // 用作反向同步:消费者通知生产者,或者生产者等待消费者;
|
||||
};
|
||||
} // namespace fa_base_matmul
|
||||
#endif
|
||||
407
csrc/attention/common/op_kernel/buffers_policy.h
Normal file
407
csrc/attention/common/op_kernel/buffers_policy.h
Normal file
@@ -0,0 +1,407 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file buffers_policy.h
|
||||
* \brief 综合管理buffer的内存和同步
|
||||
*/
|
||||
#ifndef BUFFERS_POLICY_H
|
||||
#define BUFFERS_POLICY_H
|
||||
|
||||
#include "buffer_manager.h"
|
||||
#define NUM_2 2
|
||||
#define NUM_3 3
|
||||
#define NUM_4 4
|
||||
// Q复用 KV复用
|
||||
// 申请单块buffer
|
||||
namespace fa_base_matmul {
|
||||
template<BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
class BuffersPolicySingleBuffer {
|
||||
public:
|
||||
__aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size){
|
||||
buffer_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
buffer_.Init();
|
||||
}
|
||||
|
||||
__aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager){
|
||||
buffer_.UnInit();
|
||||
bufferManager.FreeBuffer(buffer_);
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &Get(){
|
||||
return buffer_;
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetPre(){
|
||||
return Get();
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetReused(){
|
||||
return Get();
|
||||
}
|
||||
private:
|
||||
Buffer<bufferType, syncType> buffer_;
|
||||
};
|
||||
|
||||
// 申请2个buffer,乒乓轮转
|
||||
template<BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
class BuffersPolicyDB {
|
||||
public:
|
||||
__aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size){
|
||||
ping_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
pong_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
|
||||
ping_.Init();
|
||||
pong_.Init();
|
||||
}
|
||||
|
||||
__aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager){
|
||||
ping_.UnInit();
|
||||
pong_.UnInit();
|
||||
|
||||
bufferManager.FreeBuffer(ping_);
|
||||
bufferManager.FreeBuffer(pong_);
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &Get() {
|
||||
if (flag1_) { // 1
|
||||
flag1_ = 0;
|
||||
return ping_;
|
||||
} else { // 0
|
||||
flag1_ = 1;
|
||||
return pong_;
|
||||
}
|
||||
}
|
||||
|
||||
// 需要与Get联用, 首次调用Get,第二次调用GetPre(Q复用)
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetPre() {
|
||||
if (flag1_) { // 0->1
|
||||
return pong_;
|
||||
} else { // 1->0
|
||||
return ping_;
|
||||
}
|
||||
}
|
||||
|
||||
// 需要与Get,GetPre联用, 首次调用Get,第二次调用GetPre,第三次复用时GetReused(KV复用)
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetReused() {
|
||||
if (flag2_ == 0) {
|
||||
flag2_ = 1;
|
||||
return pong_;
|
||||
} else {
|
||||
flag2_ = 0;
|
||||
return ping_;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetReused(bool isNextS2IdxNoChange) {
|
||||
if (isNextS2IdxNoChange) {
|
||||
if (flag2_ == 0) {
|
||||
return pong_;
|
||||
} else {
|
||||
return ping_;
|
||||
}
|
||||
} else {
|
||||
return GetReused();
|
||||
}
|
||||
}
|
||||
|
||||
private:
|
||||
Buffer<bufferType, syncType> ping_;
|
||||
Buffer<bufferType, syncType> pong_;
|
||||
uint32_t flag1_ = 0;
|
||||
uint32_t flag2_ = 0;
|
||||
};
|
||||
|
||||
// 申请3个buffer, 轮转
|
||||
template<BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
class BuffersPolicy3buff {
|
||||
public:
|
||||
__aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size) {
|
||||
a_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
b_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
c_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
|
||||
a_.Init();
|
||||
b_.Init();
|
||||
c_.Init();
|
||||
}
|
||||
|
||||
__aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager) {
|
||||
a_.UnInit();
|
||||
b_.UnInit();
|
||||
c_.UnInit();
|
||||
|
||||
bufferManager.FreeBuffer(a_);
|
||||
bufferManager.FreeBuffer(b_);
|
||||
bufferManager.FreeBuffer(c_);
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &Get() {
|
||||
if (flag1_ == 0) {
|
||||
flag1_ = 1;
|
||||
return a_;
|
||||
} else if (flag1_ == 1) {
|
||||
flag1_ = NUM_2;
|
||||
return b_;
|
||||
} else {
|
||||
flag1_ = 0;
|
||||
return c_;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetVec() { // mixcore architecture
|
||||
if (flag1_vec1_ == 0) {
|
||||
flag1_vec1_ = 1;
|
||||
return a_;
|
||||
} else if (flag1_vec1_ == 1) {
|
||||
flag1_vec1_ = NUM_2;
|
||||
return b_;
|
||||
} else {
|
||||
flag1_vec1_ = 0;
|
||||
return c_;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetCube() { // mixcore architecture
|
||||
if (flag1_bmm2_ == 0) {
|
||||
flag1_bmm2_ = 1;
|
||||
return a_;
|
||||
} else if (flag1_bmm2_ == 1) {
|
||||
flag1_bmm2_ = NUM_2;
|
||||
return b_;
|
||||
} else {
|
||||
flag1_bmm2_ = 0;
|
||||
return c_;
|
||||
}
|
||||
}
|
||||
|
||||
// Q复用
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetPre() {
|
||||
if (flag1_ == 0) {
|
||||
return c_;
|
||||
} else if (flag1_ == 1) {
|
||||
return a_;
|
||||
} else {
|
||||
return b_;
|
||||
}
|
||||
}
|
||||
|
||||
// KV复用
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetReused() {
|
||||
if (flag2_ == 0) {
|
||||
flag2_ = 1;
|
||||
return a_;
|
||||
} else if (flag2_ == 1){
|
||||
flag2_ = NUM_2;
|
||||
return b_;
|
||||
} else {
|
||||
flag2_ = 0;
|
||||
return c_;
|
||||
}
|
||||
}
|
||||
private:
|
||||
Buffer<bufferType, syncType> a_;
|
||||
Buffer<bufferType, syncType> b_;
|
||||
Buffer<bufferType, syncType> c_;
|
||||
uint32_t flag1_ = 0;
|
||||
uint32_t flag1_vec1_ = 0;
|
||||
uint32_t flag1_bmm2_ = 0;
|
||||
uint32_t flag2_ = 0;
|
||||
};
|
||||
|
||||
// 申请4个buffer + kv复用
|
||||
template<BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
class BuffersPolicy4buff {
|
||||
public:
|
||||
__aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size) {
|
||||
a_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
b_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
c_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
d_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
|
||||
a_.Init();
|
||||
b_.Init();
|
||||
c_.Init();
|
||||
d_.Init();
|
||||
}
|
||||
|
||||
__aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager) {
|
||||
a_.UnInit();
|
||||
b_.UnInit();
|
||||
c_.UnInit();
|
||||
d_.UnInit();
|
||||
|
||||
bufferManager.FreeBuffer(a_);
|
||||
bufferManager.FreeBuffer(b_);
|
||||
bufferManager.FreeBuffer(c_);
|
||||
bufferManager.FreeBuffer(d_);
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &Get(uint32_t id) {
|
||||
uint32_t flag = id % 4;
|
||||
if (flag == 0) {
|
||||
return a_;
|
||||
} else if (flag == 1) {
|
||||
return b_;
|
||||
} else if (flag == 2) { // 2:c_
|
||||
return c_;
|
||||
} else {
|
||||
return d_;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &Get() {
|
||||
auto& buffer = Get(head_);
|
||||
head_++;
|
||||
return buffer;
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetReused() {
|
||||
auto& buffer = Get(used_);
|
||||
used_ = (used_ - tail_ + 1) % (head_ - tail_) + tail_;
|
||||
return buffer;
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetFree() {
|
||||
if (tail_ == used_) {
|
||||
used_++;
|
||||
}
|
||||
auto& buffer = Get(tail_);
|
||||
tail_++;
|
||||
return buffer;
|
||||
}
|
||||
private:
|
||||
Buffer<bufferType, syncType> a_;
|
||||
Buffer<bufferType, syncType> b_;
|
||||
Buffer<bufferType, syncType> c_;
|
||||
Buffer<bufferType, syncType> d_;
|
||||
uint32_t tail_ = 0; // 表示当前正在使用的buffer队列队尾
|
||||
uint32_t head_ = 0; // 表示当前正在使用的buffer队列队首+1
|
||||
uint32_t used_ = 0; // 表示当前正在使用的buffer,于首尾间,左闭右开
|
||||
};
|
||||
|
||||
template<BufferType bufferType, SyncType syncType = SyncType::INNER_CORE_SYNC>
|
||||
class Matrix2x2BufferPolicy { // 4buffer
|
||||
// 二维buffer管理,地址行优先,使用列优先
|
||||
// MracBuffer:memory address with row first, alloc/use/free with column first
|
||||
public:
|
||||
__aicore__ inline void Init(BufferManager<bufferType> &bufferManager, uint32_t size) {
|
||||
bufferM0k0_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
bufferM0k1_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
bufferM1k0_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
bufferM1k1_ = bufferManager.template AllocBuffer<syncType>(size);
|
||||
|
||||
bufferM0k0_.Init();
|
||||
bufferM0k1_.Init();
|
||||
bufferM1k0_.Init();
|
||||
bufferM1k1_.Init();
|
||||
}
|
||||
|
||||
__aicore__ inline void Uninit(BufferManager<bufferType> &bufferManager) {
|
||||
bufferM0k0_.UnInit();
|
||||
bufferM0k1_.UnInit();
|
||||
bufferM1k0_.UnInit();
|
||||
bufferM1k1_.UnInit();
|
||||
|
||||
bufferManager.FreeBuffer(bufferM0k0_);
|
||||
bufferManager.FreeBuffer(bufferM0k1_);
|
||||
bufferManager.FreeBuffer(bufferM1k0_);
|
||||
bufferManager.FreeBuffer(bufferM1k1_);
|
||||
}
|
||||
|
||||
__aicore__ inline void SetMExtent(int32_t mExtent) {
|
||||
aIdx_ = -1;
|
||||
amIdx_ = (amIdx_ + mSize_ - 1) % mSize_; // 翻转 0->1, 1->0
|
||||
akIdx_ = 0;
|
||||
|
||||
uIdx_ = -1;
|
||||
umIdx_ = (umIdx_ + mSize_ - 1) % mSize_;
|
||||
ukIdx_ = 0;
|
||||
|
||||
fIdx_ = -1;
|
||||
fmIdx_ = (fmIdx_ + mSize_ - 1) % mSize_;
|
||||
fkIdx_ = 0;
|
||||
|
||||
mExtent_ = mExtent;
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &AllocNext() {
|
||||
aIdx_++;
|
||||
return GetBuffer(aIdx_, amIdx_, akIdx_);
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &ReuseNext() {
|
||||
uIdx_++;
|
||||
return GetBuffer(uIdx_, umIdx_, ukIdx_);
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &FreeNext() {
|
||||
fIdx_++;
|
||||
return GetBuffer(fIdx_, fmIdx_, fkIdx_);
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &PeekNextK() { // 在Alloc阶段使用,k方向取下一个
|
||||
return PeekBuffer(amIdx_, (1 - akIdx_)); // k翻转
|
||||
}
|
||||
private:
|
||||
__aicore__ inline Buffer<bufferType, syncType> &GetBuffer(int32_t xIdx, int32_t &mIdx, int32_t &kIdx) {
|
||||
// xIdx为入参,表示当前alloc/use/free的idx,mIdx和kIdx为下标出参,移动到下一个buffer并获取
|
||||
mIdx = (mIdx + mExtent_ - 1) % mExtent_;
|
||||
kIdx = (xIdx / mExtent_) % kSize_;
|
||||
if (mIdx == 0 && kIdx == 0) {
|
||||
return bufferM0k0_;
|
||||
} else if (mIdx == 0 && kIdx == 1) {
|
||||
return bufferM0k1_;
|
||||
} else if (mIdx == 1 && kIdx == 0) {
|
||||
return bufferM1k0_;
|
||||
} else { // 该分支条件为:mIdx == 1 && kIdx == 1
|
||||
return bufferM1k1_;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline Buffer<bufferType, syncType> &PeekBuffer(int32_t mIdx, int32_t kIdx) {
|
||||
// 只访问buffer,不进行下标移动
|
||||
if (mIdx == 0 && kIdx == 0) {
|
||||
return bufferM0k0_;
|
||||
} else if (mIdx == 0 && kIdx == 1) {
|
||||
return bufferM0k1_;
|
||||
} else if ((mIdx == 1) && (kIdx == 0)) {
|
||||
return bufferM1k0_;
|
||||
} else { // mIdx == 1 && kIdx == 1
|
||||
return bufferM1k1_;
|
||||
}
|
||||
}
|
||||
|
||||
Buffer<bufferType, syncType> bufferM0k0_;
|
||||
Buffer<bufferType, syncType> bufferM0k1_;
|
||||
Buffer<bufferType, syncType> bufferM1k0_;
|
||||
Buffer<bufferType, syncType> bufferM1k1_;
|
||||
int32_t mSize_ = 2; // m的总buffer数
|
||||
int32_t kSize_ = 2; // k的总buffer数
|
||||
|
||||
// Alloc
|
||||
int32_t aIdx_ = -1; // 当前第几次Alloc Buffer
|
||||
int32_t amIdx_ = 0; // 当前Alloc Buffer的m下标
|
||||
int32_t akIdx_ = 0; // 当前Alloc Buffer的k下标
|
||||
|
||||
// Reuse
|
||||
int32_t uIdx_ = -1;
|
||||
int32_t umIdx_ = 0;
|
||||
int32_t ukIdx_ = 0;
|
||||
|
||||
// Free
|
||||
int32_t fIdx_ = -1;
|
||||
int32_t fmIdx_ = 0;
|
||||
int32_t fkIdx_ = 0;
|
||||
|
||||
int32_t mExtent_ = 0; // m实际使用的大小,可以为1或者2
|
||||
};
|
||||
}
|
||||
#endif
|
||||
1158
csrc/attention/common/op_kernel/matmul.h
Normal file
1158
csrc/attention/common/op_kernel/matmul.h
Normal file
File diff suppressed because it is too large
Load Diff
35
csrc/attention/common/op_kernel/memcopy/fa_gm_tensor.h
Normal file
35
csrc/attention/common/op_kernel/memcopy/fa_gm_tensor.h
Normal file
@@ -0,0 +1,35 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file fa_gm_tensor.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef FA_GM_TENSOR_H
|
||||
#define FA_GM_TENSOR_H
|
||||
|
||||
#if ASC_DEVKIT_MAJOR >= 9
|
||||
#include "kernel_vec_intf.h"
|
||||
#include "kernel_cube_intf.h"
|
||||
#else
|
||||
#include "kernel_operator.h"
|
||||
#endif
|
||||
#include "gm_layout.h"
|
||||
#include "offset_calculator_v2.h"
|
||||
|
||||
using AscendC::GlobalTensor;
|
||||
|
||||
template <typename Q_T, GmFormat FORMAT, typename ACTLEN_T = uint64_t>
|
||||
struct FaGmTensor {
|
||||
GlobalTensor<Q_T> gmTensor;
|
||||
OffsetCalculator<FORMAT, ACTLEN_T> offsetCalculator;
|
||||
};
|
||||
|
||||
#endif
|
||||
43
csrc/attention/common/op_kernel/memcopy/fa_l1_tensor.h
Normal file
43
csrc/attention/common/op_kernel/memcopy/fa_l1_tensor.h
Normal file
@@ -0,0 +1,43 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file fa_l1_tensor.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef FA_L1_TENSOR_H
|
||||
#define FA_L1_TENSOR_H
|
||||
|
||||
#if ASC_DEVKIT_MAJOR >= 9
|
||||
#include "kernel_vec_intf.h"
|
||||
#include "kernel_cube_intf.h"
|
||||
#else
|
||||
#include "kernel_operator.h"
|
||||
#endif
|
||||
|
||||
using AscendC::LocalTensor;
|
||||
|
||||
enum class L1Format {
|
||||
NZ = 0
|
||||
};
|
||||
|
||||
enum class ScaleTrans {
|
||||
NO_TRANS = 0,
|
||||
ND2NZ = 1,
|
||||
DN2NZ = 2
|
||||
};
|
||||
|
||||
template <typename Q_T, L1Format FORMAT>
|
||||
struct FaL1Tensor {
|
||||
LocalTensor<Q_T> tensor;
|
||||
uint32_t rowCount;
|
||||
};
|
||||
|
||||
#endif
|
||||
26
csrc/attention/common/op_kernel/memcopy/gm_coord.h
Normal file
26
csrc/attention/common/op_kernel/memcopy/gm_coord.h
Normal file
@@ -0,0 +1,26 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file gm_coord.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef GM_COORD_H
|
||||
#define GM_COORD_H
|
||||
|
||||
struct GmCoord {
|
||||
uint32_t bIdx;
|
||||
uint32_t n2Idx;
|
||||
uint32_t gS1Idx;
|
||||
uint32_t dIdx;
|
||||
uint32_t gS1DealSize;
|
||||
uint32_t dDealSize;
|
||||
};
|
||||
#endif
|
||||
427
csrc/attention/common/op_kernel/memcopy/gm_layout.h
Normal file
427
csrc/attention/common/op_kernel/memcopy/gm_layout.h
Normal file
@@ -0,0 +1,427 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file gm_layout.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef GM_LAYOUT_H
|
||||
#define GM_LAYOUT_H
|
||||
|
||||
#if ASC_DEVKIT_MAJOR >= 9
|
||||
#include "kernel_vec_intf.h"
|
||||
#include "kernel_cube_intf.h"
|
||||
#else
|
||||
#include "kernel_operator.h"
|
||||
#endif
|
||||
|
||||
// ----------------------------------------------GmLayout--------------------------------
|
||||
enum class GmFormat {
|
||||
BSNGD = 0,
|
||||
BNGSD = 1,
|
||||
NGBSD = 2,
|
||||
TNGD = 3,
|
||||
NGTD = 4,
|
||||
BSND = 5,
|
||||
BNSD = 6,
|
||||
TND = 7,
|
||||
NTD = 8,
|
||||
PA_BnBsND = 9,
|
||||
PA_BnNBsD = 10,
|
||||
PA_NZ = 11,
|
||||
NGD = 12, // post_quant
|
||||
ND = 13, //antiquant no PA
|
||||
BS2 = 14,
|
||||
BNS2 = 15,
|
||||
PA_BnBs = 16, //antiquant PA
|
||||
PA_BnNBs = 17,
|
||||
BN2GS1S2 = 18, //PSE_GmFormat
|
||||
SBNGD = 19,
|
||||
SBND = 20,
|
||||
NTGD = 21,
|
||||
PA_NZ_K_SCALE = 22,
|
||||
};
|
||||
|
||||
template <GmFormat FORMAT>
|
||||
struct GmLayout {
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::BSNGD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s, uint32_t d) {
|
||||
shape = AscendC::MakeShape(b, n, g, s, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t gStride = dStride * d;
|
||||
uint64_t nStride = gStride * g;
|
||||
uint64_t sStride = nStride * n;
|
||||
uint64_t bStride = sStride * s;
|
||||
stride = AscendC::MakeStride(bStride, nStride, gStride, sStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::BNGSD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s, uint32_t d) {
|
||||
shape = AscendC::MakeShape(b, n, g, s, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t sStride = dStride * d;
|
||||
uint64_t gStride = sStride * s;
|
||||
uint64_t nStride = gStride * g;
|
||||
uint64_t bStride = nStride * n;
|
||||
stride = AscendC::MakeStride(bStride, nStride, gStride, sStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::NGBSD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s, uint32_t d) {
|
||||
shape = AscendC::MakeShape(b, n, g, s, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t sStride = dStride * d;
|
||||
uint64_t bStride = sStride * s;
|
||||
uint64_t gStride = bStride * b;
|
||||
uint64_t nStride = gStride * g;
|
||||
stride = AscendC::MakeStride(bStride, nStride, gStride, sStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::TNGD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t g, uint32_t d) {
|
||||
shape = AscendC::MakeShape(t, n, g, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t gStride = dStride * d;
|
||||
uint64_t nStride = gStride * g;
|
||||
uint64_t tStride = nStride * n;
|
||||
stride = AscendC::MakeStride(tStride, nStride, gStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::NGTD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t g, uint32_t d) {
|
||||
shape = AscendC::MakeShape(t, n, g, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t tStride = dStride * d;
|
||||
uint64_t gStride = tStride * t;
|
||||
uint64_t nStride = gStride * g;
|
||||
stride = AscendC::MakeStride(tStride, nStride, gStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::NTGD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t g, uint32_t d) {
|
||||
shape = AscendC::MakeShape(t, n, g, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t gStride = dStride * d;
|
||||
uint64_t tStride = gStride * g;
|
||||
uint64_t nStride = tStride * t;
|
||||
stride = AscendC::MakeStride(tStride, nStride, gStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::BSND> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t s, uint32_t d) {
|
||||
shape = AscendC::MakeShape(b, n, s, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t nStride = dStride * d;
|
||||
uint64_t sStride = nStride * n;
|
||||
uint64_t bStride = sStride * s;
|
||||
stride = AscendC::MakeStride(bStride, nStride, sStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::BNSD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t s, uint32_t d) {
|
||||
shape = AscendC::MakeShape(b, n, s, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t sStride = dStride * d;
|
||||
uint64_t nStride = sStride * s;
|
||||
uint64_t bStride = nStride * n;
|
||||
stride = AscendC::MakeStride(bStride, nStride, sStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::TND> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t d) {
|
||||
shape = AscendC::MakeShape(t, n, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t nStride = dStride * d;
|
||||
uint64_t tStride = nStride * n;
|
||||
stride = AscendC::MakeStride(tStride, nStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::NTD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t t, uint32_t n, uint32_t d) {
|
||||
shape = AscendC::MakeShape(t, n, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t tStride = dStride * d;
|
||||
uint64_t nStride = tStride * t;
|
||||
stride = AscendC::MakeStride(tStride, nStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::PA_BnBsND> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize, uint32_t d) {
|
||||
shape = AscendC::MakeShape(n, blockSize, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t nStride = dStride * d;
|
||||
uint64_t bsStride = nStride * n;
|
||||
uint64_t bnStride = bsStride * blockSize;
|
||||
stride = AscendC::MakeStride(bnStride, nStride, bsStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::PA_BnNBsD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize, uint32_t d) {
|
||||
shape = AscendC::MakeShape(n, blockSize, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t bsStride = dStride * d;
|
||||
uint64_t nStride = bsStride * blockSize;
|
||||
uint64_t bnStride = nStride * n;
|
||||
stride = AscendC::MakeStride(bnStride, nStride, bsStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::PA_NZ> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize, uint32_t d1, uint32_t d0) {
|
||||
shape = AscendC::MakeShape(n, d1, blockSize, d0);
|
||||
uint64_t d0Stride = 1;
|
||||
uint64_t bsStride = d0Stride * d0;
|
||||
uint64_t d1Stride = bsStride * blockSize;
|
||||
uint64_t nStride = d1Stride * d1;
|
||||
uint64_t bnStride = nStride * n;
|
||||
stride = AscendC::MakeStride(bnStride, nStride, d1Stride, bsStride, d0Stride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::PA_NZ_K_SCALE> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize1, uint32_t d, uint32_t blockSize0) {
|
||||
shape = AscendC::MakeShape(n, blockSize1, d, blockSize0);
|
||||
uint64_t bs0Stride = 1;
|
||||
uint64_t dStride = bs0Stride * blockSize0;
|
||||
uint64_t bs1Stride = dStride * d;
|
||||
uint64_t nStride = bs1Stride * blockSize1;
|
||||
uint64_t bnStride = nStride * n;
|
||||
stride = AscendC::MakeStride(bnStride, nStride, bs1Stride, dStride, bs0Stride);
|
||||
}
|
||||
};
|
||||
|
||||
// post_quant
|
||||
template <>
|
||||
struct GmLayout<GmFormat::NGD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t n, uint32_t g, uint32_t d) {
|
||||
shape = AscendC::MakeShape(n, g, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t gStride = dStride * d;
|
||||
uint64_t nStride = gStride * g;
|
||||
stride = AscendC::MakeStride(nStride, gStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
//antiquant
|
||||
template <>
|
||||
struct GmLayout<GmFormat::ND> {
|
||||
AscendC::Shape<uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t n, uint32_t d) {
|
||||
shape = AscendC::MakeShape(n, d);
|
||||
|
||||
uint64_t dStride = 1;
|
||||
uint64_t nStride = dStride * d; //headDim
|
||||
stride = AscendC::MakeStride(nStride, dStride);
|
||||
}
|
||||
};
|
||||
template <>
|
||||
struct GmLayout<GmFormat::BS2> {
|
||||
AscendC::Shape<uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t s) {
|
||||
shape = AscendC::MakeShape(b, s);
|
||||
|
||||
uint64_t sStride = 1;
|
||||
uint64_t bStride = sStride * s;
|
||||
|
||||
stride = AscendC::MakeStride(bStride, sStride);
|
||||
}
|
||||
};
|
||||
template <>
|
||||
struct GmLayout<GmFormat::BNS2> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t s) {
|
||||
shape = AscendC::MakeShape(b, n, s);
|
||||
|
||||
uint64_t sStride = 1;
|
||||
uint64_t nStride = sStride * s;
|
||||
uint64_t bStride = nStride * n;
|
||||
|
||||
stride = AscendC::MakeStride(bStride, nStride, sStride);
|
||||
}
|
||||
};
|
||||
template <>
|
||||
struct GmLayout<GmFormat::PA_BnBs> {
|
||||
AscendC::Shape<uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t blockSize) {
|
||||
shape = AscendC::MakeShape(blockSize);
|
||||
|
||||
uint64_t bsStride = 1;
|
||||
uint64_t bnStride = bsStride * blockSize;
|
||||
stride = AscendC::MakeStride(bnStride, bsStride);
|
||||
}
|
||||
};
|
||||
template <>
|
||||
struct GmLayout<GmFormat::PA_BnNBs> {
|
||||
AscendC::Shape<uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t n, uint32_t blockSize) {
|
||||
shape = AscendC::MakeShape(n, blockSize);
|
||||
|
||||
uint64_t bsStride = 1;
|
||||
uint64_t nStride = bsStride * blockSize;
|
||||
uint64_t bnStride = nStride * n; //blockSize * kvHeadNum
|
||||
stride = AscendC::MakeStride(bnStride, nStride, bsStride);
|
||||
}
|
||||
};
|
||||
|
||||
//PSE_GmLayout
|
||||
template <>
|
||||
struct GmLayout<GmFormat::BN2GS1S2> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s1, uint32_t s2)
|
||||
{
|
||||
shape = AscendC::MakeShape(b, n, g, s1, s2);
|
||||
uint64_t s2Stride = 1;
|
||||
uint64_t s1Stride = s2Stride * s2;
|
||||
uint64_t gStride = s1Stride * s1;
|
||||
uint64_t nStride = gStride * g;
|
||||
uint64_t bStride = nStride * n;
|
||||
stride = AscendC::MakeStride(bStride, nStride, gStride, s1Stride, s2Stride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::SBNGD> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t g, uint32_t s, uint32_t d) {
|
||||
shape = AscendC::MakeShape(b, n, g, s, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t gStride = dStride * d;
|
||||
uint64_t nStride = gStride * g;
|
||||
uint64_t bStride = nStride * n;
|
||||
uint64_t sStride = bStride * b;
|
||||
stride = AscendC::MakeStride(bStride, nStride, gStride, sStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct GmLayout<GmFormat::SBND> {
|
||||
AscendC::Shape<uint32_t, uint32_t, uint32_t, uint32_t> shape;
|
||||
AscendC::Stride<uint64_t, uint64_t, uint64_t, uint64_t> stride;
|
||||
|
||||
__aicore__ inline GmLayout() = default;
|
||||
__aicore__ inline void MakeLayout(uint32_t b, uint32_t n, uint32_t s, uint32_t d) {
|
||||
shape = AscendC::MakeShape(b, n, s, d);
|
||||
uint64_t dStride = 1;
|
||||
uint64_t nStride = dStride * d;
|
||||
uint64_t bStride = nStride * n;
|
||||
uint64_t sStride = bStride * b;
|
||||
stride = AscendC::MakeStride(bStride, nStride, sStride, dStride);
|
||||
}
|
||||
};
|
||||
|
||||
#endif
|
||||
1104
csrc/attention/common/op_kernel/memcopy/offset_calculator_v2.h
Normal file
1104
csrc/attention/common/op_kernel/memcopy/offset_calculator_v2.h
Normal file
File diff suppressed because it is too large
Load Diff
140
csrc/attention/common/op_kernel/memcopy/parser.h
Normal file
140
csrc/attention/common/op_kernel/memcopy/parser.h
Normal file
@@ -0,0 +1,140 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file parser.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef PARSER_H
|
||||
#define PARSER_H
|
||||
|
||||
#if ASC_DEVKIT_MAJOR >= 9
|
||||
#include "kernel_vec_intf.h"
|
||||
#include "kernel_cube_intf.h"
|
||||
#else
|
||||
#include "kernel_operator.h"
|
||||
#endif
|
||||
|
||||
using AscendC::GlobalTensor;
|
||||
|
||||
// ----------------------------------------------ActualSeqLensParser--------------------------------
|
||||
enum class ActualSeqLensMode
|
||||
{
|
||||
BY_BATCH = 0,
|
||||
ACCUM = 1,
|
||||
};
|
||||
|
||||
template <ActualSeqLensMode MODE, typename ACTLEN_T = uint64_t>
|
||||
class ActualSeqLensParser {
|
||||
};
|
||||
|
||||
template <typename ACTLEN_T>
|
||||
class ActualSeqLensParser<ActualSeqLensMode::ACCUM, ACTLEN_T> {
|
||||
public:
|
||||
__aicore__ inline ActualSeqLensParser() = default;
|
||||
|
||||
__aicore__ inline void Init(GlobalTensor<ACTLEN_T> actualSeqLengthsGm, uint32_t actualLenDims,
|
||||
uint64_t defaultVal = 0)
|
||||
{
|
||||
this->actualSeqLengthsGm = actualSeqLengthsGm;
|
||||
this->actualLenDims = actualLenDims;
|
||||
}
|
||||
|
||||
__aicore__ inline uint64_t GetTBase(uint32_t bIdx) const
|
||||
{
|
||||
if (bIdx == 0) {
|
||||
return 0;
|
||||
}
|
||||
return actualSeqLengthsGm.GetValue(bIdx - 1);
|
||||
}
|
||||
|
||||
__aicore__ inline uint64_t GetMxVscaleTBase(uint32_t bIdx) const
|
||||
{
|
||||
if (bIdx == 0) {
|
||||
return 0;
|
||||
}
|
||||
uint64_t vScaleTBaseOffset = 0;
|
||||
for (uint32_t idx = 0; idx < bIdx; idx++) {
|
||||
vScaleTBaseOffset += ((GetActualSeqLength(idx) + 63) >> 6);
|
||||
}
|
||||
return vScaleTBaseOffset;
|
||||
}
|
||||
|
||||
__aicore__ inline uint64_t GetActualSeqLength(uint32_t bIdx) const
|
||||
{
|
||||
if (bIdx == 0) {
|
||||
return actualSeqLengthsGm.GetValue(0);
|
||||
}
|
||||
return (actualSeqLengthsGm.GetValue(bIdx) - actualSeqLengthsGm.GetValue(bIdx - 1));
|
||||
}
|
||||
|
||||
__aicore__ inline uint64_t GetTSize() const
|
||||
{
|
||||
return actualSeqLengthsGm.GetValue(actualLenDims - 1);
|
||||
}
|
||||
private:
|
||||
GlobalTensor<ACTLEN_T> actualSeqLengthsGm;
|
||||
uint32_t actualLenDims;
|
||||
};
|
||||
|
||||
template <typename ACTLEN_T>
|
||||
class ActualSeqLensParser<ActualSeqLensMode::BY_BATCH, ACTLEN_T> {
|
||||
public:
|
||||
__aicore__ inline ActualSeqLensParser() = default;
|
||||
|
||||
__aicore__ inline void Init(GlobalTensor<ACTLEN_T> actualSeqLengthsGm, uint32_t actualLenDims, uint64_t defaultVal)
|
||||
{
|
||||
this->actualSeqLengthsGm = actualSeqLengthsGm;
|
||||
this->actualLenDims = actualLenDims;
|
||||
this->defaultVal = defaultVal;
|
||||
}
|
||||
|
||||
__aicore__ inline uint64_t GetActualSeqLength(uint32_t bIdx) const
|
||||
{
|
||||
if (actualLenDims == 0) {
|
||||
return defaultVal;
|
||||
}
|
||||
if (actualLenDims == 1) {
|
||||
return actualSeqLengthsGm.GetValue(0);
|
||||
}
|
||||
return actualSeqLengthsGm.GetValue(bIdx);
|
||||
}
|
||||
|
||||
__aicore__ inline uint32_t GetActualLenDims() const
|
||||
{
|
||||
return actualLenDims;
|
||||
}
|
||||
private:
|
||||
GlobalTensor<ACTLEN_T> actualSeqLengthsGm;
|
||||
uint32_t actualLenDims = 0;
|
||||
uint64_t defaultVal = 0;
|
||||
};
|
||||
|
||||
// ----------------------------------------------BlockTableParser--------------------------------
|
||||
class BlockTableParser {
|
||||
public:
|
||||
__aicore__ inline BlockTableParser() = default;
|
||||
|
||||
__aicore__ inline void Init(GlobalTensor<int32_t> blockTableGm, uint32_t maxblockNumPerBatch)
|
||||
{
|
||||
this->blockTableGm = blockTableGm;
|
||||
this->maxblockNumPerBatch = maxblockNumPerBatch;
|
||||
}
|
||||
|
||||
__aicore__ inline int32_t GetBlockIdx(uint32_t bIdx, uint32_t blockIdxInBatch) const
|
||||
{
|
||||
return blockTableGm.GetValue(bIdx * maxblockNumPerBatch + blockIdxInBatch);
|
||||
}
|
||||
private:
|
||||
GlobalTensor<int32_t> blockTableGm;
|
||||
uint32_t maxblockNumPerBatch;
|
||||
};
|
||||
|
||||
#endif
|
||||
31
csrc/attention/common/op_kernel/offset_calculator.h
Normal file
31
csrc/attention/common/op_kernel/offset_calculator.h
Normal file
@@ -0,0 +1,31 @@
|
||||
/**
|
||||
* Copyright (c) 2025 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file offset_calculator.h
|
||||
* \brief
|
||||
*/
|
||||
#ifndef OFFSET_CALCULATOR_H
|
||||
#define OFFSET_CALCULATOR_H
|
||||
|
||||
#if ASC_DEVKIT_MAJOR >= 9
|
||||
#include "kernel_basic_intf.h"
|
||||
#else
|
||||
#include "kernel_operator.h"
|
||||
#endif
|
||||
|
||||
#include "memcopy/gm_layout.h"
|
||||
#include "memcopy/parser.h"
|
||||
#include "memcopy/offset_calculator_v2.h"
|
||||
#include "memcopy/fa_gm_tensor.h"
|
||||
#include "memcopy/fa_l1_tensor.h"
|
||||
#include "memcopy/gm_coord.h"
|
||||
|
||||
#endif
|
||||
19
csrc/attention/compressor/CMakeLists.txt
Normal file
19
csrc/attention/compressor/CMakeLists.txt
Normal file
@@ -0,0 +1,19 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
|
||||
file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
|
||||
if(NOT ENABLE_TEST AND NOT BENCHMARK)
|
||||
list(REMOVE_ITEM CURRENT_DIRS tests)
|
||||
endif()
|
||||
foreach(SUB_DIR ${CURRENT_DIRS})
|
||||
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
|
||||
add_subdirectory(${SUB_DIR})
|
||||
endif()
|
||||
endforeach()
|
||||
499
csrc/attention/compressor/README.md
Normal file
499
csrc/attention/compressor/README.md
Normal file
@@ -0,0 +1,499 @@
|
||||
# Compressor
|
||||
|
||||
## 产品支持情况
|
||||
|
||||
| 产品 | 是否支持 |
|
||||
| ------------------------------------------------------------ | :------: |
|
||||
|<term>Ascend 950PR/Ascend 950DT</term>| √ |
|
||||
|<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>| √ |
|
||||
|<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>| × |
|
||||
|<term>Atlas 200I/500 A2 推理产品</term>| × |
|
||||
|<term>Atlas 推理系列加速卡产品</term>| × |
|
||||
|<term>Atlas 训练系列产品</term>| × |
|
||||
|
||||
## 功能说明
|
||||
|
||||
- API功能:Compressor是推理场景下SAS和QLI的前处理算子,用于将每4或128个token的KV cache压缩成一个,然后每个token与这些压缩的KV cache进行DSA计算。在长序列的情况下,Compressor可以有效地减少计算开销。
|
||||
|
||||
- 计算公式:
|
||||
|
||||
压缩阶段:
|
||||
1. 计算矩阵乘法:
|
||||
- C4A: $\left[kv\_state^a, score\_state^a\right] = X @ \left[W^{aKV}, W^{aGate}\right], \left[kv\_state^b, score\_state^b\right] = X @ \left[W^{bKV}, W^{bGate}\right];$
|
||||
- C128A: $\left[kv\_state, score\_state\right] = X @ \left[W^{KV}, W^{Gate}\right]$
|
||||
2. 计算分组加法:
|
||||
- C4A: $score\_state_i^\prime = \left[score\_state_{\left[4(i-1)+1:4i,:\right]}^a; score\_state_{\left[4i+1:4(i+1),:\right]}^b\right] + Ape,~i=1,2,\cdots, \frac{s}{4};$
|
||||
- C128A: $score\_state_i^\prime = score\_state_{\left[128(i-1)+1:128i,:\right]} + Ape,~i=1,2,\cdots, \frac{s}{128};$
|
||||
3. 计算分组Softmax:
|
||||
- C4A: $S_i^\prime = softmax(score\_state_i^\prime),~i=1,2,\cdots, \frac{s}{4};$
|
||||
- C128A: $S_i^\prime = softmax(score\_state_i^\prime),~i=1,2,\cdots, \frac{s}{128};$
|
||||
4. 计算Hadamard乘积:
|
||||
- C4A: $(S_H)_i = S_i^\prime \odot \left[kv\_state^a_{\left[4(i-1)+1:4i,:\right]} ; kv\_state^b_{\left[4i+1:4(i+1),:\right]}\right],~i=1,2,\cdots, \frac{s}{4};$
|
||||
- C128A: $S_H = S_i^\prime \odot kv\_state;$
|
||||
5. 沿着压缩轴分组求和:
|
||||
- C4A: $C_{i}^{\text{Comp}} = \left[1\right]_{1\times8} @ (S_H)_i, ~i=1,2,\cdots, \frac{s}{4};$
|
||||
- C128A: $C_{i}^{\text{Comp}} = \left[1\right]_{1\times128} @ (S_H)_i, ~i=1,2,\cdots, \frac{s}{128};$
|
||||
|
||||
后处理阶段:
|
||||
|
||||
6. 计算RMSNorm:
|
||||
- $\text{RMS}(C^{\text{Comp}}) = \sqrt{\frac{1}{N} \sum_{i=j* N}^{(j+1)* N} {(C_{i}^{\text{Comp}})}^{\text{2}} + norm\_eps} ,N=head\_dim, ~j=1,2,\cdots, \frac{s}{cmp\_ratio}$
|
||||
- $\text{RmsNorm}(C^{\text{Comp}}) = norm\_weight \cdot \frac{C_{i}^{\text{Comp}}}{\text{RMS}(C^{\text{Comp}})}$
|
||||
7. 计算Rope;
|
||||
|
||||
- 主要计算过程为:
|
||||
1. 将输入$X$与$W^{KV}$做Matmul运算得到$kv\_state$,将输入$X$与$W^{Gate}$做Matmul运算后再与$Ape$做Add运算得到$score\_state$,$kv\_state$与$score\_state$根据输入的start_pos及cu_seqlens完成更新。
|
||||
2. 在coff为2的情况下对$kv\_state$和$score\_state$进行数据重排。
|
||||
3. 对$score\_state$进行softmax运算将softmax结果与$kv\_state$做Mul计算,后进行ReduceSum运算。
|
||||
4. 根据输入数据norm_weight、rope_sin、rope_cos,进行RMSNorm和Rope运算,得到$cmp\_kv$结果输出。
|
||||
|
||||
## 参数说明
|
||||
|
||||
| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |
|
||||
|----------------------------|-----------|----------------------------------------------------------------------|----------------|------------|
|
||||
| x | 输入 | 公式中的$X$,表示原始不经压缩的数据。 | FLOAT16、BFLOAT16 | ND |
|
||||
| wkv | 输入 | 公式中的$W^{KV}$,表示kv压缩权重。 | FLOAT16、BFLOAT16 | ND |
|
||||
| wgate | 输入 | 公式中的$W^{Gate}$,表示gate压缩权重。 | FLOAT16、BFLOAT16 | ND |
|
||||
| kv_state | 输入 | 公式中的$kv\_state$,表示kv\_state的历史数据。 | FLOAT32 | ND |
|
||||
| score_state | 输入 | 公式中的$score\_state$,表示score\_state中的历史数据。 | FLOAT32 | ND |
|
||||
| ape | 输入 | 公式中的$Ape$,表示positional biases。 | FLOAT32 | ND |
|
||||
| norm\_weight | 输入 | 表示计算RmsNorm时的权重系数。 | FLOAT16、BFLOAT16 | ND |
|
||||
| rope\_sin | 输入 | 表示Rope计算时sin的权重系数。 | FLOAT16、BFLOAT16 | ND |
|
||||
| rope\_cos | 输入 | 表示Rope计算时cos的权重系数。 | FLOAT16、BFLOAT16 | ND |
|
||||
| rope\_head\_dim | 属性 | 表示rope_cos和rope_sin的hidden层最小单元大小,当前仅支持64。 | INT32 | - |
|
||||
| cmp\_ratio | 属性 | 用于稀疏计算,表示数据压缩率。 | INT32 | - |
|
||||
| kv\_block\_table | 可选输入 | 表示kv\_state存储使用的block映射表。当其中元素的值为0时,表示当前位置无需进行更新kv\_state操作。 | INT32 | ND |
|
||||
| score\_block\_table | 可选输入 | 表示score\_state存储使用的block映射表。当其中元素的值为0时,表示当前位置无需进行更新score\_state操作。 | INT32 | ND |
|
||||
| cu\_seqlens | 可选输入 | 表示不同Batch中的有效token数。 | INT32 | ND |
|
||||
| seqused | 可选输入 | 表示不同Batch中实际参与压缩的token数,如果指定为None时,表示和每个Batch上的Sequence Length长度相同。 | INT32 | ND |
|
||||
| start\_pos | 可选输入 | 表示计算起始位置。 | INT32 | ND |
|
||||
| coff | 可选属性 | 默认值1,支持1/2。当coff=1时,无需进行overlap数据重排。当coff=2时,需要进行overlap数据重排。 | INT32 | - |
|
||||
| norm\_eps | 可选属性 | 表示RmsNorm计算的权重系数。默认值1e-6。 | FLOAT32 | - |
|
||||
| rotary\_mode | 可选属性 | 表示Rop计算的模式。默认值1,支持1/2。rotary\_mode为1时,代表half模式。rotary\_mode为2时,代表interleave模式。 | INT32 | - |
|
||||
| enabled\_grad | 可选属性 | 训练场景使用,表示是否参与反向更新。默认值false,支持false/true。**目前暂不支持输入true**。 | BOOL | - |
|
||||
| cmp\_kv | 输出 | 表示压缩后的数据。 | FLOAT16、BFLOAT16 | ND |
|
||||
| wkv\_proj | 可选输出 | 训练反向使用,表示wkv权重Matmul的计算结果,**目前暂不支持返回wkv\_proj**。 | FLOAT16、BFLOAT16 | ND |
|
||||
| softmax\_res | 可选输出 | 训练反向使用,表示Softmax计算结果,**目前暂不支持返回softmax\_res**。 | FLOAT16、BFLOAT16 | ND |
|
||||
| norm\_x | 可选输出 | 训练反向使用,表示Rms计算的输入,**目前暂不支持返回norm\_x**。 | FLOAT16、BFLOAT16 | ND |
|
||||
| norm\_rstd | 可选输出 | 训练反向使用,表示Rms计算的中间结果,**目前暂不支持返回norm\_rstd**。 | FLOAT16、BFLOAT16 | ND |
|
||||
|
||||
## 约束说明
|
||||
|
||||
- x参数维度含义:B(Batch Size)表示输入样本批量大小、S(Sequence Length)表示输入样本序列长度、H(Head Size)表示hidden层的大小、D(Head Dim)表示hidden层的最小单元大小、T表示所有Batch输入样本序列长度的累加和。
|
||||
- 输入shape限制:
|
||||
- wkv支持输入shape[coff* D,H]
|
||||
- wgate支持输入shape[coff* D,H]
|
||||
- kv\_state、score\_state支持输入shape[block_num,block_size,coff* D],要求block_num>0。
|
||||
- ape支持输入shape[cmp_ratio,coff* D]
|
||||
- norm\_weight支持输入shape[D,]
|
||||
- start\_pos支持输入shape[B,]
|
||||
- 若x的维度采用BS合轴,即x的输入shape为[T,H]
|
||||
- rope_sin、rope_cos要求输入shape为[min(T,T//cmp_ratio+B),rope_head_dim]。
|
||||
- cu\_seqlens输入shape必须为[B+1,]。该参数中每个元素的值表示当前batch与之前所有batch的token数总和,即前缀和,因此后一个元素的值必须大于等于前一个元素的值,且第一位必须位0。
|
||||
- seqused,支持输入shape[B,],要求每个Batch的有效token数要求小于等于对应Sequence Length长度,即seqused[n] <= cu\_seqlens[n+1] - cu\_seqlens[n],且不小于0。
|
||||
- kv\_block\_table、score\_block\_table支持输入shape[B,ceil(Smax/block_size)]。Smax为每个Batch中最大的Sequence Length,即Smax=max(start\_pos)+max(cu\_seqlens[n+1] - cu\_seqlens[n])。
|
||||
- cmp\_kv,输出shape为[min(T,T//cmp_ratio+B),D]:<batch0>compressed_tokens + <batch1>compressed_tokens + ... + <batchN>compressed_tokens + pad。
|
||||
- wkv\_proj,输出shape为[T,coff* D]。
|
||||
- norm\_x,输出shape为[min(T,T//cmp_ratio+B),D]。
|
||||
- norm\_rstd,输出shape为[min(T,T//cmp_ratio+B)]。
|
||||
- 若x的维度不采用BS合轴,即x的输入shape为[B,S,H]
|
||||
- rope_sin、rope_cos要求输入shape为[B,ceil(S/cmp_ratio),rope_head_dim]。
|
||||
- cu\_seqlens,参数必须为空。
|
||||
- seqused,支持输入shape[B,],要求每个Batch的有效token数要求小于等于对应Sequence Length长度,即要求seqused[n] <= S,且不小于0。
|
||||
- kv\_block\_table、score\_block\_table支持输入shape[B,ceil(Smax/block_size)]。Smax为每个Batch中最大的Sequence Length,即Smax=max(start\_pos)+S。
|
||||
- cmp\_kv,输出shape为[B,ceil(S/cmp_ratio),D]:(<batch0>compressed_tokens+pad0) + (<batch1>compressed_tokens+pad1) + ... + (<batchN>compressed_tokens+padN)。
|
||||
- wkv\_proj,输出shape为[B,S,coff* D]。
|
||||
- norm\_x,输出shape为[B,ceil(S/cmp_ratio),D]。
|
||||
- norm\_rstd,输出shape为[B,ceil(S/cmp_ratio)]。
|
||||
- 输入值域限制:
|
||||
- 该接口支持B、S泛化,且存在如下场景限制:
|
||||
- 部分长序列场景下,如果计算量过大可能会导致出现超过NPU内存的报错,注:这里计算量会受x输入shape的影响,值越大计算量越大。典型的长序列(即B、S的乘积或T较大)场景包括但不限于:
|
||||
<div style="overflow-x: auto;">
|
||||
<table style="undefined;table-layout: fixed; width: 400px"><colgroup>
|
||||
<col style="width: 100px">
|
||||
<col style="width: 100px">
|
||||
</colgroup><thead>
|
||||
<tr>
|
||||
<th>B</th>
|
||||
<th>S</th>
|
||||
<th>H</th>
|
||||
</tr></thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>100</td>
|
||||
<td>65525</td>
|
||||
<td>4096</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>25</td>
|
||||
<td>261120</td>
|
||||
<td>4096</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>100</td>
|
||||
<td>131072</td>
|
||||
<td>4096</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>100</td>
|
||||
<td>261120</td>
|
||||
<td>4096</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
- 输入属性限制:
|
||||
- 支持D为128/512。
|
||||
- 支持H为1K~10K,512对齐。
|
||||
- 泛化支持block_size小于等于1024,16对齐。
|
||||
- 支持cmp_ratio为4/128。支持如下三种情况:
|
||||
- C4A: D=512, coff=2, cmp_ratio=4;
|
||||
- C4Li: D=128, coff=2, cmp_ratio=4;
|
||||
- C128A: D=512, coff=1, cmp_ratio=128。
|
||||
- 支持rotary_mode为2,Rope计算模式为interleave。
|
||||
|
||||
## Atlas A3 推理系列产品 调用说明
|
||||
|
||||
- 单算子模式调用
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch_npu
|
||||
import numpy as np
|
||||
import custom_ops
|
||||
import torch.nn as nn
|
||||
import math
|
||||
|
||||
def get_seq_used_by_batch(batch_idx, S, seqused, cu_seqlens):
|
||||
if seqused is not None:
|
||||
return seqused[batch_idx]
|
||||
else:
|
||||
if cu_seqlens is not None:
|
||||
return cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]
|
||||
else:
|
||||
return S
|
||||
|
||||
data_type = torch.bfloat16
|
||||
hidden_size = 4096
|
||||
rope_head_dim = 64
|
||||
norm_eps = 1e-6
|
||||
coff = 1 # 1:no overlap 2:overlap
|
||||
cmp_ratio = 128
|
||||
rotary_mode = 2
|
||||
head_dim = 512
|
||||
cu_seqlens = [0, 1]
|
||||
# -------------
|
||||
B = 1
|
||||
S = 1
|
||||
S_max = 0
|
||||
block_size = 128
|
||||
start_pos = [8191] * B # (B,)
|
||||
start_p=8191
|
||||
seqused = None # (B,), None时cu_seqlens的数据全部参与计算,否则按传参实际值计算
|
||||
|
||||
# BS是否合轴
|
||||
bs_combine_flag = True
|
||||
update_flag = 1
|
||||
|
||||
if seqused is not None:
|
||||
seqused = torch.tensor(seqused).to(torch.int32)
|
||||
if start_pos is not None:
|
||||
start_pos = torch.tensor(start_pos).to(torch.int32)
|
||||
else:
|
||||
start_pos = torch.full((B,), start_p, dtype=torch.int32)
|
||||
|
||||
if bs_combine_flag:
|
||||
if cu_seqlens is None:
|
||||
T = B * S
|
||||
if T !=0:
|
||||
cu_seqlens = torch.arange(0, T + 1, S, dtype=torch.int32)
|
||||
else:
|
||||
cu_seqlens = torch.zeros((B+1), dtype=torch.int32)
|
||||
else:
|
||||
cu_seqlens = torch.tensor(cu_seqlens).to(torch.int32)
|
||||
for i in range(B):
|
||||
if start_pos[i] + cu_seqlens[i + 1] - cu_seqlens[i] > S_max:
|
||||
S_max = start_pos[i] + cu_seqlens[i + 1] - cu_seqlens[i]
|
||||
else:
|
||||
cu_seqlens = None
|
||||
S_max = max(start_pos) + S
|
||||
### ======================== gen input data start =============================
|
||||
# page state
|
||||
max_block_num_per_batch = (S_max + block_size - 1) // block_size
|
||||
block_num = B * max_block_num_per_batch
|
||||
next_block_id = 1
|
||||
print(f"max_block_num_per_batch: {max_block_num_per_batch}")
|
||||
block_table = torch.zeros(size=(B, max_block_num_per_batch), dtype=torch.int32)
|
||||
for i in range(B):
|
||||
# 需要读取state的范围
|
||||
cur_start = start_pos[i] // cmp_ratio * cmp_ratio - cmp_ratio
|
||||
cur_end = start_pos[i] // cmp_ratio * cmp_ratio + cmp_ratio
|
||||
if start_pos[i] % cmp_ratio == 0:
|
||||
cur_end = start_pos[i]
|
||||
cur_end = min(cur_end, start_pos[i] + S)
|
||||
cur_start_block_id = (cur_start // block_size) if cur_start >= 0 else 0
|
||||
cur_end_block_id = (cur_end - 1) // block_size
|
||||
for j in range(cur_start_block_id, cur_end_block_id + 1):
|
||||
block_table[i][j] = next_block_id
|
||||
next_block_id = next_block_id + 1
|
||||
# 需要写入state的范围
|
||||
end_pos = get_seq_used_by_batch(i, S, seqused, cu_seqlens)
|
||||
next_start = (start_pos[i] + end_pos) // cmp_ratio * cmp_ratio - cmp_ratio
|
||||
next_end = (start_pos[i] + end_pos) // cmp_ratio * cmp_ratio + cmp_ratio
|
||||
if (start_pos[i] + end_pos) % cmp_ratio == 0:
|
||||
next_end = start_pos[i] + end_pos
|
||||
next_end = min(next_end, start_pos[i] + end_pos)
|
||||
next_start_block_id = (next_start // block_size) if next_start >= 0 else 0
|
||||
next_end_block_id = (next_end - 1) // block_size
|
||||
for j in range(next_start_block_id, next_end_block_id + 1):
|
||||
if block_table[i][j] == 0:
|
||||
block_table[i][j] = next_block_id
|
||||
next_block_id = next_block_id + 1
|
||||
|
||||
if B==0:
|
||||
kv_state = torch.tensor(np.random.uniform(-10, 10, (0, block_size, coff * head_dim))).to(torch.float32)
|
||||
score_state = torch.tensor(np.random.uniform(-10, 10, (0, block_size, coff * head_dim))).to(torch.float32)
|
||||
else:
|
||||
kv_state = torch.tensor(np.random.uniform(-10, 10, (torch.max(block_table) + 1, block_size, coff * head_dim))).to(torch.float32)
|
||||
score_state = torch.tensor(np.random.uniform(-10, 10, (torch.max(block_table) + 1, block_size, coff * head_dim))).to(torch.float32)
|
||||
|
||||
# other input
|
||||
if bs_combine_flag:
|
||||
x_shape = (cu_seqlens[-1], hidden_size)
|
||||
rope_sin_shape = (min(x_shape[0], x_shape[0] // cmp_ratio + B), rope_head_dim)
|
||||
rope_cos_shape = rope_sin_shape
|
||||
else:
|
||||
x_shape = (B, S, hidden_size)
|
||||
rope_sin_shape = (B, (S + cmp_ratio - 1) // cmp_ratio, rope_head_dim)
|
||||
rope_cos_shape = rope_sin_shape
|
||||
|
||||
x = torch.tensor(np.random.uniform(-10.0, 10.0, x_shape)).to(data_type).npu()
|
||||
wkv = torch.tensor(np.random.uniform(-10, 10, (coff * head_dim, hidden_size))).to(data_type).npu()
|
||||
wgate = torch.tensor(np.random.uniform(-10, 10, (coff * head_dim, hidden_size))).to(data_type).npu()
|
||||
ape = torch.tensor(np.random.uniform(-10, 10, (cmp_ratio, coff * head_dim))).to(torch.float32).npu()
|
||||
norm_weight = torch.tensor(np.random.uniform(-10, 10, (head_dim))).to(data_type).npu()
|
||||
rope_sin = torch.tensor(np.random.uniform(-1, 1, rope_sin_shape)).to(data_type).npu()
|
||||
rope_cos = torch.tensor(np.random.uniform(-1, 1, rope_cos_shape)).to(data_type).npu()
|
||||
kv_state = kv_state.npu()
|
||||
score_state = score_state.npu()
|
||||
block_table = block_table.npu()
|
||||
start_pos = torch.tensor(start_pos).to(torch.int32).npu()
|
||||
if cu_seqlens is not None:
|
||||
cu_seqlens = torch.tensor(cu_seqlens).to(torch.int32).npu()
|
||||
if seqused is not None:
|
||||
seqused = torch.tensor(seqused).to(torch.int32).npu()
|
||||
|
||||
cmp_kv,_ ,_ ,_ ,_ = (
|
||||
torch.ops.custom.compressor(
|
||||
x,
|
||||
wkv,
|
||||
wgate,
|
||||
kv_state,
|
||||
score_state,
|
||||
ape,
|
||||
norm_weight,
|
||||
rope_sin,
|
||||
rope_cos,
|
||||
kv_block_table = block_table,
|
||||
score_block_table = block_table,
|
||||
cu_seqlens = cu_seqlens,
|
||||
seqused = seqused,
|
||||
start_pos = start_pos,
|
||||
rope_head_dim = rope_head_dim,
|
||||
cmp_ratio = cmp_ratio,
|
||||
coff = coff,
|
||||
norm_eps = norm_eps,
|
||||
rotary_mode = rotary_mode
|
||||
)
|
||||
)
|
||||
```
|
||||
- aclgraph调用
|
||||
|
||||
```python
|
||||
import torch
|
||||
import torch_npu
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
import torchair
|
||||
import custom_ops
|
||||
import math
|
||||
|
||||
def get_seq_used_by_batch(batch_idx, S, seqused, cu_seqlens):
|
||||
if seqused is not None:
|
||||
return seqused[batch_idx]
|
||||
else:
|
||||
if cu_seqlens is not None:
|
||||
return cu_seqlens[batch_idx + 1] - cu_seqlens[batch_idx]
|
||||
else:
|
||||
return S
|
||||
|
||||
data_type = torch.bfloat16
|
||||
hidden_size = 4096
|
||||
rope_head_dim = 64
|
||||
norm_eps = 1e-6
|
||||
coff = 1 # 1:no overlap 2:overlap
|
||||
cmp_ratio = 128
|
||||
rotary_mode = 2
|
||||
head_dim = 512
|
||||
cu_seqlens = [0, 1]
|
||||
# -------------
|
||||
B = 1
|
||||
S = 1
|
||||
S_max = 0
|
||||
block_size = 128
|
||||
start_pos = [8191] * B # (B,)
|
||||
start_p=8191
|
||||
seqused = None # (B,), None时cu_seqlens的数据全部参与计算,否则按传参实际值计算
|
||||
|
||||
# BS是否合轴
|
||||
bs_combine_flag = True
|
||||
update_flag = 1
|
||||
|
||||
if seqused is not None:
|
||||
seqused = torch.tensor(seqused).to(torch.int32)
|
||||
if start_pos is not None:
|
||||
start_pos = torch.tensor(start_pos).to(torch.int32)
|
||||
else:
|
||||
start_pos = torch.full((B,), start_p, dtype=torch.int32)
|
||||
|
||||
if bs_combine_flag:
|
||||
if cu_seqlens is None:
|
||||
T = B * S
|
||||
if T !=0:
|
||||
cu_seqlens = torch.arange(0, T + 1, S, dtype=torch.int32)
|
||||
else:
|
||||
cu_seqlens = torch.zeros((B+1), dtype=torch.int32)
|
||||
else:
|
||||
cu_seqlens = torch.tensor(cu_seqlens).to(torch.int32)
|
||||
for i in range(B):
|
||||
if start_pos[i] + cu_seqlens[i + 1] - cu_seqlens[i] > S_max:
|
||||
S_max = start_pos[i] + cu_seqlens[i + 1] - cu_seqlens[i]
|
||||
else:
|
||||
cu_seqlens = None
|
||||
S_max = max(start_pos) + S
|
||||
### ======================== gen input data start =============================
|
||||
# page state
|
||||
max_block_num_per_batch = (S_max + block_size - 1) // block_size
|
||||
block_num = B * max_block_num_per_batch
|
||||
next_block_id = 1
|
||||
print(f"max_block_num_per_batch: {max_block_num_per_batch}")
|
||||
block_table = torch.zeros(size=(B, max_block_num_per_batch), dtype=torch.int32)
|
||||
for i in range(B):
|
||||
# 需要读取state的范围
|
||||
cur_start = start_pos[i] // cmp_ratio * cmp_ratio - cmp_ratio
|
||||
cur_end = start_pos[i] // cmp_ratio * cmp_ratio + cmp_ratio
|
||||
if start_pos[i] % cmp_ratio == 0:
|
||||
cur_end = start_pos[i]
|
||||
cur_end = min(cur_end, start_pos[i] + S)
|
||||
cur_start_block_id = (cur_start // block_size) if cur_start >= 0 else 0
|
||||
cur_end_block_id = (cur_end - 1) // block_size
|
||||
for j in range(cur_start_block_id, cur_end_block_id + 1):
|
||||
block_table[i][j] = next_block_id
|
||||
next_block_id = next_block_id + 1
|
||||
# 需要写入state的范围
|
||||
end_pos = get_seq_used_by_batch(i, S, seqused, cu_seqlens)
|
||||
next_start = (start_pos[i] + end_pos) // cmp_ratio * cmp_ratio - cmp_ratio
|
||||
next_end = (start_pos[i] + end_pos) // cmp_ratio * cmp_ratio + cmp_ratio
|
||||
if (start_pos[i] + end_pos) % cmp_ratio == 0:
|
||||
next_end = start_pos[i] + end_pos
|
||||
next_end = min(next_end, start_pos[i] + end_pos)
|
||||
next_start_block_id = (next_start // block_size) if next_start >= 0 else 0
|
||||
next_end_block_id = (next_end - 1) // block_size
|
||||
for j in range(next_start_block_id, next_end_block_id + 1):
|
||||
if block_table[i][j] == 0:
|
||||
block_table[i][j] = next_block_id
|
||||
next_block_id = next_block_id + 1
|
||||
|
||||
if B==0:
|
||||
kv_state = torch.tensor(np.random.uniform(-10, 10, (0, block_size, coff * head_dim))).to(torch.float32)
|
||||
score_state = torch.tensor(np.random.uniform(-10, 10, (0, block_size, coff * head_dim))).to(torch.float32)
|
||||
else:
|
||||
kv_state = torch.tensor(np.random.uniform(-10, 10, (torch.max(block_table) + 1, block_size, coff * head_dim))).to(torch.float32)
|
||||
score_state = torch.tensor(np.random.uniform(-10, 10, (torch.max(block_table) + 1, block_size, coff * head_dim))).to(torch.float32)
|
||||
|
||||
# other input
|
||||
if bs_combine_flag:
|
||||
x_shape = (cu_seqlens[-1], hidden_size)
|
||||
rope_sin_shape = (min(x_shape[0], x_shape[0] // cmp_ratio + B), rope_head_dim)
|
||||
rope_cos_shape = rope_sin_shape
|
||||
else:
|
||||
x_shape = (B, S, hidden_size)
|
||||
rope_sin_shape = (B, (S + cmp_ratio - 1) // cmp_ratio, rope_head_dim)
|
||||
rope_cos_shape = rope_sin_shape
|
||||
|
||||
x = torch.tensor(np.random.uniform(-10.0, 10.0, x_shape)).to(data_type).npu()
|
||||
wkv = torch.tensor(np.random.uniform(-10, 10, (coff * head_dim, hidden_size))).to(data_type).npu()
|
||||
wgate = torch.tensor(np.random.uniform(-10, 10, (coff * head_dim, hidden_size))).to(data_type).npu()
|
||||
ape = torch.tensor(np.random.uniform(-10, 10, (cmp_ratio, coff * head_dim))).to(torch.float32).npu()
|
||||
norm_weight = torch.tensor(np.random.uniform(-10, 10, (head_dim))).to(data_type).npu()
|
||||
rope_sin = torch.tensor(np.random.uniform(-1, 1, rope_sin_shape)).to(data_type).npu()
|
||||
rope_cos = torch.tensor(np.random.uniform(-1, 1, rope_cos_shape)).to(data_type).npu()
|
||||
kv_state = kv_state.npu()
|
||||
score_state = score_state.npu()
|
||||
block_table = block_table.npu()
|
||||
start_pos = torch.tensor(start_pos).to(torch.int32).npu()
|
||||
if cu_seqlens is not None:
|
||||
cu_seqlens = torch.tensor(cu_seqlens).to(torch.int32).npu()
|
||||
if seqused is not None:
|
||||
seqused = torch.tensor(seqused).to(torch.int32).npu()
|
||||
|
||||
class CompressorNetwork(nn.Module):
|
||||
def __init__(self):
|
||||
super(CompressorNetwork, self).__init__()
|
||||
|
||||
def forward(self, x, wkv, wgate, kv_state, score_state, ape, norm_weight, rope_sin,
|
||||
rope_cos, rope_head_dim, cmp_ratio, kv_block_table = None, score_block_table = None, cu_seqlens = None,
|
||||
seqused = None, start_pos = None, coff = 1, norm_eps = 1e-6, rotary_mode = 1):
|
||||
cmp_kv,_ ,_ ,_ ,_ = (
|
||||
torch.ops.custom.compressor(
|
||||
x,
|
||||
wkv,
|
||||
wgate,
|
||||
kv_state,
|
||||
score_state,
|
||||
ape,
|
||||
norm_weight,
|
||||
rope_sin,
|
||||
rope_cos,
|
||||
kv_block_table = kv_block_table,
|
||||
score_block_table = score_block_table,
|
||||
cu_seqlens = cu_seqlens,
|
||||
seqused = seqused,
|
||||
start_pos = start_pos,
|
||||
rope_head_dim = rope_head_dim,
|
||||
cmp_ratio = cmp_ratio,
|
||||
coff = coff,
|
||||
norm_eps = norm_eps,
|
||||
rotary_mode = rotary_mode
|
||||
)
|
||||
)
|
||||
return cmp_kv
|
||||
|
||||
from torchair.configs.compiler_config import CompilerConfig
|
||||
config = CompilerConfig()
|
||||
npu_backend = torchair.get_npu_backend(compiler_config=config)
|
||||
torch._dynamo.reset()
|
||||
npu_mode = torch.compile(CompressorNetwork(), fullgraph=True, backend=npu_backend, dynamic=False)
|
||||
cmp_kv = npu_mode(
|
||||
x,
|
||||
wkv,
|
||||
wgate,
|
||||
kv_state,
|
||||
score_state,
|
||||
ape,
|
||||
norm_weight,
|
||||
rope_sin,
|
||||
rope_cos,
|
||||
kv_block_table = block_table,
|
||||
score_block_table = block_table,
|
||||
cu_seqlens = cu_seqlens,
|
||||
seqused = seqused,
|
||||
start_pos = start_pos,
|
||||
rope_head_dim = rope_head_dim,
|
||||
cmp_ratio = cmp_ratio,
|
||||
coff = coff,
|
||||
norm_eps = norm_eps,
|
||||
rotary_mode = rotary_mode)
|
||||
```
|
||||
|
||||
更多使用示例见[pytest示例](./tests/pytest/README.md)。
|
||||
40
csrc/attention/compressor/op_host/CMakeLists.txt
Normal file
40
csrc/attention/compressor/op_host/CMakeLists.txt
Normal file
@@ -0,0 +1,40 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
add_op_to_compiled_list()
|
||||
|
||||
if (BUILD_OPEN_PROJECT)
|
||||
target_sources(op_host_aclnn PRIVATE
|
||||
compressor_def.cpp
|
||||
)
|
||||
endif()
|
||||
|
||||
add_ops_compile_options(
|
||||
OP_NAME Compressor
|
||||
OPTIONS --cce-auto-sync=off
|
||||
-Wno-deprecated-declarations
|
||||
-mllvm -cce-aicore-hoist-movemask=false
|
||||
--op_relocatable_kernel_binary=true
|
||||
)
|
||||
|
||||
if (NOT BUILD_OPS_RTY_KERNEL)
|
||||
set(SUPPORTED_ARCHS arch32 arch35)
|
||||
add_modules_sources(OPTYPE compressor ACLNNTYPE aclnn)
|
||||
add_tiling_modules()
|
||||
|
||||
foreach(ARCH ${ARCH_DIRECTORY})
|
||||
if(ARCH IN_LIST SUPPORTED_ARCHS)
|
||||
target_sources(${OPHOST_NAME}_tiling_obj PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/${ARCH}/compressor_tiling.cpp
|
||||
)
|
||||
endif()
|
||||
endforeach()
|
||||
|
||||
endif()
|
||||
|
||||
1027
csrc/attention/compressor/op_host/arch32/compressor_tiling.cpp
Normal file
1027
csrc/attention/compressor/op_host/arch32/compressor_tiling.cpp
Normal file
File diff suppressed because it is too large
Load Diff
381
csrc/attention/compressor/op_host/arch32/compressor_tiling.h
Normal file
381
csrc/attention/compressor/op_host/arch32/compressor_tiling.h
Normal file
@@ -0,0 +1,381 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_tiling.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_TILING_H
|
||||
#define COMPRESSOR_TILING_H
|
||||
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <map>
|
||||
#include <unordered_map>
|
||||
#include <set>
|
||||
#include <sstream>
|
||||
#include "register/tilingdata_base.h"
|
||||
#include "tiling/tiling_api.h"
|
||||
#include "exe_graph/runtime/tiling_context.h"
|
||||
#include "register/op_def_registry.h"
|
||||
#include "../../op_kernel/arch32/compressor_template_tiling_key.h"
|
||||
#include "../../op_kernel/arch32/compressor_tiling_data.h"
|
||||
#include "platform/platform_info.h"
|
||||
|
||||
#ifdef ASCENDC_OP_TEST
|
||||
#define CMP_EXTERN_C extern "C"
|
||||
#else
|
||||
#define CMP_EXTERN_C
|
||||
#endif
|
||||
// #define DAY0_SCOPE
|
||||
|
||||
namespace optiling {
|
||||
|
||||
// INPUT
|
||||
constexpr uint32_t TOKEN_X_INPUT_INDEX = 0;
|
||||
constexpr uint32_t WEIGHT_KV_INPUT_INDEX = 1;
|
||||
constexpr uint32_t WEIGHT_WGATE_INPUT_INDEX = 2;
|
||||
constexpr uint32_t STATE_CACHE_INPUT_INDEX = 3;
|
||||
constexpr uint32_t APE_INPUT_INDEX = 4;
|
||||
constexpr uint32_t NORM_WEIGHT_INPUT_INDEX = 5;
|
||||
constexpr uint32_t ROPE_SIN_INPUT_INDEX = 6;
|
||||
constexpr uint32_t ROPE_COS_INPUT_INDEX = 7;
|
||||
|
||||
// INPUT(OPTION)
|
||||
constexpr uint32_t STATE_BLOCK_TABLE_INPUT_INDEX = 8;
|
||||
constexpr uint32_t CU_SEQ_LEN_INPUT_INDEX = 9;
|
||||
constexpr uint32_t SEQ_USED_INPUT_INDEX = 10;
|
||||
constexpr uint32_t START_POS_INPUT_INDEX = 11;
|
||||
|
||||
// ATTR
|
||||
constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 0;
|
||||
constexpr uint32_t CMP_RATIO_ATTR_INDEX = 1;
|
||||
constexpr uint32_t COFF_ATTR_INDEX = 2;
|
||||
constexpr uint32_t NORM_EPS_ATTR_INDEX = 3;
|
||||
constexpr uint32_t ROTARY_MODE_ATTR_INDEX = 4;
|
||||
constexpr uint32_t CACHE_MODE_ATTR_INDEX = 5;
|
||||
constexpr uint32_t STATE_CACHE_STRIDE_DIM0_ATTR_INDEX = 6;
|
||||
|
||||
// OUTPUT
|
||||
constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0;
|
||||
|
||||
constexpr uint32_t COMPRESSOR_DIM_NUM_1 = 1;
|
||||
constexpr uint32_t COMPRESSOR_DIM_NUM_2 = 2;
|
||||
constexpr uint32_t COMPRESSOR_DIM_NUM_3 = 3;
|
||||
constexpr uint32_t COMPRESSOR_DIM_NUM_4 = 4;
|
||||
constexpr uint32_t COMPRESSOR_DIM_INDEX_0 = 0;
|
||||
constexpr uint32_t COMPRESSOR_DIM_INDEX_1 = 1;
|
||||
constexpr uint32_t COMPRESSOR_DIM_INDEX_2 = 2;
|
||||
constexpr uint32_t COMPRESSOR_DIM_INDEX_3 = 3;
|
||||
|
||||
// CONSTRAINTS
|
||||
constexpr uint32_t MAX_HIDDEN_SIZE = 10240;
|
||||
constexpr uint32_t MIN_HIDDEN_SIZE = 1024;
|
||||
constexpr uint32_t ALIGN_FACTOR_HIDDEN_SIZE = 512;
|
||||
constexpr uint32_t MIN_BLOCK_SIZE = 1;
|
||||
|
||||
constexpr uint32_t BATCH_MODE_SCHEDULE = 1;
|
||||
|
||||
static const std::string X_NAME = "query";
|
||||
static const std::string WKV_NAME = "wkv";
|
||||
static const std::string WGATE_NAME = "wgate";
|
||||
static const std::string STATE_CACHE_NAME = "state_cache";
|
||||
static const std::string APE_NAME = "ape";
|
||||
static const std::string NORM_WEIGHT_NAME = "norm_weight";
|
||||
static const std::string ROPE_SIN_NAME = "rope_sin";
|
||||
static const std::string ROPE_COS_NAME = "rope_cos";
|
||||
static const std::string STATE_BLOCK_TABLE_NAME = "state_block_table";
|
||||
static const std::string CU_SEQLENS_NAME = "cu_seqlens";
|
||||
static const std::string SEQUSED_NAME = "seq_used";
|
||||
static const std::string START_POS_NAME = "start_pos";
|
||||
static const std::string ROPE_HEAD_DIM_NAME = "rope_head_dim";
|
||||
static const std::string CMP_RATIO_NAME = "cmp_ratio";
|
||||
static const std::string COFF_NAME = "coff";
|
||||
static const std::string NORM_EPS_NAME = "nrom_eps";
|
||||
static const std::string ROTARY_MODE_NAME = "rotary_mode";
|
||||
static const std::string CACHE_MODE_NAME = "cache_mode";
|
||||
static const std::string CMP_KV_NAME = "cmp_kv";
|
||||
|
||||
static std::string DataTypeToSerialString(ge::DataType type);
|
||||
|
||||
const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP = {
|
||||
{X_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
|
||||
{WKV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
|
||||
{WGATE_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
|
||||
{STATE_CACHE_NAME, {ge::DT_FLOAT}},
|
||||
{APE_NAME, {ge::DT_FLOAT}},
|
||||
{NORM_WEIGHT_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
|
||||
{ROPE_SIN_NAME, {ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT}},
|
||||
{ROPE_COS_NAME, {ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT}},
|
||||
{STATE_BLOCK_TABLE_NAME, {ge::DT_INT32}},
|
||||
{CU_SEQLENS_NAME, {ge::DT_INT32}},
|
||||
{SEQUSED_NAME, {ge::DT_INT32}},
|
||||
{START_POS_NAME, {ge::DT_INT32}},
|
||||
{CMP_KV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}
|
||||
};
|
||||
|
||||
const std::map<std::string, std::vector<uint32_t>> DIM_NUM_MAP = {
|
||||
{X_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},
|
||||
{WKV_NAME, {COMPRESSOR_DIM_NUM_2}},
|
||||
{WGATE_NAME, {COMPRESSOR_DIM_NUM_2}},
|
||||
{STATE_CACHE_NAME, {COMPRESSOR_DIM_NUM_3}},
|
||||
{APE_NAME, {COMPRESSOR_DIM_NUM_2}},
|
||||
{NORM_WEIGHT_NAME, {COMPRESSOR_DIM_NUM_1}},
|
||||
{ROPE_SIN_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},
|
||||
{ROPE_COS_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},
|
||||
{STATE_BLOCK_TABLE_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_1}},
|
||||
{CU_SEQLENS_NAME, {COMPRESSOR_DIM_NUM_1}},
|
||||
{SEQUSED_NAME, {COMPRESSOR_DIM_NUM_1}},
|
||||
{START_POS_NAME, {COMPRESSOR_DIM_NUM_1}},
|
||||
{CMP_KV_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}
|
||||
};
|
||||
|
||||
static const std::map<std::string, uint32_t> LAYOUT_DIM_MAP = {
|
||||
{"BSH", COMPRESSOR_DIM_NUM_3},
|
||||
{"TH", COMPRESSOR_DIM_NUM_2},
|
||||
};
|
||||
|
||||
const std::map<ge::DataType, std::string> DATATYPE_TO_STRING_MAP = {
|
||||
{ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set.
|
||||
{ge::DT_FLOAT, "DT_FLOAT"}, // float type
|
||||
{ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type
|
||||
{ge::DT_INT8, "DT_INT8"}, // int8 type
|
||||
{ge::DT_INT16, "DT_INT16"}, // int16 type
|
||||
{ge::DT_UINT16, "DT_UINT16"}, // uint16 type
|
||||
{ge::DT_UINT8, "DT_UINT8"}, // uint8 type
|
||||
{ge::DT_INT32, "DT_INT32"}, // uint32 type
|
||||
{ge::DT_INT64, "DT_INT64"}, // int64 type
|
||||
{ge::DT_UINT32, "DT_UINT32"}, // unsigned int32
|
||||
{ge::DT_UINT64, "DT_UINT64"}, // unsigned int64
|
||||
{ge::DT_BOOL, "DT_BOOL"}, // bool type
|
||||
{ge::DT_DOUBLE, "DT_DOUBLE"}, // double type
|
||||
{ge::DT_DUAL, "DT_DUAL"}, // dual output type
|
||||
{ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type
|
||||
{ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type
|
||||
{ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type
|
||||
{ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type
|
||||
{ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type
|
||||
{ge::DT_QINT8, "DT_QINT8"}, // qint8 type
|
||||
{ge::DT_QINT16, "DT_QINT16"}, // qint16 type
|
||||
{ge::DT_QINT32, "DT_QINT32"}, // qint32 type
|
||||
{ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type
|
||||
{ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type
|
||||
{ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type
|
||||
{ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type
|
||||
{ge::DT_STRING, "DT_STRING"}, // string type
|
||||
{ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type
|
||||
{ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type
|
||||
{ge::DT_INT4, "DT_INT4"}, // dt_variant type
|
||||
{ge::DT_UINT1, "DT_UINT1"}, // dt_variant type
|
||||
{ge::DT_INT2, "DT_INT2"}, // dt_variant type
|
||||
{ge::DT_UINT2, "DT_UINT2"} // dt_variant type
|
||||
};
|
||||
|
||||
struct CompressorCompileInfo {
|
||||
int64_t core_num;
|
||||
};
|
||||
|
||||
struct RequiredParaInfo {
|
||||
const gert::CompileTimeTensorDesc *desc;
|
||||
const gert::StorageShape *shape;
|
||||
};
|
||||
|
||||
struct OptionalParaInfo {
|
||||
const gert::CompileTimeTensorDesc *desc;
|
||||
const gert::StorageShape *shape;
|
||||
const gert::Tensor *tensor;
|
||||
};
|
||||
|
||||
enum class LayoutType {
|
||||
LAYOUT_BSH,
|
||||
LAYOUT_TH
|
||||
};
|
||||
|
||||
enum class TemplateId:uint8_t {
|
||||
NORMAL = 0,
|
||||
EMPTY_X = 1,
|
||||
PERF = 2
|
||||
};
|
||||
|
||||
CMP_EXTERN_C ge::graphStatus TilingCompressor(gert::TilingContext *context);
|
||||
struct CompressorBaseShapeInfo {
|
||||
uint32_t bSize = 0; // B
|
||||
uint32_t sSize = 0; // S
|
||||
uint32_t hSize = 0; // Hidden size
|
||||
uint32_t tSize = 0; // T
|
||||
uint32_t nSize = 0; // N
|
||||
uint32_t dSize = 0; // D
|
||||
uint32_t coffSize = 0; // Coff: 1 or 2
|
||||
uint32_t csSize = 0; // Compress sequence len
|
||||
uint32_t rSize = 0; // Compress ratio
|
||||
uint32_t cgSize = 0; // Compress group size
|
||||
uint32_t drSize = 0; // Dr
|
||||
};
|
||||
|
||||
const std::vector<int> ROPE_HEAD_DIM {64};
|
||||
const std::vector<int> COFF {1, 2};
|
||||
#ifdef DAY0_SCOPE
|
||||
const std::vector<int> CMP_RATIO {4, 128};
|
||||
const std::vector<int> ROTARY_MODE {2};
|
||||
#else
|
||||
const std::vector<int> CMP_RATIO {2, 4, 8, 16, 32, 64, 128};
|
||||
const std::vector<int> ROTARY_MODE {1, 2};
|
||||
#endif
|
||||
const std::vector<uint32_t> HEAD_DIM {128, 512};
|
||||
const std::vector<int> CACHE_MODE {1};
|
||||
|
||||
enum class ROTARY_MODE:uint8_t {
|
||||
HALF = 1,
|
||||
INTERLEAVE = 2
|
||||
};
|
||||
|
||||
enum class CACHE_MODE:uint8_t {
|
||||
CONTINUOUS = 1,
|
||||
CYCLE = 2
|
||||
};
|
||||
|
||||
struct CompressorContext {
|
||||
const char *opName;
|
||||
const char *opType;
|
||||
fe::PlatFormInfos *platformInfo;
|
||||
|
||||
RequiredParaInfo x;
|
||||
RequiredParaInfo wkv;
|
||||
RequiredParaInfo wgate;
|
||||
RequiredParaInfo stateCache;
|
||||
RequiredParaInfo ape;
|
||||
RequiredParaInfo normWeight;
|
||||
RequiredParaInfo ropeSin;
|
||||
RequiredParaInfo ropeCos;
|
||||
OptionalParaInfo stateBlockTable;
|
||||
OptionalParaInfo cuSeqlens;
|
||||
OptionalParaInfo seqUsed;
|
||||
OptionalParaInfo startPos;
|
||||
RequiredParaInfo cmpKv;
|
||||
|
||||
const int *ropeHeadDim;
|
||||
const int *coff;
|
||||
const int *cmpRatio;
|
||||
const float *normEps;
|
||||
const int *rotaryMode;
|
||||
const int *cacheMode;
|
||||
const int *stateCacheStrideDim0;
|
||||
TemplateId templateId;
|
||||
|
||||
ge::DataType dtype = ge::DT_BF16;
|
||||
LayoutType layout = LayoutType::LAYOUT_BSH;
|
||||
|
||||
size_t *workSpaces;
|
||||
uint64_t tilingKey;
|
||||
uint32_t blockDim;
|
||||
};
|
||||
|
||||
class CompressorTiling {
|
||||
public:
|
||||
explicit CompressorTiling(CompressorContext *context) : context_(context) {}
|
||||
~CompressorTiling() = default;
|
||||
|
||||
static ge::graphStatus ConvertContext(gert::TilingContext &context, CompressorContext &compressorContext);
|
||||
ge::graphStatus RunBigKernelTiling(CompressorTilingData* tilingData);
|
||||
|
||||
private:
|
||||
static void ConvertRequiredParams(gert::TilingContext &context, CompressorContext &compressorContext);
|
||||
|
||||
static void ConvertOptionalParams(gert::TilingContext &context, CompressorContext &compressorContext);
|
||||
ge::graphStatus GetNpuInfo();
|
||||
ge::graphStatus SetBaseInfo();
|
||||
ge::graphStatus SetPageAttentionInfo();
|
||||
ge::graphStatus SetWorkSpaceInfo();
|
||||
ge::graphStatus SetScenarioInfo();
|
||||
ge::graphStatus SetTemplateId();
|
||||
ge::graphStatus SetInnerSplitInfo();
|
||||
ge::graphStatus CalcWorkSpace();
|
||||
ge::graphStatus CheckSinglePara() const;
|
||||
ge::graphStatus GenTilingKey() const;
|
||||
template <typename T>
|
||||
ge::graphStatus CheckFeatureValueSupport(const T *featureValue, const std::vector<T> &expectFeatureValList,
|
||||
const std::string &name) const;
|
||||
template <typename T>
|
||||
ge::graphStatus CheckAttrValueSupport(const T *attrValue, const std::vector<T> &expectAttrValList,
|
||||
const std::string &name) const;
|
||||
template <typename T>
|
||||
void LogErrorNumberSupport(const std::vector<T> &expectNumberList, const T &actualValue, const std::string &name,
|
||||
const std::string subName) const;
|
||||
ge::graphStatus CheckDimNumInLayoutSupport(const std::string &layout, const gert::StorageShape *shape,
|
||||
const std::string &name) const;
|
||||
ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, const std::string &name) const;
|
||||
void LogErrorDtypeSupport(const std::vector<ge::DataType> &expectDtypeList, const ge::DataType &actualDtype,
|
||||
const std::string &name) const;
|
||||
ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape, const std::string &name) const;
|
||||
ge::graphStatus LogErrorShapeConsistency(const std::string &name, const gert::StorageShape *shape,
|
||||
const uint32_t &dimNum, const std::string &subName,
|
||||
const uint32_t &expectNum) const;
|
||||
ge::graphStatus CheckSingleParaX() const;
|
||||
ge::graphStatus CheckSingleParaWkv() const;
|
||||
ge::graphStatus CheckSingleParaWgate() const;
|
||||
ge::graphStatus CheckSingleParaStateCache() const;
|
||||
ge::graphStatus CheckSingleParaApe() const;
|
||||
ge::graphStatus CheckSingleParaNormWeight() const;
|
||||
ge::graphStatus CheckSingleParaRopeSin() const;
|
||||
ge::graphStatus CheckSingleParaRopeCos() const;
|
||||
ge::graphStatus CheckSingleParaStateBlockTable() const;
|
||||
ge::graphStatus CheckSingleParaCuSeqlens() const;
|
||||
ge::graphStatus CheckSingleParaSeqused() const;
|
||||
ge::graphStatus CheckSingleParaStartPos() const;
|
||||
ge::graphStatus CheckSingleParaCmpKv() const;
|
||||
ge::graphStatus CheckSingleParaRopeHeadDim() const;
|
||||
ge::graphStatus CheckSingleParaCmpRatio() const;
|
||||
ge::graphStatus CheckSingleParaCoff() const;
|
||||
ge::graphStatus CheckSingleParaNormEps() const;
|
||||
ge::graphStatus CheckSingleParaRotaryMode() const;
|
||||
ge::graphStatus CheckSingleParaCacheMode() const;
|
||||
ge::graphStatus CheckRequiredParaExistence() const;
|
||||
ge::graphStatus CheckRequiredInOutExistence() const;
|
||||
ge::graphStatus CheckRequiredAttrExistence() const;
|
||||
ge::graphStatus CheckFeature() const;
|
||||
ge::graphStatus CheckShapeConsistency() const;
|
||||
ge::graphStatus CheckShapeConsistencyRope() const;
|
||||
ge::graphStatus CheckDtypeConsistencyX(const gert::CompileTimeTensorDesc *desc, const std::string &name) const;
|
||||
ge::graphStatus CheckDtypeConsistencyRope() const;
|
||||
ge::graphStatus CheckDtypeConsistency() const;
|
||||
ge::graphStatus CheckMultiParaConsistency() const;
|
||||
ge::graphStatus CheckDimNumConsistency() const;
|
||||
ge::graphStatus CheckEmptyTensor() const;
|
||||
ge::graphStatus CheckScenarioConsistency() const;
|
||||
ge::graphStatus CheckBlockDimConstrain() const;
|
||||
|
||||
size_t ubSize_ = 0;
|
||||
size_t l1Size_ = 0;
|
||||
size_t l0cSize_ = 0;
|
||||
size_t l0bSize_ = 0;
|
||||
uint32_t coreNum_ = 0;
|
||||
uint32_t aicNum_ = 0;
|
||||
uint32_t aivNum_ = 0;
|
||||
platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B;
|
||||
size_t libapiSize_ = 0;
|
||||
size_t workspaceSize_ = 0;
|
||||
uint8_t coff = 1;
|
||||
|
||||
uint32_t mBaseSize = 0;
|
||||
uint32_t dbaseSize = 0;
|
||||
|
||||
CompressorBaseShapeInfo baseShapeInfo_;
|
||||
CompressorContext *context_ = nullptr;
|
||||
CompressorBaseParams *baseParams_ = nullptr;
|
||||
CompressorPageAttentionParams *pageAttentionParams_ = nullptr;
|
||||
CompressorInnerSplitParams *innerSplitParams_ = nullptr;
|
||||
CompressorWorkspaceParams *workspaceParams_ = nullptr;
|
||||
};
|
||||
|
||||
} // optiling
|
||||
|
||||
#endif
|
||||
1071
csrc/attention/compressor/op_host/arch35/compressor_tiling.cpp
Normal file
1071
csrc/attention/compressor/op_host/arch35/compressor_tiling.cpp
Normal file
File diff suppressed because it is too large
Load Diff
375
csrc/attention/compressor/op_host/arch35/compressor_tiling.h
Normal file
375
csrc/attention/compressor/op_host/arch35/compressor_tiling.h
Normal file
@@ -0,0 +1,375 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_tiling.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_TILING_H
|
||||
#define COMPRESSOR_TILING_H
|
||||
|
||||
#include <cstdint>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <map>
|
||||
#include <unordered_map>
|
||||
#include <set>
|
||||
#include <sstream>
|
||||
#include "register/tilingdata_base.h"
|
||||
#include "tiling/tiling_api.h"
|
||||
#include "exe_graph/runtime/tiling_context.h"
|
||||
#include "register/op_def_registry.h"
|
||||
#include "../../op_kernel/arch35/compressor_template_tiling_key.h"
|
||||
#include "../../op_kernel/arch35/compressor_tiling_data.h"
|
||||
#include "platform/platform_info.h"
|
||||
|
||||
#ifdef ASCENDC_OP_TEST
|
||||
#define CMP_EXTERN_C extern "C"
|
||||
#else
|
||||
#define CMP_EXTERN_C
|
||||
#endif
|
||||
|
||||
namespace optiling {
|
||||
|
||||
// INPUT
|
||||
constexpr uint32_t TOKEN_X_INPUT_INDEX = 0;
|
||||
constexpr uint32_t WEIGHT_KV_INPUT_INDEX = 1;
|
||||
constexpr uint32_t WEIGHT_WGATE_INPUT_INDEX = 2;
|
||||
constexpr uint32_t STATE_CACHE_INPUT_INDEX = 3;
|
||||
constexpr uint32_t APE_INPUT_INDEX = 4;
|
||||
constexpr uint32_t NORM_WEIGHT_INPUT_INDEX = 5;
|
||||
constexpr uint32_t ROPE_SIN_INPUT_INDEX = 6;
|
||||
constexpr uint32_t ROPE_COS_INPUT_INDEX = 7;
|
||||
|
||||
// INPUT(OPTION)
|
||||
constexpr uint32_t STATE_BLOCK_TABLE_INPUT_INDEX = 8;
|
||||
constexpr uint32_t CU_SEQ_LEN_INPUT_INDEX = 9;
|
||||
constexpr uint32_t SEQ_USED_INPUT_INDEX = 10;
|
||||
constexpr uint32_t START_POS_INPUT_INDEX = 11;
|
||||
|
||||
// ATTR
|
||||
constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 0;
|
||||
constexpr uint32_t CMP_RATIO_ATTR_INDEX = 1;
|
||||
constexpr uint32_t COFF_ATTR_INDEX = 2;
|
||||
constexpr uint32_t NORM_EPS_ATTR_INDEX = 3;
|
||||
constexpr uint32_t ROTARY_MODE_ATTR_INDEX = 4;
|
||||
constexpr uint32_t CACHE_MODE_ATTR_INDEX = 5;
|
||||
constexpr uint32_t STATE_CACHE_STRIDE_DIM0_ATTR_INDEX = 6;
|
||||
|
||||
// OUTPUT
|
||||
constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0;
|
||||
|
||||
constexpr uint32_t COMPRESSOR_DIM_NUM_1 = 1;
|
||||
constexpr uint32_t COMPRESSOR_DIM_NUM_2 = 2;
|
||||
constexpr uint32_t COMPRESSOR_DIM_NUM_3 = 3;
|
||||
constexpr uint32_t COMPRESSOR_DIM_NUM_4 = 4;
|
||||
constexpr uint32_t COMPRESSOR_DIM_INDEX_0 = 0;
|
||||
constexpr uint32_t COMPRESSOR_DIM_INDEX_1 = 1;
|
||||
constexpr uint32_t COMPRESSOR_DIM_INDEX_2 = 2;
|
||||
constexpr uint32_t COMPRESSOR_DIM_INDEX_3 = 3;
|
||||
|
||||
// CONSTRAINTS
|
||||
constexpr uint32_t MAX_HIDDEN_SIZE = 10240;
|
||||
constexpr uint32_t MIN_HIDDEN_SIZE = 1024;
|
||||
constexpr uint32_t ALIGN_FACTOR_HIDDEN_SIZE = 512;
|
||||
constexpr uint32_t MIN_BLOCK_SIZE = 1;
|
||||
|
||||
constexpr uint32_t BATCH_MODE_SCHEDULE = 1;
|
||||
|
||||
static const std::string X_NAME = "query";
|
||||
static const std::string WKV_NAME = "wkv";
|
||||
static const std::string WGATE_NAME = "wgate";
|
||||
static const std::string STATE_CACHE_NAME = "state_cache";
|
||||
static const std::string APE_NAME = "ape";
|
||||
static const std::string NORM_WEIGHT_NAME = "norm_weight";
|
||||
static const std::string ROPE_SIN_NAME = "rope_sin";
|
||||
static const std::string ROPE_COS_NAME = "rope_cos";
|
||||
static const std::string STATE_BLOCK_TABLE_NAME = "state_block_table";
|
||||
static const std::string CU_SEQLENS_NAME = "cu_seqlens";
|
||||
static const std::string SEQUSED_NAME = "seq_used";
|
||||
static const std::string START_POS_NAME = "start_pos";
|
||||
static const std::string ROPE_HEAD_DIM_NAME = "rope_head_dim";
|
||||
static const std::string CMP_RATIO_NAME = "cmp_ratio";
|
||||
static const std::string COFF_NAME = "coff";
|
||||
static const std::string NORM_EPS_NAME = "nrom_eps";
|
||||
static const std::string ROTARY_MODE_NAME = "rotary_mode";
|
||||
static const std::string CACHE_MODE_NAME = "cache_mode";
|
||||
static const std::string CMP_KV_NAME = "cmp_kv";
|
||||
|
||||
static std::string DataTypeToSerialString(ge::DataType type);
|
||||
|
||||
const std::map<std::string, std::vector<ge::DataType>> DTYPE_SUPPORT_MAP = {
|
||||
{X_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
|
||||
{WKV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
|
||||
{WGATE_NAME, {ge::DT_BF16, ge::DT_FLOAT16}},
|
||||
{STATE_CACHE_NAME, {ge::DT_FLOAT}},
|
||||
{APE_NAME, {ge::DT_FLOAT}},
|
||||
{NORM_WEIGHT_NAME, {ge::DT_FLOAT}},
|
||||
{ROPE_SIN_NAME, {ge::DT_FLOAT}},
|
||||
{ROPE_COS_NAME, {ge::DT_FLOAT}},
|
||||
{STATE_BLOCK_TABLE_NAME, {ge::DT_INT32}},
|
||||
{CU_SEQLENS_NAME, {ge::DT_INT32}},
|
||||
{SEQUSED_NAME, {ge::DT_INT32}},
|
||||
{START_POS_NAME, {ge::DT_INT32}},
|
||||
{CMP_KV_NAME, {ge::DT_BF16, ge::DT_FLOAT16}}
|
||||
};
|
||||
|
||||
const std::map<std::string, std::vector<uint32_t>> DIM_NUM_MAP = {
|
||||
{X_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},
|
||||
{WKV_NAME, {COMPRESSOR_DIM_NUM_2}},
|
||||
{WGATE_NAME, {COMPRESSOR_DIM_NUM_2}},
|
||||
{STATE_CACHE_NAME, {COMPRESSOR_DIM_NUM_3}},
|
||||
{APE_NAME, {COMPRESSOR_DIM_NUM_2}},
|
||||
{NORM_WEIGHT_NAME, {COMPRESSOR_DIM_NUM_1}},
|
||||
{ROPE_SIN_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},
|
||||
{ROPE_COS_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}},
|
||||
{STATE_BLOCK_TABLE_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_1}},
|
||||
{CU_SEQLENS_NAME, {COMPRESSOR_DIM_NUM_1}},
|
||||
{SEQUSED_NAME, {COMPRESSOR_DIM_NUM_1}},
|
||||
{START_POS_NAME, {COMPRESSOR_DIM_NUM_1}},
|
||||
{CMP_KV_NAME, {COMPRESSOR_DIM_NUM_2, COMPRESSOR_DIM_NUM_3}}
|
||||
};
|
||||
|
||||
static const std::map<std::string, uint32_t> LAYOUT_DIM_MAP = {
|
||||
{"BSH", COMPRESSOR_DIM_NUM_3},
|
||||
{"TH", COMPRESSOR_DIM_NUM_2},
|
||||
};
|
||||
|
||||
const std::map<ge::DataType, std::string> DATATYPE_TO_STRING_MAP = {
|
||||
{ge::DT_UNDEFINED, "DT_UNDEFINED"}, // Used to indicate a DataType field has not been set.
|
||||
{ge::DT_FLOAT, "DT_FLOAT"}, // float type
|
||||
{ge::DT_FLOAT16, "DT_FLOAT16"}, // fp16 type
|
||||
{ge::DT_INT8, "DT_INT8"}, // int8 type
|
||||
{ge::DT_INT16, "DT_INT16"}, // int16 type
|
||||
{ge::DT_UINT16, "DT_UINT16"}, // uint16 type
|
||||
{ge::DT_UINT8, "DT_UINT8"}, // uint8 type
|
||||
{ge::DT_INT32, "DT_INT32"}, // uint32 type
|
||||
{ge::DT_INT64, "DT_INT64"}, // int64 type
|
||||
{ge::DT_UINT32, "DT_UINT32"}, // unsigned int32
|
||||
{ge::DT_UINT64, "DT_UINT64"}, // unsigned int64
|
||||
{ge::DT_BOOL, "DT_BOOL"}, // bool type
|
||||
{ge::DT_DOUBLE, "DT_DOUBLE"}, // double type
|
||||
{ge::DT_DUAL, "DT_DUAL"}, // dual output type
|
||||
{ge::DT_DUAL_SUB_INT8, "DT_DUAL_SUB_INT8"}, // dual output int8 type
|
||||
{ge::DT_DUAL_SUB_UINT8, "DT_DUAL_SUB_UINT8"}, // dual output uint8 type
|
||||
{ge::DT_COMPLEX32, "DT_COMPLEX32"}, // complex32 type
|
||||
{ge::DT_COMPLEX64, "DT_COMPLEX64"}, // complex64 type
|
||||
{ge::DT_COMPLEX128, "DT_COMPLEX128"}, // complex128 type
|
||||
{ge::DT_QINT8, "DT_QINT8"}, // qint8 type
|
||||
{ge::DT_QINT16, "DT_QINT16"}, // qint16 type
|
||||
{ge::DT_QINT32, "DT_QINT32"}, // qint32 type
|
||||
{ge::DT_QUINT8, "DT_QUINT8"}, // quint8 type
|
||||
{ge::DT_QUINT16, "DT_QUINT16"}, // quint16 type
|
||||
{ge::DT_RESOURCE, "DT_RESOURCE"}, // resource type
|
||||
{ge::DT_STRING_REF, "DT_STRING_REF"}, // string ref type
|
||||
{ge::DT_STRING, "DT_STRING"}, // string type
|
||||
{ge::DT_VARIANT, "DT_VARIANT"}, // dt_variant type
|
||||
{ge::DT_BF16, "DT_BFLOAT16"}, // dt_bfloat16 type
|
||||
{ge::DT_INT4, "DT_INT4"}, // dt_variant type
|
||||
{ge::DT_UINT1, "DT_UINT1"}, // dt_variant type
|
||||
{ge::DT_INT2, "DT_INT2"}, // dt_variant type
|
||||
{ge::DT_UINT2, "DT_UINT2"} // dt_variant type
|
||||
};
|
||||
|
||||
struct CompressorCompileInfo {
|
||||
int64_t core_num;
|
||||
};
|
||||
|
||||
struct RequiredParaInfo {
|
||||
const gert::CompileTimeTensorDesc *desc;
|
||||
const gert::StorageShape *shape;
|
||||
};
|
||||
|
||||
struct OptionalParaInfo {
|
||||
const gert::CompileTimeTensorDesc *desc;
|
||||
const gert::StorageShape *shape;
|
||||
const gert::Tensor *tensor;
|
||||
};
|
||||
|
||||
enum class LayoutType {
|
||||
LAYOUT_BSH,
|
||||
LAYOUT_TH
|
||||
};
|
||||
|
||||
enum class TemplateId:uint8_t {
|
||||
NORMAL = 0,
|
||||
EMPTY_X = 1,
|
||||
FULL_LOAD = 2
|
||||
};
|
||||
|
||||
CMP_EXTERN_C ge::graphStatus TilingCompressor(gert::TilingContext *context);
|
||||
struct CompressorBaseShapeInfo {
|
||||
uint32_t bSize = 0; // B
|
||||
uint32_t sSize = 0; // S
|
||||
uint32_t hSize = 0; // Hidden size
|
||||
uint32_t tSize = 0; // T
|
||||
uint32_t nSize = 0; // N
|
||||
uint32_t dSize = 0; // D
|
||||
uint32_t coffSize = 0; // Coff: 1 or 2
|
||||
uint32_t csSize = 0; // Compress sequence len
|
||||
uint32_t rSize = 0; // Compress ratio
|
||||
uint32_t cgSize = 0; // Compress group size
|
||||
uint32_t drSize = 0; // Dr
|
||||
};
|
||||
|
||||
const std::vector<int> ROPE_HEAD_DIM {64};
|
||||
const std::vector<int> COFF {1, 2};
|
||||
const std::vector<int> CMP_RATIO {2, 4, 8, 16, 32, 64, 128};
|
||||
const std::vector<int> ROTARY_MODE {1, 2};
|
||||
const std::vector<uint32_t> HEAD_DIM {128, 512};
|
||||
const std::vector<int> CACHE_MODE {1, 2};
|
||||
|
||||
enum class ROTARY_MODE:uint8_t {
|
||||
HALF = 1,
|
||||
INTERLEAVE = 2
|
||||
};
|
||||
|
||||
enum class CACHE_MODE:uint8_t {
|
||||
CONTINUOUS = 1,
|
||||
CYCLE = 2
|
||||
};
|
||||
|
||||
struct CompressorContext {
|
||||
const char *opName;
|
||||
const char *opType;
|
||||
fe::PlatFormInfos *platformInfo;
|
||||
|
||||
RequiredParaInfo x;
|
||||
RequiredParaInfo wkv;
|
||||
RequiredParaInfo wgate;
|
||||
RequiredParaInfo stateCache;
|
||||
RequiredParaInfo ape;
|
||||
RequiredParaInfo normWeight;
|
||||
RequiredParaInfo ropeSin;
|
||||
RequiredParaInfo ropeCos;
|
||||
OptionalParaInfo stateBlockTable;
|
||||
OptionalParaInfo cuSeqlens;
|
||||
OptionalParaInfo seqUsed;
|
||||
OptionalParaInfo startPos;
|
||||
RequiredParaInfo cmpKv;
|
||||
|
||||
const int *ropeHeadDim;
|
||||
const int *coff;
|
||||
const int *cmpRatio;
|
||||
const float *normEps;
|
||||
const int *rotaryMode;
|
||||
const int *cacheMode;
|
||||
const int *stateCacheStrideDim0;
|
||||
TemplateId templateId;
|
||||
|
||||
ge::DataType dtype = ge::DT_BF16;
|
||||
LayoutType layout = LayoutType::LAYOUT_BSH;
|
||||
|
||||
size_t *workSpaces;
|
||||
uint64_t tilingKey;
|
||||
uint32_t blockDim;
|
||||
};
|
||||
|
||||
class CompressorTiling {
|
||||
public:
|
||||
explicit CompressorTiling(CompressorContext *context) : context_(context) {}
|
||||
~CompressorTiling() = default;
|
||||
|
||||
static ge::graphStatus ConvertContext(gert::TilingContext &context, CompressorContext &compressorContext);
|
||||
ge::graphStatus RunBigKernelTiling(CompressorTilingData* tilingData);
|
||||
|
||||
private:
|
||||
static void ConvertRequiredParams(gert::TilingContext &context, CompressorContext &compressorContext);
|
||||
|
||||
static void ConvertOptionalParams(gert::TilingContext &context, CompressorContext &compressorContext);
|
||||
ge::graphStatus GetNpuInfo();
|
||||
ge::graphStatus SetBaseInfo();
|
||||
ge::graphStatus SetPageAttentionInfo();
|
||||
ge::graphStatus SetWorkSpaceInfo();
|
||||
ge::graphStatus SetScenarioInfo();
|
||||
ge::graphStatus SetTemplateId();
|
||||
ge::graphStatus SetInnerSplitInfo();
|
||||
ge::graphStatus CalcWorkSpace();
|
||||
ge::graphStatus CheckSinglePara() const;
|
||||
ge::graphStatus GenTilingKey() const;
|
||||
template <typename T>
|
||||
ge::graphStatus CheckFeatureValueSupport(const T *featureValue, const std::vector<T> &expectFeatureValList,
|
||||
const std::string &name) const;
|
||||
template <typename T>
|
||||
ge::graphStatus CheckAttrValueSupport(const T *attrValue, const std::vector<T> &expectAttrValList,
|
||||
const std::string &name) const;
|
||||
template <typename T>
|
||||
void LogErrorNumberSupport(const std::vector<T> &expectNumberList, const T &actualValue, const std::string &name,
|
||||
const std::string subName) const;
|
||||
ge::graphStatus CheckDimNumInLayoutSupport(const std::string &layout, const gert::StorageShape *shape,
|
||||
const std::string &name) const;
|
||||
ge::graphStatus CheckDtypeSupport(const gert::CompileTimeTensorDesc *desc, const std::string &name) const;
|
||||
void LogErrorDtypeSupport(const std::vector<ge::DataType> &expectDtypeList, const ge::DataType &actualDtype,
|
||||
const std::string &name) const;
|
||||
ge::graphStatus CheckDimNumSupport(const gert::StorageShape *shape, const std::string &name) const;
|
||||
ge::graphStatus LogErrorShapeConsistency(const std::string &name, const gert::StorageShape *shape,
|
||||
const uint32_t &dimNum, const std::string &subName,
|
||||
const uint32_t &expectNum) const;
|
||||
ge::graphStatus CheckSingleParaX() const;
|
||||
ge::graphStatus CheckSingleParaWkv() const;
|
||||
ge::graphStatus CheckSingleParaWgate() const;
|
||||
ge::graphStatus CheckSingleParaStateCache() const;
|
||||
ge::graphStatus CheckSingleParaApe() const;
|
||||
ge::graphStatus CheckSingleParaNormWeight() const;
|
||||
ge::graphStatus CheckSingleParaRopeSin() const;
|
||||
ge::graphStatus CheckSingleParaRopeCos() const;
|
||||
ge::graphStatus CheckSingleParaStateBlockTable() const;
|
||||
ge::graphStatus CheckSingleParaCuSeqlens() const;
|
||||
ge::graphStatus CheckSingleParaSeqused() const;
|
||||
ge::graphStatus CheckSingleParaStartPos() const;
|
||||
ge::graphStatus CheckSingleParaCmpKv() const;
|
||||
ge::graphStatus CheckSingleParaRopeHeadDim() const;
|
||||
ge::graphStatus CheckSingleParaCmpRatio() const;
|
||||
ge::graphStatus CheckSingleParaCoff() const;
|
||||
ge::graphStatus CheckSingleParaNormEps() const;
|
||||
ge::graphStatus CheckSingleParaRotaryMode() const;
|
||||
ge::graphStatus CheckSingleParaCacheMode() const;
|
||||
ge::graphStatus CheckRequiredParaExistence() const;
|
||||
ge::graphStatus CheckRequiredInOutExistence() const;
|
||||
ge::graphStatus CheckRequiredAttrExistence() const;
|
||||
ge::graphStatus CheckFeature() const;
|
||||
ge::graphStatus CheckShapeConsistency() const;
|
||||
ge::graphStatus CheckShapeConsistencyRope() const;
|
||||
ge::graphStatus CheckDtypeConsistencyX(const gert::CompileTimeTensorDesc *desc, const std::string &name) const;
|
||||
ge::graphStatus CheckDtypeConsistencyFp32(const gert::CompileTimeTensorDesc *desc, const std::string &name) const;
|
||||
ge::graphStatus CheckDtypeConsistency() const;
|
||||
ge::graphStatus CheckMultiParaConsistency() const;
|
||||
ge::graphStatus CheckDimNumConsistency() const;
|
||||
ge::graphStatus CheckEmptyTensor() const;
|
||||
ge::graphStatus CheckScenarioConsistency() const;
|
||||
ge::graphStatus CheckBlockDimConstrain() const;
|
||||
|
||||
size_t ubSize_ = 0;
|
||||
size_t l1Size_ = 0;
|
||||
size_t l0cSize_ = 0;
|
||||
size_t l0bSize_ = 0;
|
||||
uint32_t coreNum_ = 0;
|
||||
uint32_t aicNum_ = 0;
|
||||
uint32_t aivNum_ = 0;
|
||||
platform_ascendc::SocVersion socVersion_ = platform_ascendc::SocVersion::ASCEND910B;
|
||||
size_t libapiSize_ = 0;
|
||||
size_t workspaceSize_ = 0;
|
||||
uint8_t coff = 1;
|
||||
|
||||
uint32_t mBaseSize = 0;
|
||||
uint32_t dbaseSize = 0;
|
||||
|
||||
CompressorBaseShapeInfo baseShapeInfo_;
|
||||
CompressorContext *context_ = nullptr;
|
||||
CompressorBaseParams *baseParams_ = nullptr;
|
||||
CompressorPageAttentionParams *pageAttentionParams_ = nullptr;
|
||||
CompressorInnerSplitParams *innerSplitParams_ = nullptr;
|
||||
CompressorWorkspaceParams *workspaceParams_ = nullptr;
|
||||
};
|
||||
|
||||
} // optiling
|
||||
|
||||
#endif
|
||||
191
csrc/attention/compressor/op_host/compressor_def.cpp
Normal file
191
csrc/attention/compressor/op_host/compressor_def.cpp
Normal file
@@ -0,0 +1,191 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
#include "register/op_def_registry.h"
|
||||
|
||||
namespace ops {
|
||||
class Compressor : public OpDef {
|
||||
public:
|
||||
static constexpr uint32_t ROPE_HEAD_DIM_VALUE = 64;
|
||||
static constexpr uint32_t CMP_RATIO_VALUE = 4;
|
||||
static constexpr uint32_t COFF_VALUE = 1;
|
||||
static constexpr uint32_t ROTARY_MODE_VALUE = 1;
|
||||
static constexpr uint32_t CACHE_MODE_VALUE = 1;
|
||||
static constexpr uint32_t STATE_CACHE_STRIDE_DIM0 = 0;
|
||||
|
||||
explicit Compressor(const char *name) : OpDef(name)
|
||||
{
|
||||
this->Input("x")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("wkv")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("wgate")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("state_cache")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.IgnoreContiguous();
|
||||
this->Input("ape")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("norm_weight")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("rope_sin")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("rope_cos")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("state_block_table")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("cu_seqlens")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("seqused")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("start_pos")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Output("cmp_kv")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND});
|
||||
this->Output("state_cache")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND});
|
||||
this->Attr("rope_head_dim").AttrType(REQUIRED).Int(ROPE_HEAD_DIM_VALUE);
|
||||
this->Attr("cmp_ratio").AttrType(REQUIRED).Int(CMP_RATIO_VALUE);
|
||||
this->Attr("coff").AttrType(OPTIONAL).Int(COFF_VALUE);
|
||||
this->Attr("norm_eps").AttrType(OPTIONAL).Float(1e-6f);
|
||||
this->Attr("rotary_mode").AttrType(OPTIONAL).Int(ROTARY_MODE_VALUE);
|
||||
this->Attr("cache_mode").AttrType(OPTIONAL).Int(CACHE_MODE_VALUE);
|
||||
this->Attr("state_cache_stride_dim0").AttrType(OPTIONAL).Int(STATE_CACHE_STRIDE_DIM0);
|
||||
OpAICoreConfig aicore_config;
|
||||
aicore_config.DynamicCompileStaticFlag(true)
|
||||
.DynamicFormatFlag(true)
|
||||
.DynamicRankSupportFlag(true)
|
||||
.DynamicShapeSupportFlag(true)
|
||||
.NeedCheckSupportFlag(false)
|
||||
.PrecisionReduceFlag(true)
|
||||
.ExtendCfgInfo("aclnnSupport.value", "support_aclnn"); // set value of aclnn support
|
||||
|
||||
OpAICoreConfig config910;
|
||||
config910.Input("x")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("wkv")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("wgate")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("state_cache")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.IgnoreContiguous();
|
||||
config910.Input("ape")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("norm_weight")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("rope_sin")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("rope_cos")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT, ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("state_block_table")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("cu_seqlens")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("seqused")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Input("start_pos")
|
||||
.ParamType(OPTIONAL)
|
||||
.DataTypeList({ge::DT_INT32})
|
||||
.FormatList({ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
config910.Output("cmp_kv")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16})
|
||||
.FormatList({ge::FORMAT_ND});
|
||||
config910.Output("state_cache")
|
||||
.ParamType(REQUIRED)
|
||||
.DataTypeList({ge::DT_FLOAT})
|
||||
.FormatList({ge::FORMAT_ND});
|
||||
config910.DynamicCompileStaticFlag(true)
|
||||
.DynamicFormatFlag(true)
|
||||
.DynamicRankSupportFlag(true)
|
||||
.DynamicShapeSupportFlag(true)
|
||||
.NeedCheckSupportFlag(false)
|
||||
.PrecisionReduceFlag(true)
|
||||
.ExtendCfgInfo("aclnnSupport.value", "support_aclnn");
|
||||
this->AICore().AddConfig("ascend910b", config910);
|
||||
this->AICore().AddConfig("ascend910_93", config910);
|
||||
this->AICore().AddConfig("ascend950", aicore_config);
|
||||
}
|
||||
};
|
||||
OP_ADD(Compressor, optiling::CompressorCompileInfo);
|
||||
} // namespace ops
|
||||
174
csrc/attention/compressor/op_host/compressor_proto.cpp
Normal file
174
csrc/attention/compressor/op_host/compressor_proto.cpp
Normal file
@@ -0,0 +1,174 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
#include <graph/utils/type_utils.h>
|
||||
#include <register/op_impl_registry.h>
|
||||
#include "log/ops_log.h"
|
||||
|
||||
using namespace ge;
|
||||
|
||||
namespace ops {
|
||||
// INPUT
|
||||
constexpr uint32_t TOKEN_X_INPUT_INDEX = 0;
|
||||
constexpr uint32_t WEIGHT_KV_INPUT_INDEX = 1;
|
||||
constexpr uint32_t WEIGHT_WGATE_INPUT_INDEX = 2;
|
||||
|
||||
constexpr uint32_t STATE_CACHE_INPUT_INDEX = 3;
|
||||
|
||||
constexpr uint32_t APE_INPUT_INDEX = 4;
|
||||
constexpr uint32_t NORM_WEIGHT_INPUT_INDEX = 5;
|
||||
constexpr uint32_t ROPE_SIN_INPUT_INDEX = 6;
|
||||
constexpr uint32_t ROPE_COS_INPUT_INDEX = 7;
|
||||
|
||||
// INPUT(OPTION)
|
||||
constexpr uint32_t STATE_BLOCK_TABLE_INPUT_INDEX = 8;
|
||||
|
||||
constexpr uint32_t CU_SEQ_LEN_INPUT_INDEX = 9;
|
||||
constexpr uint32_t SEQ_USED_INPUT_INDEX = 10;
|
||||
constexpr uint32_t START_POS_INPUT_INDEX = 11;
|
||||
|
||||
// ATTR
|
||||
constexpr uint32_t ROPE_HEAD_DIM_ATTR_INDEX = 0;
|
||||
constexpr uint32_t CMP_RATIO_ATTR_INDEX = 1;
|
||||
constexpr uint32_t COFF_ATTR_INDEX = 2;
|
||||
constexpr uint32_t NORM_EPS_ATTR_INDEX = 3;
|
||||
constexpr uint32_t ROTARY_MODE_ATTR_INDEX = 4;
|
||||
constexpr uint32_t CACHE_MODE_ATTR_INDEX = 5;
|
||||
constexpr uint32_t STATE_CACHE_STRIDE_DIM0_ATTR_INDEX = 6;
|
||||
|
||||
// OUTPUT
|
||||
constexpr uint32_t CMP_KV_OUTPUT_INDEX = 0;
|
||||
|
||||
// ATTR DEFAULT VALUE
|
||||
constexpr uint32_t CMP_RATIO_VALUE = 4;
|
||||
constexpr uint32_t COFF_VALUE = 1;
|
||||
|
||||
struct CompressorProtoShapeParam {
|
||||
bool isBsMerge { false };
|
||||
int64_t B { 0 };
|
||||
int64_t T { 0 };
|
||||
int64_t S { 0 };
|
||||
int64_t Sr { 0 };
|
||||
int64_t H { 0 };
|
||||
int64_t D { 0 };
|
||||
};
|
||||
|
||||
// tmp
|
||||
constexpr uint32_t DIM_NUM_1 = 1;
|
||||
constexpr uint32_t DIM_NUM_2 = 2;
|
||||
constexpr uint32_t DIM_NUM_3 = 3;
|
||||
constexpr uint32_t DIM_NUM_4 = 4;
|
||||
constexpr uint32_t DIM_INDEX_0 = 0;
|
||||
constexpr uint32_t DIM_INDEX_1 = 1;
|
||||
constexpr uint32_t DIM_INDEX_2 = 2;
|
||||
constexpr uint32_t DIM_INDEX_3 = 3;
|
||||
|
||||
ge::graphStatus GetCompressorShapeDim(const gert::InferShapeContext* context, CompressorProtoShapeParam &shapeParam)
|
||||
{
|
||||
auto xShape = context->GetRequiredInputShape(TOKEN_X_INPUT_INDEX); // (B, S, H) | (T, H)
|
||||
OPS_LOG_E_IF_NULL(context, xShape, return ge::GRAPH_FAILED)
|
||||
auto wkvShape = context->GetRequiredInputShape(WEIGHT_KV_INPUT_INDEX); // (coff * D, H)
|
||||
OPS_LOG_E_IF_NULL(context, wkvShape, return ge::GRAPH_FAILED)
|
||||
auto wgateShape = context->GetRequiredInputShape(WEIGHT_WGATE_INPUT_INDEX); // (coff * D, H)
|
||||
OPS_LOG_E_IF_NULL(context, wgateShape, return ge::GRAPH_FAILED)
|
||||
|
||||
auto stateCacheShape = context->GetRequiredInputShape(STATE_CACHE_INPUT_INDEX); // (block_num, block_size, 2 * coff * D) | (B, tokrn_size, 2 * coff * D)
|
||||
OPS_LOG_E_IF_NULL(context, stateCacheShape, return ge::GRAPH_FAILED)
|
||||
|
||||
auto apeShape = context->GetRequiredInputShape(APE_INPUT_INDEX); // (r, coff * D)
|
||||
OPS_LOG_E_IF_NULL(context, apeShape, return ge::GRAPH_FAILED)
|
||||
auto normWeightShape = context->GetRequiredInputShape(NORM_WEIGHT_INPUT_INDEX); // (D)
|
||||
OPS_LOG_E_IF_NULL(context, normWeightShape, return ge::GRAPH_FAILED)
|
||||
auto ropeSinShape = context->GetRequiredInputShape(ROPE_SIN_INPUT_INDEX); // (B, ceil(S / r), rD) | (min(T, T/r + B), rD)
|
||||
OPS_LOG_E_IF_NULL(context, ropeSinShape, return ge::GRAPH_FAILED)
|
||||
auto ropeCosShape = context->GetRequiredInputShape(ROPE_COS_INPUT_INDEX); // (B, ceil(S / r), rD) | (min(T, T/r + B), rD)
|
||||
OPS_LOG_E_IF_NULL(context, ropeCosShape, return ge::GRAPH_FAILED)
|
||||
|
||||
auto stateBlockTableShape = context->GetRequiredInputShape(STATE_BLOCK_TABLE_INPUT_INDEX); // (B, sMax/block_size) | (B, )
|
||||
OPS_LOG_E_IF_NULL(context, stateBlockTableShape, return ge::GRAPH_FAILED)
|
||||
|
||||
auto cuSeqlensShape = context->GetRequiredInputShape(CU_SEQ_LEN_INPUT_INDEX); // (B+1,)
|
||||
OPS_LOG_E_IF_NULL(context, cuSeqlensShape, return ge::GRAPH_FAILED)
|
||||
auto seqUsedShape = context->GetRequiredInputShape(SEQ_USED_INPUT_INDEX); // (B,)
|
||||
OPS_LOG_E_IF_NULL(context, seqUsedShape, return ge::GRAPH_FAILED)
|
||||
auto startPosShape = context->GetRequiredInputShape(START_POS_INPUT_INDEX); // (B,)
|
||||
OPS_LOG_E_IF_NULL(context, startPosShape, return ge::GRAPH_FAILED)
|
||||
|
||||
if (xShape->GetDimNum() == DIM_NUM_3) { // BS
|
||||
shapeParam.isBsMerge = false;
|
||||
shapeParam.B = xShape->GetDim(DIM_INDEX_0);
|
||||
shapeParam.S = xShape->GetDim(DIM_INDEX_1);
|
||||
shapeParam.H = xShape->GetDim(DIM_INDEX_2);
|
||||
shapeParam.T = shapeParam.B * shapeParam.S;
|
||||
} else { // T
|
||||
shapeParam.isBsMerge = true;
|
||||
shapeParam.T = xShape->GetDim(DIM_INDEX_0);
|
||||
shapeParam.H = xShape->GetDim(DIM_INDEX_1);
|
||||
}
|
||||
|
||||
shapeParam.D = normWeightShape->GetDim(DIM_INDEX_0);
|
||||
shapeParam.Sr = ropeSinShape->GetDim(DIM_INDEX_1);
|
||||
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus SetCompressorShapeDim(const CompressorProtoShapeParam &shapeParam, gert::InferShapeContext* context)
|
||||
{
|
||||
auto cmpKvShape = context->GetOutputShape(CMP_KV_OUTPUT_INDEX); // query: (B, S, N, Hckv) | (T, N, Hckv)
|
||||
OPS_LOG_E_IF_NULL(context, cmpKvShape, return ge::GRAPH_FAILED)
|
||||
auto attr = context->GetAttrs();
|
||||
const uint32_t *cmpRatioPtr = attr->GetAttrPointer<uint32_t>(CMP_RATIO_ATTR_INDEX);
|
||||
uint32_t cmpRatio = (cmpRatioPtr != nullptr) ? *cmpRatioPtr : CMP_RATIO_VALUE;
|
||||
const uint32_t *coffPtr = attr->GetAttrPointer<uint32_t>(COFF_ATTR_INDEX);
|
||||
uint32_t coff = (coffPtr != nullptr) ? *coffPtr : COFF_VALUE;
|
||||
// Set output shape
|
||||
if (!shapeParam.isBsMerge) {
|
||||
cmpKvShape->SetDimNum(DIM_NUM_3); // (B, Sr, H)
|
||||
cmpKvShape->SetDim(DIM_INDEX_0, shapeParam.B);
|
||||
cmpKvShape->SetDim(DIM_INDEX_1, shapeParam.Sr);
|
||||
cmpKvShape->SetDim(DIM_INDEX_2, shapeParam.H);
|
||||
} else {
|
||||
cmpKvShape->SetDimNum(DIM_NUM_2); // (T, N, Hckv)
|
||||
cmpKvShape->SetDim(DIM_INDEX_0, shapeParam.Sr);
|
||||
cmpKvShape->SetDim(DIM_INDEX_1, shapeParam.H);
|
||||
}
|
||||
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus InferDataTypeCompressor(gert::InferDataTypeContext* context)
|
||||
{
|
||||
OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Compressor", "Context is nullptr."),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_LOG_I(context->GetNodeName(), "Enter Compressor inferDataType impl.");
|
||||
|
||||
context->SetOutputDataType(CMP_KV_OUTPUT_INDEX, context->GetRequiredInputDataType(TOKEN_X_INPUT_INDEX));
|
||||
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
ge::graphStatus InferShapeCompressor(gert::InferShapeContext* context)
|
||||
{
|
||||
OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("Compressor", "Context is nullptr."),
|
||||
return ge::GRAPH_FAILED);
|
||||
OPS_LOG_I(context->GetNodeName(), "Enter Compressor infershape impl.");
|
||||
|
||||
CompressorProtoShapeParam shapeParam {};
|
||||
auto apiRet = GetCompressorShapeDim(context, shapeParam);
|
||||
OPS_LOG_E_IF((apiRet != GRAPH_SUCCESS), context, return ge::GRAPH_FAILED, "Context get input shape failed");
|
||||
|
||||
apiRet = SetCompressorShapeDim(shapeParam, context);
|
||||
OPS_LOG_E_IF((apiRet != GRAPH_SUCCESS), context, return ge::GRAPH_FAILED, "Context set output shape failed");
|
||||
|
||||
return GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
IMPL_OP_INFERSHAPE(Compressor).InferShape(InferShapeCompressor).InferDataType(InferDataTypeCompressor);
|
||||
} // namespace ops
|
||||
@@ -0,0 +1,408 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_block_cube_perf.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_BLOCK_CUBE_PERF_H
|
||||
#define COMPRESSOR_BLOCK_CUBE_PERF_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_tools.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
template<typename COMP> class CompressorBlockCubePerf {
|
||||
using MM1_OUT_T = float;
|
||||
public:
|
||||
__aicore__ inline CompressorBlockCubePerf(){};
|
||||
__aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools<COMP> &tools);
|
||||
__aicore__ inline void Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut);
|
||||
__aicore__ inline void InitBuffers(TPipe *pipe);
|
||||
__aicore__ inline void InitGlobalBuffers(const GlobalTensor<MM1_OUT_T>& kvMm1ResGm, const GlobalTensor<MM1_OUT_T>& scoreMm1ResGm);
|
||||
__aicore__ inline void AllocEventID(TPipe *pipe);
|
||||
__aicore__ inline void FreeEventID(TPipe *pipe);
|
||||
__aicore__ inline void ComputeMm1(const RunInfo &info);
|
||||
|
||||
private:
|
||||
using T = float;
|
||||
using X_T = typename AscendC::Conditional<COMP::xDtype == X_DTYPE::BF16, bfloat16_t, half>::type;
|
||||
|
||||
__aicore__ inline uint32_t GetMSize(const RunInfo &info, uint32_t coffId);
|
||||
__aicore__ inline void CopyXGmToL1(const RunInfo &info, LocalTensor<X_T> xL1Tensor, uint32_t hIdx, uint32_t kBase);
|
||||
__aicore__ inline void CopyWeightGmToL1(LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase, uint32_t coffId);
|
||||
__aicore__ inline void LoadAToL0(const RunInfo &info, LocalTensor<X_T> aL0Tensor, LocalTensor<X_T> xL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize);
|
||||
__aicore__ inline void LoadBToL0(LocalTensor<X_T> bL0Tensor, LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase);
|
||||
__aicore__ inline void MatrixMmad(LocalTensor<T> cL0Tensor, LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C);
|
||||
__aicore__ inline void CopyOutMm1Res(const RunInfo &info, LocalTensor<T> cL0Tensor,
|
||||
uint32_t coffId, uint32_t mStart, uint32_t mDealSize);
|
||||
|
||||
ConstInfo constInfo_ = {};
|
||||
CompressorTools<COMP> tools_;
|
||||
|
||||
// GM
|
||||
GlobalTensor<X_T> xGm_;
|
||||
GlobalTensor<X_T> wkvGm_;
|
||||
GlobalTensor<X_T> wgateGm_;
|
||||
GlobalTensor<MM1_OUT_T>kvMm1ResGm;
|
||||
GlobalTensor<MM1_OUT_T>scoreMm1ResGm;
|
||||
GlobalTensor<int32_t> cuSeqlensGm_;
|
||||
GlobalTensor<int32_t> sequsedGm_;
|
||||
GlobalTensor<int32_t> startPosGm_;
|
||||
bool isExistSeqUsed = false;
|
||||
|
||||
// =================================L1 Buffer=================================
|
||||
static constexpr uint32_t L1_X_SIZE = 128 * 1024;
|
||||
static constexpr uint32_t L1_W_SIZE = 64 * 1024;
|
||||
// L1 Buffer
|
||||
TBuf<TPosition::A1> xBufL1;
|
||||
TBuf<TPosition::A1> wBufL1;
|
||||
// =================================L0 Buffer=================================
|
||||
// L0 buffer size
|
||||
static constexpr uint32_t L0A_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k
|
||||
static constexpr uint32_t L0B_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k
|
||||
static constexpr uint32_t L0C_PP_SIZE = 64 * 1024; // (128 * 2) * 64 * 4 = 64k
|
||||
// L0_A
|
||||
TBuf<TPosition::A2> tmpBufL0A;
|
||||
// L0_B
|
||||
TBuf<TPosition::B2> tmpBufL0B;
|
||||
// L0_C
|
||||
TBuf<TPosition::CO1> tmpBufL0C;
|
||||
// =================================Event&Buffer ID===========================
|
||||
// mte2 <> mte1 EventID
|
||||
static constexpr uint32_t X_EVENT0 = EVENT_ID0;
|
||||
static constexpr uint32_t X_EVENT1 = EVENT_ID1;
|
||||
uint32_t xBufId = 0; // 用于DB计数
|
||||
static constexpr uint32_t W_EVENT0 = EVENT_ID4;
|
||||
static constexpr uint32_t W_EVENT1 = EVENT_ID5;
|
||||
static constexpr uint32_t W_EVENT2 = EVENT_ID6;
|
||||
static constexpr uint32_t W_EVENT3 = EVENT_ID7;
|
||||
uint32_t wBufId = 0; // 用于DB计数
|
||||
// mte1 <> mmad EventID
|
||||
static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3;
|
||||
static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4;
|
||||
uint32_t l0abBufId = 0;
|
||||
// mmad <> fixpipe EventID
|
||||
static constexpr uint32_t L0C_EVENT0 = EVENT_ID0; // 每块L0C单独分配EVENT_ID
|
||||
static constexpr uint32_t L0C_EVENT1 = EVENT_ID1;
|
||||
uint32_t l0cBufId = 0;
|
||||
|
||||
// =================================Loop======================================
|
||||
uint32_t curBIdx_ = 0;
|
||||
uint32_t curSIdx_ = 0;
|
||||
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::InitParams(const ConstInfo &constInfo, const CompressorTools<COMP> &tools)
|
||||
{
|
||||
this->constInfo_ = constInfo;
|
||||
this->tools_ = tools;
|
||||
}
|
||||
|
||||
template <typename COMP> __aicore__ inline void CompressorBlockCubePerf<COMP>::Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut)
|
||||
{
|
||||
xGm_.SetGlobalBuffer((__gm__ X_T *)x);
|
||||
wkvGm_.SetGlobalBuffer((__gm__ X_T *)wKv);
|
||||
wgateGm_.SetGlobalBuffer((__gm__ X_T *)wGate);
|
||||
startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos);
|
||||
isExistSeqUsed = (seqUsed != nullptr);
|
||||
if (isExistSeqUsed) {
|
||||
sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed);
|
||||
}
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::InitBuffers(TPipe *pipe)
|
||||
{
|
||||
// L1
|
||||
// 1. coff=1时, mBase=256, kL1=256, X单次拷贝到L1的数据量最大为mBase*kL1*sizeof(BF16/FP16)=256*256*2=128K
|
||||
// 2. coff=2时, mBase=128, kL1=256, r最大为128, X单次拷贝到L1的最大数据量为(128+r)*kL1*sizeof(BF16/FP16)<=128K
|
||||
pipe->InitBuffer(xBufL1, L1_X_SIZE * 2);
|
||||
// dBaseSize<=64, wkv和wgate各一份, kL1=256, 右矩阵为dBaseSize*2*sizeof(BF16/FP16)<=64K
|
||||
// cur和pre循环使用, 2份buffer就足够
|
||||
pipe->InitBuffer(wBufL1, L1_W_SIZE * 4);
|
||||
|
||||
// L0
|
||||
pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2);
|
||||
pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2);
|
||||
pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 2);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::InitGlobalBuffers(const GlobalTensor<MM1_OUT_T>& kvMm1ResGm, const GlobalTensor<MM1_OUT_T>& scoreMm1ResGm)
|
||||
{
|
||||
this->kvMm1ResGm = kvMm1ResGm;
|
||||
this->scoreMm1ResGm = scoreMm1ResGm;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::AllocEventID(TPipe *pipe)
|
||||
{
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT0);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT1);
|
||||
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT0);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT1);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT2);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT3);
|
||||
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT0);
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT1);
|
||||
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT0);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT1);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::FreeEventID(TPipe *pipe)
|
||||
{
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT0);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT1);
|
||||
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT0);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT1);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT2);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT3);
|
||||
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT0);
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT1);
|
||||
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT0);
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT1);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::CopyXGmToL1(const RunInfo &info, LocalTensor<X_T> xL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase)
|
||||
{
|
||||
uint32_t tStart = tools_.GetTIdxByBatch(info.bStart) + info.sStart; // 此基本块在整个序列中的位置
|
||||
uint32_t copySeqCnt = info.dealSeqCnt; // 此基本块处理的长度
|
||||
|
||||
uint32_t xL1Offset = 0 * (32 / sizeof(X_T));
|
||||
uint64_t sIdx = tStart; // 起始s在整个T的起始点
|
||||
uint64_t gmOffset = sIdx * constInfo_.hSize + hIdx;
|
||||
uint32_t nValue = copySeqCnt;
|
||||
uint32_t dValue = kBase; // 拷贝的列数kBase
|
||||
uint32_t srcDValue = constInfo_.hSize;
|
||||
uint32_t dstNzC0Stride = (copySeqCnt + 15) / 16 * 16; // 1行变2行的行方向的偏移,需要16对齐
|
||||
CopySingleMatrixNDToNZ(xL1Tensor[xL1Offset], xGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::CopyWeightGmToL1(LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase, uint32_t coffId)
|
||||
{
|
||||
// coffId=0, 搬运左矩阵的数据; coffId=1, 搬运右矩阵的数据
|
||||
uint64_t gmOffset = coffId * constInfo_.headDim * constInfo_.hSize + constInfo_.dIdx * constInfo_.hSize + hIdx;
|
||||
uint32_t wkvL1Offset = 0;
|
||||
uint32_t wgateL1Offset = constInfo_.dBaseSize * (32 / sizeof(X_T)); // wgate与wkv的起始点相隔dBaseSize个32B
|
||||
uint32_t nValue = constInfo_.dBaseSize;
|
||||
uint32_t dValue = kBase;
|
||||
uint32_t srcDValue = constInfo_.hSize;
|
||||
uint32_t dstNzC0Stride = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
|
||||
CopySingleMatrixNDToNZ(wL1Tensor[wkvL1Offset], wkvGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
CopySingleMatrixNDToNZ(wL1Tensor[wgateL1Offset], wgateGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::LoadAToL0(const RunInfo &info, LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> xL1Tensor, uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize)
|
||||
{
|
||||
uint32_t mSize = info.dealSeqCnt;
|
||||
|
||||
uint32_t mSizeAlign = Align(mSize, 16U);
|
||||
uint32_t xTensorOffset = kStart * mSizeAlign + mStart * (32 / sizeof(X_T));
|
||||
uint32_t mLoop = Align(mDealSize, 16U) / 16;
|
||||
|
||||
for (uint32_t i = 0; i < mLoop; i++) {
|
||||
LoadData2DParams loadData2DParams;
|
||||
loadData2DParams.startIndex = i;
|
||||
loadData2DParams.repeatTimes = kBase / (32 / sizeof(X_T));
|
||||
loadData2DParams.srcStride = mSizeAlign / 16;
|
||||
loadData2DParams.dstGap = 0;
|
||||
loadData2DParams.ifTranspose = false;
|
||||
LoadData(aL0Tensor[i * 16 * kBase], xL1Tensor[xTensorOffset], loadData2DParams); // 16: 一个分型的行数
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::LoadBToL0(LocalTensor<X_T> bL0Tensor, LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase)
|
||||
{
|
||||
uint32_t rowCnt = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
|
||||
uint64_t wTensorOffset = rowCnt * kStart;
|
||||
LoadData2DParams loadData2DParams;
|
||||
loadData2DParams.startIndex = 0;
|
||||
loadData2DParams.repeatTimes = (rowCnt / 16) * (kBase / (32 / sizeof(X_T)));
|
||||
loadData2DParams.srcStride = 1;
|
||||
loadData2DParams.dstGap = 0;
|
||||
loadData2DParams.ifTranspose = false;
|
||||
LoadData(bL0Tensor, wL1Tensor[wTensorOffset], loadData2DParams);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::MatrixMmad(LocalTensor<T> cL0Tensor, LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C)
|
||||
{
|
||||
MmadParams mmadParams;
|
||||
mmadParams.m = (mActSize + 15) / 16 * 16;
|
||||
mmadParams.n = nDealSize;
|
||||
mmadParams.k = kActSize;
|
||||
mmadParams.cmatrixInitVal = isInitL0C;
|
||||
mmadParams.cmatrixSource = false;
|
||||
Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams);
|
||||
PipeBarrier<PIPE_M>();
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::CopyOutMm1Res(const RunInfo &info, LocalTensor<T> cL0Tensor,
|
||||
uint32_t coffId, uint32_t mStart, uint32_t mDealSize)
|
||||
{
|
||||
// coffId=0, 存左矩阵的数据; coffId=1, 存右矩阵的数据
|
||||
FixpipeParamsV220 fixParams;
|
||||
fixParams.mSize = mDealSize;
|
||||
fixParams.nSize = constInfo_.dBaseSize;
|
||||
fixParams.srcStride = (mDealSize + 15) / 16 * 16; // 需要16对齐
|
||||
fixParams.dstStride = (uint32_t)COMP::coff * constInfo_.headDim;
|
||||
fixParams.ndNum = 1;
|
||||
|
||||
uint64_t dbOffset = info.cubeDbIdx * constInfo_.dbSize;
|
||||
uint64_t gmOffset = coffId * constInfo_.headDim + constInfo_.dIdx + mStart * fixParams.dstStride + dbOffset;
|
||||
uint32_t kvOffset = 0;
|
||||
uint32_t scoreOffset = (mDealSize + 15) / 16 * 16 * constInfo_.dBaseSize;
|
||||
|
||||
Fixpipe(kvMm1ResGm[gmOffset], cL0Tensor[kvOffset], fixParams);
|
||||
Fixpipe(scoreMm1ResGm[gmOffset], cL0Tensor[scoreOffset], fixParams);
|
||||
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorBlockCubePerf<COMP>::GetMSize(const RunInfo &info, uint32_t coffId)
|
||||
{
|
||||
return info.dealSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubePerf<COMP>::ComputeMm1(const RunInfo &info)
|
||||
{
|
||||
static constexpr uint32_t K_SIZE = 512;
|
||||
static constexpr uint32_t K_L1_BASE = 256;
|
||||
static constexpr uint32_t M_L0_BASE = 128;
|
||||
static constexpr uint32_t K_L0_BASE = 128;
|
||||
uint32_t nCoff = (uint32_t)COMP::coff;
|
||||
|
||||
// hSize为K_SIZE=512的倍数
|
||||
uint32_t hSize = constInfo_.hSize;
|
||||
uint32_t hIdxStart = (constInfo_.aiCoreIdx % constInfo_.dBasicBlockNum) * K_L1_BASE; // 每组核内的h循环起始不同
|
||||
for (uint32_t h = 0; h < hSize; h += K_SIZE) {
|
||||
for (uint32_t k = 0; k < K_SIZE; k += K_L1_BASE) {
|
||||
bool isFirst = (h == 0 && k == 0);
|
||||
bool isLast = ((h + K_SIZE >= hSize) && (k + K_L1_BASE >= K_SIZE));
|
||||
uint32_t hIdx = (h + k + hIdxStart) % hSize; // h方向错位搬运
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT0 + xBufId);
|
||||
LocalTensor<X_T> xL1Tensor = xBufL1.GetWithOffset<X_T>(L1_X_SIZE / sizeof(X_T), xBufId * L1_X_SIZE);
|
||||
CopyXGmToL1(info, xL1Tensor, hIdx, K_L1_BASE);
|
||||
SetFlag<HardEvent::MTE2_MTE1>(X_EVENT0 + xBufId);
|
||||
WaitFlag<HardEvent::MTE2_MTE1>(X_EVENT0 + xBufId);
|
||||
for (uint32_t i = nCoff; i > 0; i--) {
|
||||
// coffId=0, 计算pre数据; coffId=1, 计算cur数据
|
||||
uint32_t coffId = i - 1;
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT0 + wBufId);
|
||||
LocalTensor<X_T> wL1Tensor = wBufL1.GetWithOffset<X_T>(L1_W_SIZE / sizeof(X_T), wBufId * L1_W_SIZE);
|
||||
CopyWeightGmToL1(wL1Tensor, hIdx, K_L1_BASE, coffId);
|
||||
SetFlag<HardEvent::MTE2_MTE1>(W_EVENT0 + wBufId);
|
||||
WaitFlag<HardEvent::MTE2_MTE1>(W_EVENT0 + wBufId);
|
||||
|
||||
uint32_t mSize = GetMSize(info, coffId);
|
||||
uint32_t actMDealSize = M_L0_BASE;
|
||||
for (uint32_t mL0 = 0; mL0 < mSize; mL0 += M_L0_BASE) {
|
||||
if (mL0 + M_L0_BASE > mSize) {
|
||||
actMDealSize = mSize - mL0;
|
||||
}
|
||||
|
||||
l0cBufId = coffId + (mL0 / M_L0_BASE);
|
||||
LocalTensor<T> cL0Tensor = tmpBufL0C.GetWithOffset<T>((L0C_PP_SIZE / sizeof(T)), l0cBufId * L0C_PP_SIZE);
|
||||
if (isFirst) {
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT0 + l0cBufId);
|
||||
}
|
||||
uint32_t nDealSize = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
|
||||
for (uint32_t kL0 = 0; kL0 < K_L1_BASE; kL0 += K_L0_BASE) {
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT0 + l0abBufId);
|
||||
LocalTensor<X_T> aL0Tensor = tmpBufL0A.GetWithOffset<X_T>(L0A_PP_SIZE / sizeof(X_T), l0abBufId * L0A_PP_SIZE);
|
||||
LocalTensor<X_T> bL0Tensor = tmpBufL0B.GetWithOffset<X_T>(L0B_PP_SIZE / sizeof(X_T), l0abBufId * L0B_PP_SIZE);
|
||||
LoadAToL0(info, aL0Tensor, xL1Tensor, kL0, K_L0_BASE, mL0, actMDealSize);
|
||||
LoadBToL0(bL0Tensor, wL1Tensor, kL0, K_L0_BASE);
|
||||
SetFlag<HardEvent::MTE1_M>(L0AB_EVENT0 + l0abBufId);
|
||||
WaitFlag<HardEvent::MTE1_M>(L0AB_EVENT0 + l0abBufId);
|
||||
bool isInitL0C = isFirst && (kL0 == 0);
|
||||
MatrixMmad(cL0Tensor, aL0Tensor, bL0Tensor, actMDealSize, nDealSize, K_L0_BASE, isInitL0C);
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT0 + l0abBufId);
|
||||
l0abBufId = (l0abBufId + 1) % 2;
|
||||
}
|
||||
if (isLast) {
|
||||
SetFlag<HardEvent::M_FIX>(L0C_EVENT0 + l0cBufId);
|
||||
WaitFlag<HardEvent::M_FIX>(L0C_EVENT0 + l0cBufId);
|
||||
CopyOutMm1Res(info, cL0Tensor, coffId, mL0, actMDealSize);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT0 + l0cBufId);
|
||||
}
|
||||
}
|
||||
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT0 + wBufId);
|
||||
wBufId = (wBufId + 1) % 4;
|
||||
}
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT0 + xBufId);
|
||||
xBufId = (xBufId + 1) % 2;
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif // COMPRESSOR_BLOCK_CUBE_PERF_H
|
||||
File diff suppressed because it is too large
Load Diff
341
csrc/attention/compressor/op_kernel/arch32/compressor_comm.h
Normal file
341
csrc/attention/compressor/op_kernel/arch32/compressor_comm.h
Normal file
@@ -0,0 +1,341 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_comm.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_COMM_H
|
||||
#define COMPRESSOR_COMM_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include "kernel_operator_list_tensor_intf.h"
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
#include "lib/matrix/matmul/tiling.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
template <typename T>
|
||||
__aicore__ inline T CeilDivT(T num1, T num2)
|
||||
{
|
||||
if (num2 == 0) {
|
||||
return static_cast<T>(0);
|
||||
}
|
||||
return (num1 + num2 - 1) / num2;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Align(T num, T rnd)
|
||||
{
|
||||
return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd)));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Trunc(T num, T rnd)
|
||||
{
|
||||
return ((rnd) == 0) ? 0 : (((num) / (rnd) * (rnd)));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T FloorPow2(T num)
|
||||
{
|
||||
if (num == 0) return 1;
|
||||
for(uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) {
|
||||
num |= (num >> i);
|
||||
}
|
||||
return num - (num >> 1);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T CeilPow2(T num)
|
||||
{
|
||||
if (num <= 1) return 1;
|
||||
num --;
|
||||
for(uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) {
|
||||
num |= (num >> i);
|
||||
}
|
||||
num ++;
|
||||
return num;
|
||||
}
|
||||
|
||||
enum class X_LAYOUT : std::uint8_t {
|
||||
BSH = static_cast<std::uint8_t>(0),
|
||||
TH = static_cast<std::uint8_t>(1)
|
||||
};
|
||||
|
||||
enum class X_DTYPE : std::uint8_t {
|
||||
BF16 = static_cast<std::uint8_t>(0),
|
||||
FP16 = static_cast<std::uint8_t>(1)
|
||||
};
|
||||
|
||||
enum class ROPE_DTYPE : std::uint8_t {
|
||||
SAME_AS_X = static_cast<std::uint8_t>(0),
|
||||
FP32 = static_cast<std::uint8_t>(1)
|
||||
};
|
||||
|
||||
enum class COFF : std::uint8_t {
|
||||
DISABLE = static_cast<std::uint8_t>(1),
|
||||
OVERLAP = static_cast<std::uint8_t>(2)
|
||||
};
|
||||
|
||||
enum class ROTARY_MODE : std::uint8_t {
|
||||
HALF = static_cast<std::uint8_t>(1),
|
||||
INTERLEAVE = static_cast<std::uint8_t>(2)
|
||||
};
|
||||
|
||||
enum class CACHE_MODE : std::uint8_t {
|
||||
CONTINUOUS = static_cast<std::uint8_t>(1),
|
||||
CYCLE = static_cast<std::uint8_t>(2)
|
||||
};
|
||||
|
||||
enum class TEMPLATE_ID : uint8_t {
|
||||
NORMAL = 0,
|
||||
EMPTY_X = 1,
|
||||
PERF = 2
|
||||
};
|
||||
|
||||
template <X_LAYOUT X_L, X_DTYPE X_T, ROPE_DTYPE R_T, COFF C, ROTARY_MODE Rotary_Mode, typename... Args>
|
||||
struct COMPType {
|
||||
static constexpr X_LAYOUT xLayout = X_L;
|
||||
static constexpr X_DTYPE xDtype = X_T;
|
||||
static constexpr ROPE_DTYPE ropeDtype = R_T;
|
||||
static constexpr COFF coff = C;
|
||||
static constexpr ROTARY_MODE rotaryMode = Rotary_Mode;
|
||||
};
|
||||
|
||||
struct ConstInfo {
|
||||
// 整个AICORE的任务信息, 左闭右开区间[ (bStart, s2Start), (bEnd, s2End) )
|
||||
uint32_t bStart = 0U;
|
||||
uint32_t sStart = 0U;
|
||||
uint32_t bEnd = 0U;
|
||||
uint32_t sEnd = 0U;
|
||||
|
||||
// 分核相关
|
||||
uint32_t usedCoreNum = 0;
|
||||
uint32_t dBaseSize = 0;
|
||||
uint32_t mBaseSize = 0;
|
||||
uint32_t tcSize = 0;
|
||||
uint32_t tcBaseSize = 0;
|
||||
uint32_t tcBasicBlockNum = 0;
|
||||
uint32_t dBasicBlockNum = 0;
|
||||
uint32_t coreGroupNum = 0;
|
||||
uint32_t singleCoreDealTcBasicNum = 0;
|
||||
uint32_t dIdx = 0;
|
||||
uint32_t bIdxOfLastTc = 0;
|
||||
uint32_t sIdxOfLastTc = 0;
|
||||
|
||||
// shape及参数
|
||||
uint32_t batchSize = 0;
|
||||
uint32_t hSize = 0;
|
||||
uint32_t sSize = 0;
|
||||
uint32_t headDim = 0;
|
||||
uint32_t ropeHeadDim = 0;
|
||||
uint32_t cmpRatio = 0;
|
||||
float normEps = 1e-6;
|
||||
float reciprocalD = 0;
|
||||
|
||||
uint32_t curGroupIdx = 0;
|
||||
uint32_t tailGroupIdx = 0;
|
||||
uint32_t tailBasicBlockNum = 0;
|
||||
uint32_t realDealBasicBlockNum = 0;
|
||||
|
||||
// pageAttention
|
||||
uint32_t blockNum = 0;
|
||||
uint32_t blockSize = 0;
|
||||
uint32_t maxBlockNumPerBatch = 0;
|
||||
uint64_t stateCacheStrideDim0 = 0;
|
||||
|
||||
// workSpace
|
||||
uint32_t dbWorkspaceRatio = 1;
|
||||
uint32_t mm1KvResSize = 0;
|
||||
uint32_t mm1ScoreResSize = 0;
|
||||
uint32_t vec1TailCacheSize = 0;
|
||||
uint32_t vec1ResSize = 0;
|
||||
uint32_t mm1ResSize = 0; // 所有cube输出kv/score结果的总大小
|
||||
|
||||
uint32_t aiCoreIdx = 0;
|
||||
uint32_t nSize = 0;
|
||||
|
||||
uint32_t dbSize = 0;
|
||||
};
|
||||
|
||||
struct RunInfo {
|
||||
bool isValid = false;
|
||||
uint32_t cubeDbIdx = 0; // kernel主循环索引
|
||||
|
||||
// 增加字段
|
||||
uint32_t dealTcNum = 0;
|
||||
// 右边相关信息
|
||||
uint32_t bStart = 0;
|
||||
uint32_t sStart = 0;
|
||||
uint32_t dealSeqCnt = 0;
|
||||
// 左边相关信息
|
||||
uint32_t preBStart = 0;
|
||||
uint32_t preSStart = 0;
|
||||
uint32_t preDealSeqCnt = 0; // 左边需要处理的s大小
|
||||
uint32_t preFirstSeqCnt = 0; // 左边首块大小
|
||||
|
||||
|
||||
uint32_t bEnd = 0;
|
||||
uint32_t sEnd = 0;
|
||||
uint32_t bStartSeqIdx = 0;
|
||||
uint32_t bEndSeqIdx = 0;
|
||||
|
||||
// v2分核信息 sc是左闭右开
|
||||
uint32_t scStart = 0;
|
||||
uint32_t scEnd = 0;
|
||||
uint32_t dealScSize = 0;
|
||||
|
||||
// vec1Res offset
|
||||
uint64_t vec1ResOffset = 0;
|
||||
};
|
||||
|
||||
struct Vec1RunInfo {
|
||||
// vec相关信息,一次syncAll需处理数据的起始索引
|
||||
bool resetResFlag = false; // v1积攒N轮 是否是N轮的起始轮
|
||||
uint32_t c1v1DbIdx = 0; // vec1 doubleBuffer索引
|
||||
uint32_t v1v2DbIdx = 0; // v1v2 doubleBuffer索引
|
||||
uint32_t bStart = 0;
|
||||
uint32_t sStart = 0;
|
||||
uint32_t dealTcNum = 0;
|
||||
uint32_t dealScSize = 0;
|
||||
};
|
||||
|
||||
struct Vec2RunInfo {
|
||||
// uint32_t bStart = 0;
|
||||
uint32_t v2DbIdx = 0; // v2 doubleBuffer索引
|
||||
uint32_t sStart = 0;
|
||||
uint32_t bEnd = 0;
|
||||
uint32_t sEnd = 0;
|
||||
// v2分核信息 sc是左闭右开
|
||||
uint32_t scStart = 0;
|
||||
uint32_t scEnd = 0;
|
||||
// uint32_t dealScSize = 0;
|
||||
|
||||
// 增加字段
|
||||
uint32_t bStart = 0;
|
||||
uint32_t compressedId = 0;
|
||||
uint32_t bCompressedId = 0;
|
||||
uint32_t dealScSize = 0;
|
||||
};
|
||||
|
||||
struct MSplitInfo {
|
||||
uint32_t vecStartB = 0U;
|
||||
uint32_t vecStartS = 0U;
|
||||
uint32_t vecEndB = 0U;
|
||||
uint32_t vecEndS = 0U;
|
||||
uint32_t dealTcNum = 0U;
|
||||
// vec1Res offset
|
||||
uint64_t vec1StartOffset = 0;
|
||||
uint64_t vec1ResOffset = 0;
|
||||
};
|
||||
|
||||
struct BlockInfo {
|
||||
__aicore__ inline BlockInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealSeqSize) :
|
||||
bIdx(bIdx), sIdx(sIdx), dealSeqSize(dealSeqSize) {};
|
||||
uint32_t bIdx = 0U;
|
||||
uint32_t sIdx = 0U;
|
||||
uint32_t dealSeqSize = 0;
|
||||
|
||||
uint32_t isFirst = true;
|
||||
uint32_t bSeqUsed = 0U;
|
||||
uint32_t bStartPos = 0U;
|
||||
uint32_t headHolderSeqCnt = 0U;
|
||||
uint32_t validSeqCnt = 0U;
|
||||
uint32_t tailHolderSeqCnt = 0U;
|
||||
uint32_t dealTcSize = 0U;
|
||||
uint32_t tailValidSeqCnt = 0U;
|
||||
uint32_t compressTcSize = 0U;
|
||||
};
|
||||
|
||||
// BUFFER的字节数
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_64K = 65536;
|
||||
|
||||
// BLOCK和REPEAT的字节数
|
||||
inline constexpr uint64_t BYTE_BLOCK = 32UL;
|
||||
inline constexpr uint32_t REPEAT_BLOCK_BYTE = 256U;
|
||||
// BLOCK和REPEAT的FP32元素数
|
||||
inline constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); // 8
|
||||
inline constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); // 64
|
||||
inline constexpr uint32_t REPEAT_STRIDE_NUM = REPEAT_BLOCK_BYTE / BYTE_BLOCK; // 8
|
||||
inline constexpr uint32_t REPEAT_MAX_NUM = 255;
|
||||
inline constexpr uint32_t BRCB_NUM = 8;
|
||||
inline constexpr uint32_t MAX_R = 256;
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void CopySingleMatrixNDToNZ(LocalTensor<T> l1Tensor, const GlobalTensor<T> gmTensor,
|
||||
uint32_t nValue, uint32_t dValue, uint32_t srcDValue, uint32_t dstNzC0Stride)
|
||||
{
|
||||
Nd2NzParams nd2nzPara;
|
||||
nd2nzPara.ndNum = 1;
|
||||
nd2nzPara.nValue = nValue; // nd矩阵的行数
|
||||
if constexpr (IsSameType<T, int4b_t>::value) {
|
||||
constexpr uint32_t HALF_SIZE_DIVISOR = 2;
|
||||
nd2nzPara.dValue = dValue / HALF_SIZE_DIVISOR;
|
||||
nd2nzPara.srcDValue = srcDValue / HALF_SIZE_DIVISOR;
|
||||
} else {
|
||||
nd2nzPara.dValue = dValue; // nd矩阵的列数
|
||||
nd2nzPara.srcDValue = srcDValue; // 同一nd矩阵相邻行起始地址间的偏移
|
||||
}
|
||||
nd2nzPara.dstNzC0Stride = dstNzC0Stride;
|
||||
nd2nzPara.dstNzNStride = 1;
|
||||
nd2nzPara.srcNdMatrixStride = 0;
|
||||
nd2nzPara.dstNzMatrixStride = 0;
|
||||
DataCopy(l1Tensor, gmTensor, nd2nzPara);
|
||||
}
|
||||
template <typename T>
|
||||
__aicore__ inline void DumpTensorForDim2(GlobalTensor<T> tensor, uint32_t desc, uint32_t dumpSize, uint32_t row, uint32_t col)
|
||||
{
|
||||
uint32_t array2[] = {static_cast<uint32_t>(row), static_cast<uint32_t>(col)};
|
||||
AscendC::ShapeInfo shapeInfo(2, array2);
|
||||
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DumpTensorForDim2(LocalTensor<T> tensor, uint32_t desc, uint32_t dumpSize, uint32_t row, uint32_t col)
|
||||
{
|
||||
uint32_t array2[] = {static_cast<uint32_t>(row), static_cast<uint32_t>(col)};
|
||||
AscendC::ShapeInfo shapeInfo(2, array2);
|
||||
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DumpTensorForDim2(LocalTensor<T> tensor, uint32_t desc, uint32_t dumpSize)
|
||||
{
|
||||
uint32_t col = 32 / sizeof(T);
|
||||
uint32_t array2[] = {static_cast<uint32_t>(dumpSize / col), static_cast<uint32_t>(col)};
|
||||
AscendC::ShapeInfo shapeInfo(2, array2);
|
||||
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DumpTensorForDim2(GlobalTensor<T> tensor, uint32_t desc, uint32_t dumpSize)
|
||||
{
|
||||
uint32_t col = 32 / sizeof(T);
|
||||
uint32_t array2[] = {static_cast<uint32_t>(dumpSize / col), static_cast<uint32_t>(col)};
|
||||
AscendC::ShapeInfo shapeInfo(2, array2);
|
||||
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
#endif
|
||||
@@ -0,0 +1,86 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_kernel.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_KERNEL
|
||||
#define COMPRESSOR_KERNEL
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_template_tiling_key.h"
|
||||
#include "compressor_kernel_perf.h"
|
||||
#include "compressor_tiling_data.h"
|
||||
#include "compressor_tools.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorKernel {
|
||||
public:
|
||||
__aicore__ inline CompressorKernel(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData)
|
||||
: pipe_(pipe), tilingData_(tilingData) {}
|
||||
|
||||
__aicore__ inline void Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *workspace);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
// ==============================TilingData&TPipe==============================
|
||||
TPipe* pipe_;
|
||||
const optiling::CompressorTilingData* __restrict tilingData_;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *workspace)
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::Process()
|
||||
{
|
||||
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif // COMPRESSOR_KERNEL
|
||||
@@ -0,0 +1,695 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_kernel_perf.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_KERNEL_PERF_H
|
||||
#define COMPRESSOR_KERNEL_PERF_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_template_tiling_key.h"
|
||||
#include "compressor_tiling_data.h"
|
||||
#include "compressor_tools.h"
|
||||
#include "compressor_block_cube_perf.h"
|
||||
#include "compressor_block_vec_perf.h"
|
||||
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
struct CmpBlockInfo {
|
||||
__aicore__ inline CmpBlockInfo() {};
|
||||
__aicore__ inline CmpBlockInfo(uint32_t bIdx, uint32_t sIdx, bool needReset = false) : bIdx(bIdx), sIdx(sIdx), needReset(needReset) {};
|
||||
|
||||
uint32_t bIdx = 0U;
|
||||
uint32_t sIdx = 0U;
|
||||
uint32_t bSeqUsed = 0U;
|
||||
uint32_t bStartPos = 0U;
|
||||
bool needReset = false;
|
||||
bool isFirst = true;
|
||||
|
||||
uint32_t headSeqCnt = 0U;
|
||||
uint32_t validSeqCnt = 0U;
|
||||
uint32_t tailSeqCnt = 0U;
|
||||
bool isCompress = 0U;
|
||||
};
|
||||
|
||||
struct BasicBlockInfo {
|
||||
uint32_t bIdx = 0;
|
||||
uint32_t sIdx = 0;
|
||||
uint32_t compressedTcNum = 0;
|
||||
uint32_t dealSeqCnt = 0;
|
||||
uint32_t dealTcNum = 0;
|
||||
};
|
||||
|
||||
struct BatchInfo {
|
||||
uint32_t tcNum = 0;
|
||||
uint32_t compressedTcNum = 0;
|
||||
uint32_t remSeqCnt = 0;
|
||||
uint32_t seqCnt = 0;
|
||||
uint32_t seqUsedCnt = 0;
|
||||
uint32_t headHolderSeq = 0;
|
||||
uint32_t bStartPos = 0;
|
||||
uint32_t bIdx = 0;
|
||||
uint32_t sIdx = 0;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorKernelPerf {
|
||||
public:
|
||||
__aicore__ inline CompressorKernelPerf(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData)
|
||||
: pipe_(pipe), tilingData_(tilingData) {}
|
||||
|
||||
__aicore__ inline void Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *workspace);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
// ================================Init functions==================================
|
||||
__aicore__ inline void InitWorkspace(__gm__ uint8_t *workspace);
|
||||
// ================================Process functions================================
|
||||
__aicore__ inline void InitTilingData();
|
||||
__aicore__ inline void SetBaseSize();
|
||||
// 获取基本块数量
|
||||
__aicore__ inline uint32_t GetLoopTimes();
|
||||
__aicore__ inline void SkipInvalidBatch(BatchInfo &batchInfo);
|
||||
__aicore__ inline void UpdateCurGroup(BasicBlockInfo &basicBlockInfo, BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq);
|
||||
__aicore__ inline BasicBlockInfo SkipOneLoop(BatchInfo &batchInfo);
|
||||
// 计算分核基本信息
|
||||
__aicore__ inline void CalcSplitCoreInfo();
|
||||
|
||||
__aicore__ inline void AllocEventID();
|
||||
__aicore__ inline void FreeEventID();
|
||||
__aicore__ inline void ComputeMm1(const RunInfo &info, bool isNeedExcute);
|
||||
__aicore__ inline void ComputeVec1(const Vec1RunInfo &info);
|
||||
__aicore__ inline void ComputeVec2(const Vec2RunInfo &info);
|
||||
|
||||
__aicore__ inline bool IsNeedExcuteC1(RunInfo info);
|
||||
__aicore__ inline bool IsNeedSyncAll(uint32_t curBasicBlockIdx);
|
||||
__aicore__ inline void CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx);
|
||||
__aicore__ inline void UpdateVec2Info(Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info);
|
||||
__aicore__ inline bool IsNeedExcuteV2(Vec2RunInfo &vec2Info);
|
||||
|
||||
using X_T = typename AscendC::Conditional<COMP::xDtype == X_DTYPE::BF16, bfloat16_t, half>::type;
|
||||
using T = float;
|
||||
using MM1_OUT_T = T;
|
||||
using VEC1_OUT_T = T;
|
||||
|
||||
// 常量
|
||||
static constexpr uint64_t SYNC_MODE0 = 0;
|
||||
static constexpr uint64_t SYNC_MODE2 = 2;
|
||||
static constexpr uint32_t SYNC_C1_FLAG = 3;
|
||||
static constexpr uint32_t SYNC_V1_FLAG = 4;
|
||||
static constexpr uint32_t SYNC_V1_FLAG2 = 5;
|
||||
static constexpr uint32_t SYNC_C1_V1_FLAG = 6;
|
||||
static constexpr uint32_t SYNC_V1_C1_FLAG = 8;
|
||||
|
||||
// ==============================TilingData&TPipe==============================
|
||||
TPipe* pipe_;
|
||||
const optiling::CompressorTilingData* __restrict tilingData_;
|
||||
// ===========================Workspace Global Tensor===========================
|
||||
GlobalTensor<MM1_OUT_T> mm1KvResGm;
|
||||
GlobalTensor<MM1_OUT_T> mm1ScoreResGm;
|
||||
GlobalTensor<MM1_OUT_T> vec1KvCacheGm;
|
||||
GlobalTensor<MM1_OUT_T> vec1ScoreCacheGm;
|
||||
GlobalTensor<MM1_OUT_T> Vec1InputKvGm;
|
||||
GlobalTensor<MM1_OUT_T> Vec1InputScoreGm;
|
||||
GlobalTensor<VEC1_OUT_T> vec1ResGm;
|
||||
GlobalTensor<VEC1_OUT_T> vec2InputGm;
|
||||
// ================================Task Info====================================
|
||||
CompressorTools<COMP> tools_;
|
||||
ConstInfo constInfo{};
|
||||
uint32_t aiCoreIdx = 0;
|
||||
|
||||
// ==============================Service Define==============================
|
||||
CompressorBlockCubePerf<COMP> blockCube_;
|
||||
CompressorBlockVectorPerf<COMP> blockVec_;
|
||||
|
||||
uint32_t allCompressedTcNum_ = 0;
|
||||
uint32_t curCompressedTcNum_ = 0;
|
||||
uint32_t accDealSize = 0;
|
||||
uint32_t loopTimes = 0;
|
||||
uint32_t cubeLoop = 0;
|
||||
uint32_t vec1Loop = 0;
|
||||
uint32_t vec2Loop = 0;
|
||||
bool isFirstUpdateCurGroup = true;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *workspace)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
constInfo.aiCoreIdx = GetBlockIdx() / 2;
|
||||
} else {
|
||||
constInfo.aiCoreIdx = GetBlockIdx();
|
||||
}
|
||||
InitTilingData();
|
||||
// init tools
|
||||
tools_.toolParams_.seqSize = tilingData_->baseParams.seqSize;
|
||||
tools_.toolParams_.cmpRatio = tilingData_->baseParams.cmpRatio;
|
||||
tools_.Init(startPos, seqUsed, cuSeqlens);
|
||||
|
||||
// 剔除尾部的无效batch
|
||||
for (; constInfo.batchSize > 0; --constInfo.batchSize) {
|
||||
uint32_t bSeqUsed = tools_.GetSeqLength(constInfo.batchSize - 1);
|
||||
if (bSeqUsed > 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 所有batch的有效序列都为0时, 直接退出
|
||||
if (constInfo.batchSize == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// 0. 计算最后一个Tc块的起始位置
|
||||
constInfo.bIdxOfLastTc = constInfo.batchSize - 1;
|
||||
// 1. 计算head_dim的切分大小, 构建ConstInfo的其他信息
|
||||
SetBaseSize(); // 设置基本块大小
|
||||
CalcSplitCoreInfo();
|
||||
// 2. 计算循环次数
|
||||
loopTimes = GetLoopTimes();
|
||||
// 3. 初始化workspace
|
||||
InitWorkspace(workspace);
|
||||
// 4. 初始化block层
|
||||
if ASCEND_IS_AIC {
|
||||
#if __CCE_AICORE__ == 310
|
||||
blockCube_.InitParams(constInfo, tools_);
|
||||
#else
|
||||
blockCube_.InitParams(constInfo, tools_);
|
||||
#endif
|
||||
blockCube_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos,
|
||||
stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut);
|
||||
blockCube_.InitBuffers(pipe_);
|
||||
#if __CCE_AICORE__ == 310
|
||||
blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);
|
||||
#else
|
||||
blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);
|
||||
#endif
|
||||
} else {
|
||||
blockVec_.InitParams(constInfo, tools_);
|
||||
blockVec_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable,
|
||||
cuSeqlens, seqUsed, startPos, cmpKvOut);
|
||||
blockVec_.InitBuffers(pipe_);
|
||||
#if __CCE_AICORE__ == 310
|
||||
blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm);
|
||||
#else
|
||||
blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm);
|
||||
#endif
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::InitTilingData() {
|
||||
constInfo.cmpRatio = tilingData_->baseParams.cmpRatio;
|
||||
constInfo.batchSize = tilingData_->baseParams.batchSize;
|
||||
constInfo.mBaseSize = tilingData_->innerSplitParams.mBaseSize;
|
||||
constInfo.headDim = tilingData_->baseParams.headDim;
|
||||
constInfo.hSize = tilingData_->baseParams.hiddenSize;
|
||||
constInfo.sSize = tilingData_->baseParams.seqSize;
|
||||
constInfo.ropeHeadDim = tilingData_->baseParams.ropeHeadDim;
|
||||
constInfo.normEps = tilingData_->baseParams.normEps;
|
||||
constInfo.reciprocalD = tilingData_->baseParams.reciprocalD;
|
||||
constInfo.usedCoreNum = tilingData_->baseParams.usedCoreNum;
|
||||
|
||||
constInfo.blockNum = tilingData_->pageAttentionParams.blockNum;
|
||||
constInfo.blockSize = tilingData_->pageAttentionParams.blockSize;
|
||||
constInfo.maxBlockNumPerBatch = tilingData_->pageAttentionParams.maxBlockNumPerBatch;
|
||||
constInfo.stateCacheStrideDim0 = tilingData_->baseParams.stateCacheStrideDim0;
|
||||
|
||||
constInfo.nSize = tilingData_->baseParams.nSize;
|
||||
constInfo.vec1TailCacheSize = tilingData_->workspaceParams.vec1TailCacheSize;
|
||||
constInfo.dbWorkspaceRatio = tilingData_->workspaceParams.dbWorkspaceRatio;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::SetBaseSize()
|
||||
{
|
||||
uint32_t mSize = 0;
|
||||
uint32_t minMBaseSize = 0;
|
||||
bool sameSeqUsed = true;
|
||||
uint32_t firstBatchSeqUsed = tools_.GetSeqLength(0);
|
||||
for (uint32_t i = 0; i < constInfo.batchSize; i++) {
|
||||
uint32_t bSeqUsed = tools_.GetSeqLength(i);
|
||||
uint32_t bStartPos = tools_.GetStartPos(i);
|
||||
// 获取m大小
|
||||
mSize += bSeqUsed;
|
||||
// 获取是否等长
|
||||
if (sameSeqUsed && (bSeqUsed != firstBatchSeqUsed)) {
|
||||
sameSeqUsed = false;
|
||||
}
|
||||
// 获取m轴最小切分大小
|
||||
if (minMBaseSize != constInfo.cmpRatio) {
|
||||
uint32_t startCmpIdx = bStartPos / constInfo.cmpRatio;
|
||||
uint32_t endCmpIdx = (bStartPos + bSeqUsed) / constInfo.cmpRatio;
|
||||
if (startCmpIdx == endCmpIdx) {
|
||||
if (bSeqUsed > minMBaseSize) {
|
||||
minMBaseSize = bSeqUsed;
|
||||
}
|
||||
} else if (startCmpIdx + 1 == endCmpIdx) {
|
||||
uint32_t startCmpValidSeqCnt = constInfo.cmpRatio - (bStartPos % constInfo.cmpRatio);
|
||||
uint32_t endCmpValidSeqCnt = (bStartPos + bSeqUsed) % constInfo.cmpRatio;
|
||||
if (startCmpValidSeqCnt > minMBaseSize) {
|
||||
minMBaseSize = startCmpValidSeqCnt;
|
||||
}
|
||||
if (endCmpValidSeqCnt > minMBaseSize) {
|
||||
minMBaseSize = endCmpValidSeqCnt;
|
||||
}
|
||||
} else {
|
||||
minMBaseSize = constInfo.cmpRatio;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t aiCoreNum = constInfo.usedCoreNum;
|
||||
constInfo.dBaseSize = 64;
|
||||
uint32_t dBaseBlockNum = constInfo.headDim / constInfo.dBaseSize;
|
||||
if (sameSeqUsed && mSize <= (constInfo.mBaseSize * (aiCoreNum / dBaseBlockNum))) {
|
||||
if constexpr (COMP::coff == COFF::OVERLAP) {
|
||||
if (constInfo.headDim == 128) {
|
||||
dBaseBlockNum = 8;
|
||||
} else if (constInfo.headDim == 512) {
|
||||
dBaseBlockNum = 16;
|
||||
}
|
||||
} else {
|
||||
if (constInfo.headDim == 128) {
|
||||
dBaseBlockNum = 8;
|
||||
} else if (constInfo.headDim == 512) {
|
||||
dBaseBlockNum = 16;
|
||||
}
|
||||
}
|
||||
// 核数足够时, 修改才生效
|
||||
if (aiCoreNum >= dBaseBlockNum) {
|
||||
constInfo.dBaseSize = constInfo.headDim / dBaseBlockNum;
|
||||
// 开启全核
|
||||
uint32_t coreGroupNum = aiCoreNum / dBaseBlockNum;
|
||||
uint32_t newMBaseSize = (constInfo.batchSize + coreGroupNum - 1) / coreGroupNum * firstBatchSeqUsed;
|
||||
if (newMBaseSize > minMBaseSize && newMBaseSize < constInfo.mBaseSize) {
|
||||
constInfo.mBaseSize = newMBaseSize;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::SkipInvalidBatch(BatchInfo &batchInfo)
|
||||
{
|
||||
for (; batchInfo.bIdx < constInfo.batchSize; ++batchInfo.bIdx) {
|
||||
batchInfo.seqCnt = tools_.GetSeqLength(batchInfo.bIdx);
|
||||
if (batchInfo.seqCnt > 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
batchInfo.remSeqCnt = batchInfo.seqCnt;
|
||||
if (tools_.isExistSeqUsed_) {
|
||||
batchInfo.seqUsedCnt = tools_.GetSeqUsed(batchInfo.bIdx);
|
||||
} else {
|
||||
batchInfo.seqUsedCnt = batchInfo.seqCnt;
|
||||
}
|
||||
if (batchInfo.bIdx < constInfo.batchSize) {
|
||||
batchInfo.bStartPos = tools_.GetStartPos(batchInfo.bIdx);
|
||||
batchInfo.sIdx = 0;
|
||||
batchInfo.headHolderSeq = batchInfo.bStartPos & (constInfo.cmpRatio - 1);
|
||||
batchInfo.tcNum = (batchInfo.bStartPos + batchInfo.seqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio;
|
||||
batchInfo.compressedTcNum = (batchInfo.bStartPos + batchInfo.seqUsedCnt) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::UpdateCurGroup(BasicBlockInfo &basicBlockInfo,
|
||||
BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq)
|
||||
{
|
||||
// 更新当前组的信息
|
||||
if (curGroupQuota == 0 && !isFirstUpdateCurGroup) {
|
||||
return;
|
||||
}
|
||||
isFirstUpdateCurGroup = false;
|
||||
basicBlockInfo.bIdx = batchInfo.bIdx;
|
||||
uint32_t curGroupDealSeq = curGroupQuota < curDealSeq ? curGroupQuota : curDealSeq;
|
||||
basicBlockInfo.sIdx = batchInfo.sIdx + curGroupDealSeq;
|
||||
basicBlockInfo.dealSeqCnt += curGroupDealSeq;
|
||||
curGroupQuota -= curGroupDealSeq;
|
||||
// 结尾需要跳batch,需要考虑在当前组起始为末尾,或者当前组起始大于整个M轴
|
||||
if ((curGroupQuota == 0 || basicBlockInfo.bIdx == constInfo.batchSize - 1) && basicBlockInfo.sIdx == batchInfo.seqCnt) {
|
||||
basicBlockInfo.sIdx = 0;
|
||||
for (basicBlockInfo.bIdx++; basicBlockInfo.bIdx < constInfo.batchSize; ++basicBlockInfo.bIdx) {
|
||||
uint32_t seqCnt = tools_.GetSeqLength(basicBlockInfo.bIdx);
|
||||
if (seqCnt > 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline BasicBlockInfo CompressorKernelPerf<COMP>::SkipOneLoop(BatchInfo &batchInfo)
|
||||
{
|
||||
BasicBlockInfo basicBlockInfo{};
|
||||
isFirstUpdateCurGroup = true;
|
||||
uint32_t curGroupQuota = constInfo.mBaseSize * constInfo.curGroupIdx; // m轴当前组起始
|
||||
bool curGroupStartFlag = false;
|
||||
uint32_t quota = constInfo.coreGroupNum * constInfo.mBaseSize;
|
||||
|
||||
for (; batchInfo.bIdx < constInfo.batchSize;) {
|
||||
uint32_t curDealSeq = 0;
|
||||
uint32_t curDealTcNum = 0;
|
||||
uint32_t curDealCompressedTcNum = 0;
|
||||
// 无法处理完当前整个batch
|
||||
if (quota < batchInfo.remSeqCnt) {
|
||||
// 向下对齐r,
|
||||
if (quota > constInfo.cmpRatio - batchInfo.headHolderSeq) {
|
||||
uint32_t delta = (batchInfo.bStartPos + batchInfo.sIdx + quota) & (constInfo.cmpRatio - 1); // 超出对齐的部分
|
||||
curDealSeq = quota - delta;
|
||||
quota -= curDealSeq;
|
||||
curDealTcNum = (curDealSeq + constInfo.cmpRatio - 1) / constInfo.cmpRatio;
|
||||
curDealCompressedTcNum = min(curDealTcNum, batchInfo.compressedTcNum);
|
||||
// 更新当前组所需信息
|
||||
UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq);
|
||||
// 更新batch信息
|
||||
batchInfo.remSeqCnt = batchInfo.remSeqCnt - curDealSeq;
|
||||
batchInfo.sIdx = batchInfo.sIdx + curDealSeq;
|
||||
batchInfo.compressedTcNum -= curDealCompressedTcNum;
|
||||
batchInfo.tcNum -= curDealTcNum;
|
||||
// 更新loop信息
|
||||
basicBlockInfo.dealTcNum += curDealTcNum;
|
||||
basicBlockInfo.compressedTcNum += curDealCompressedTcNum;
|
||||
}
|
||||
break;
|
||||
} else {
|
||||
// 处理整个batch
|
||||
quota -= batchInfo.remSeqCnt;
|
||||
curDealSeq = batchInfo.remSeqCnt;
|
||||
curDealTcNum = batchInfo.tcNum;
|
||||
// 更新当前组所需信息
|
||||
UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq);
|
||||
// 更新batch和loop信息
|
||||
batchInfo.remSeqCnt = 0;
|
||||
basicBlockInfo.dealTcNum += batchInfo.tcNum;
|
||||
basicBlockInfo.compressedTcNum += batchInfo.compressedTcNum;
|
||||
batchInfo.bIdx++;
|
||||
SkipInvalidBatch(batchInfo);
|
||||
}
|
||||
}
|
||||
uint32_t totalDataSize = constInfo.coreGroupNum * constInfo.mBaseSize - quota;
|
||||
// 2. 当前组的起始偏移
|
||||
uint32_t currentGroupStart = constInfo.curGroupIdx * constInfo.mBaseSize;
|
||||
|
||||
// 3. 安全判断
|
||||
if (currentGroupStart >= totalDataSize) {
|
||||
// 超出尾块
|
||||
basicBlockInfo.dealSeqCnt = 0;
|
||||
} else {
|
||||
// 还在有效范围内,计算剩余量
|
||||
uint32_t remaining = totalDataSize - currentGroupStart;
|
||||
basicBlockInfo.dealSeqCnt = (remaining < constInfo.mBaseSize) ? remaining : constInfo.mBaseSize;
|
||||
}
|
||||
|
||||
return basicBlockInfo;
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorKernelPerf<COMP>::GetLoopTimes()
|
||||
{
|
||||
// 计算主循环次数
|
||||
uint32_t loopTimes = 0;
|
||||
BatchInfo batchInfo{};
|
||||
SkipInvalidBatch(batchInfo);
|
||||
for (;batchInfo.bIdx < constInfo.batchSize; ++loopTimes) {
|
||||
SkipOneLoop(batchInfo);
|
||||
}
|
||||
return loopTimes;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::CalcSplitCoreInfo()
|
||||
{
|
||||
// D方向的基本块数量
|
||||
constInfo.dBasicBlockNum = constInfo.headDim / constInfo.dBaseSize;
|
||||
// 核的组数
|
||||
constInfo.coreGroupNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum;
|
||||
// 每个核处理的d方向的索引
|
||||
constInfo.dIdx = (constInfo.aiCoreIdx % constInfo.dBasicBlockNum) * constInfo.dBaseSize;
|
||||
// 当前组id
|
||||
constInfo.curGroupIdx = constInfo.aiCoreIdx / constInfo.dBasicBlockNum;
|
||||
|
||||
constInfo.mm1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.coreGroupNum;
|
||||
|
||||
uint32_t coff = (uint32_t)COMP::coff;
|
||||
constInfo.mm1KvResSize = constInfo.mBaseSize * constInfo.headDim * coff;
|
||||
constInfo.mm1ScoreResSize = constInfo.mBaseSize * constInfo.headDim * coff;
|
||||
constInfo.vec1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.nSize;
|
||||
|
||||
constInfo.dbSize = constInfo.coreGroupNum * constInfo.mm1KvResSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::InitWorkspace(__gm__ uint8_t *workspace) {
|
||||
uint64_t offset = 0;
|
||||
uint64_t mm1KvResStartOffset = offset;
|
||||
// mm1KvResGm
|
||||
mm1KvResGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + offset +
|
||||
constInfo.curGroupIdx * constInfo.mm1KvResSize * sizeof(MM1_OUT_T)));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1KvResSize * sizeof(MM1_OUT_T);
|
||||
|
||||
uint64_t mm1ScoreResStartOffset = offset;
|
||||
// mm1ScoreResGm
|
||||
mm1ScoreResGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + offset +
|
||||
constInfo.curGroupIdx * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T)));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T);
|
||||
|
||||
Vec1InputKvGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + mm1KvResStartOffset));
|
||||
|
||||
Vec1InputScoreGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + mm1ScoreResStartOffset));
|
||||
|
||||
vec1KvCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T);
|
||||
|
||||
vec1ScoreCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T);
|
||||
|
||||
uint64_t beforeVecOffset = offset;
|
||||
|
||||
// vec1Res
|
||||
vec1ResGm.SetGlobalBuffer(
|
||||
(__gm__ VEC1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.vec1ResSize * sizeof(VEC1_OUT_T);
|
||||
// vec2Input
|
||||
vec2InputGm.SetGlobalBuffer(
|
||||
(__gm__ VEC1_OUT_T *)(workspace + beforeVecOffset));
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::ComputeMm1(const RunInfo &info, bool isNeedExcute) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_FIX>(SYNC_V1_C1_FLAG + info.cubeDbIdx);
|
||||
if (isNeedExcute) {
|
||||
blockCube_.ComputeMm1(info);
|
||||
}
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_FIX>(SYNC_C1_FLAG);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_FIX>(SYNC_C1_FLAG);
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_FIX>(SYNC_C1_V1_FLAG + info.cubeDbIdx);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::ComputeVec1(const Vec1RunInfo &info) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_C1_V1_FLAG + info.c1v1DbIdx);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG2 + info.c1v1DbIdx);
|
||||
blockVec_.ComputeVec1(info);
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG);
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_V1_C1_FLAG + info.c1v1DbIdx);
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE3>(SYNC_V1_FLAG2 + (info.c1v1DbIdx + 1) % constInfo.dbWorkspaceRatio);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::ComputeVec2(const Vec2RunInfo &info) {
|
||||
blockVec_.ComputeVec2(info);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::AllocEventID()
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
blockCube_.AllocEventID(pipe_);
|
||||
} else {
|
||||
blockVec_.AllocEventID();
|
||||
for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) {
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_V1_C1_FLAG + i);
|
||||
}
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE3>(SYNC_V1_FLAG2);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::FreeEventID()
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_FIX>(SYNC_V1_C1_FLAG + i);
|
||||
}
|
||||
blockCube_.FreeEventID(pipe_);
|
||||
} else {
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG2 + loopTimes % constInfo.dbWorkspaceRatio);
|
||||
blockVec_.FreeEventID();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernelPerf<COMP>::IsNeedExcuteC1(RunInfo info)
|
||||
{
|
||||
// B超出范围则cube不执行
|
||||
return info.bStart < constInfo.batchSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx)
|
||||
{
|
||||
vec1Info.bStart = batchInfo.bIdx;
|
||||
vec1Info.sStart = batchInfo.sIdx;
|
||||
vec1Info.resetResFlag = (loopIdx & (constInfo.nSize - 1)) == 0;
|
||||
vec1Info.c1v1DbIdx = (vec1Loop++ & (constInfo.dbWorkspaceRatio - 1));
|
||||
vec1Info.v1v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1));
|
||||
BasicBlockInfo basicBlockInfo = SkipOneLoop(batchInfo);
|
||||
info.cubeDbIdx = (cubeLoop++ & (constInfo.dbWorkspaceRatio - 1));
|
||||
info.dealSeqCnt = basicBlockInfo.dealSeqCnt;
|
||||
info.dealTcNum = basicBlockInfo.dealTcNum;
|
||||
info.bStart = basicBlockInfo.bIdx;
|
||||
info.sStart = basicBlockInfo.sIdx;
|
||||
vec1Info.dealTcNum = basicBlockInfo.dealTcNum;
|
||||
vec1Info.dealScSize = basicBlockInfo.compressedTcNum;
|
||||
allCompressedTcNum_ += basicBlockInfo.compressedTcNum;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernelPerf<COMP>::IsNeedExcuteV2(Vec2RunInfo &vec2Info)
|
||||
{
|
||||
return (vec2Info.dealScSize > 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernelPerf<COMP>::IsNeedSyncAll(uint32_t curBasicBlockIdx)
|
||||
{
|
||||
if (allCompressedTcNum_ == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint32_t cnt = curBasicBlockIdx + 1;
|
||||
if ((cnt == loopTimes) || (cnt % constInfo.nSize == 0)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::UpdateVec2Info(
|
||||
Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info)
|
||||
{
|
||||
// nSize轮起始先重置v2Info信息
|
||||
if (curBasicBlockIdx % constInfo.nSize == 0) {
|
||||
vec2Info.v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1));
|
||||
vec2Info.bStart = info.bStart;
|
||||
vec2Info.sStart = info.sStart;
|
||||
// 将sStart转成bCompressedId
|
||||
uint32_t startPos = tools_.GetStartPos(info.bStart);
|
||||
if (tools_.isExistSeqUsed_) {
|
||||
uint32_t seqUsed = tools_.GetSeqUsed(info.bStart);
|
||||
if (vec2Info.sStart >= seqUsed) {
|
||||
vec2Info.bStart++;
|
||||
vec2Info.sStart = 0;
|
||||
}
|
||||
}
|
||||
vec2Info.bCompressedId = (startPos + vec2Info.sStart) / constInfo.cmpRatio - startPos / constInfo.cmpRatio;
|
||||
|
||||
vec2Info.dealScSize = 0;
|
||||
} else if ((curBasicBlockIdx + 1) % constInfo.nSize == 0) {
|
||||
vec2Loop++;
|
||||
}
|
||||
vec2Info.dealScSize += info.dealScSize;
|
||||
vec2Info.compressedId += info.dealScSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelPerf<COMP>::Process()
|
||||
{
|
||||
// 所有batch的有效序列都为0时, 直接退出
|
||||
if (constInfo.batchSize == 0) {
|
||||
return;
|
||||
}
|
||||
AllocEventID();
|
||||
|
||||
BatchInfo batchInfo{};
|
||||
|
||||
RunInfo extraInfo[1];
|
||||
Vec1RunInfo vec1Info{};
|
||||
Vec2RunInfo vec2Info{};
|
||||
SkipInvalidBatch(batchInfo);
|
||||
for (uint32_t i = 0; i < loopTimes; ++i) {
|
||||
RunInfo &extraInfo0 = extraInfo[0];
|
||||
CalcC1V1Params(extraInfo0, vec1Info, batchInfo, i);
|
||||
bool isNeedExcuteC1 = IsNeedExcuteC1(extraInfo0);
|
||||
|
||||
if ASCEND_IS_AIC {
|
||||
ComputeMm1(extraInfo0, isNeedExcuteC1);
|
||||
} else {
|
||||
ComputeVec1(vec1Info);
|
||||
UpdateVec2Info(vec2Info, i, vec1Info);
|
||||
|
||||
if (IsNeedSyncAll(i)) {
|
||||
SyncAll();
|
||||
if (IsNeedExcuteV2(vec2Info)) {
|
||||
ComputeVec2(vec2Info);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
FreeEventID();
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif // COMPRESSOR_KERNEL_PERF_H
|
||||
@@ -0,0 +1,56 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file COMPRESSOR_template_tiling_key.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_TEMPLATE_TILING_KEY_H
|
||||
#define COMPRESSOR_TEMPLATE_TILING_KEY_H
|
||||
|
||||
#include "ascendc/host_api/tiling/template_argument.h"
|
||||
|
||||
#define ASCENDC_TPL_1_BW 1 // 每个参数占用1个bit位
|
||||
#define ASCENDC_TPL_2_BW 2 // 每个参数占用2个bit位
|
||||
#define ASCENDC_TPL_4_BW 4 // 每个参数占用4个bit位
|
||||
|
||||
// 可表示的tilingkey范围为64bit,注意不可超过限制
|
||||
ASCENDC_TPL_ARGS_DECL(compressor, // 算子唯一标识,与opType保持一致
|
||||
// 可能需要切分之后的headdim
|
||||
// bit:0 LAYOUT 0:BSH 1:TH
|
||||
ASCENDC_TPL_UINT_DECL(X_LAYOUT, ASCENDC_TPL_1_BW, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
// bit:1-4 x的dtype 0:BF16 1:FP16
|
||||
ASCENDC_TPL_UINT_DECL(X_DTYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
// bit:5-6 coff 1:无需overlap 2:需要overlap
|
||||
ASCENDC_TPL_UINT_DECL(COFF, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
// bit:7-8 rotary_mode 1:half 2:interleave
|
||||
ASCENDC_TPL_UINT_DECL(ROTARY_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
// bit:9-10 cache_mode 1:CONTINUOUS 2:cycle
|
||||
ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
// bit:11-12 template_id 0:empty_tensor 1:normal 2:full load
|
||||
ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2),
|
||||
// bit:13 rope dtype 0:same as x 1:fp32
|
||||
ASCENDC_TPL_UINT_DECL(ROPE_DTYPE, ASCENDC_TPL_1_BW, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
);
|
||||
|
||||
ASCENDC_TPL_SEL(
|
||||
|
||||
ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_UINT_SEL(X_LAYOUT, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
ASCENDC_TPL_UINT_SEL(X_DTYPE, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
ASCENDC_TPL_UINT_SEL(ROTARY_MODE, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2),
|
||||
ASCENDC_TPL_UINT_SEL(ROPE_DTYPE, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)),
|
||||
);
|
||||
|
||||
#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H
|
||||
@@ -0,0 +1,69 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file COMPRESSOR_tiling_datay.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_TILING_DATA_H
|
||||
#define COMPRESSOR_TILING_DATA_H
|
||||
#include <cstdint>
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
|
||||
const uint32_t CMP_MAX_AIC_CORE_NUM = 26; // 25 + 1 保证数组8字节对齐
|
||||
|
||||
namespace optiling {
|
||||
// 1. 基础参数结构体
|
||||
struct CompressorBaseParams {
|
||||
uint32_t batchSize = 0; // bastch size(批大小)
|
||||
uint32_t seqSize = 0; // sequence size(kvs大小)
|
||||
uint32_t hiddenSize = 0; // hidden size(隐藏层大小)
|
||||
uint32_t tokenSize = 0; // token size = batchSize * seqSize(token总数:批大小x序列1长度)
|
||||
uint32_t headDim = 0; // head size of kv
|
||||
uint32_t ropeHeadDim = 64; // dim size per rope head 64(单个带RoPE头的维度)
|
||||
uint32_t csSize = 0; // Compress sequence len
|
||||
uint32_t cmpRatio = 4; // Compress ratio
|
||||
uint32_t cgSize = 0; // Compress group size
|
||||
float normEps = 1e-6; // RMSNorm eps
|
||||
float reciprocalD = 0; // 1分之D
|
||||
uint32_t usedCoreNum = 0; // 使用核数
|
||||
uint32_t nSize = 0; // 控制v2积攒的轮数
|
||||
uint64_t stateCacheStrideDim0 = 0; // stateCache第0维的stride
|
||||
};
|
||||
|
||||
struct CompressorPageAttentionParams {
|
||||
uint32_t blockNum = 0;
|
||||
uint32_t blockSize = 1;
|
||||
uint32_t maxBlockNumPerBatch = 1;
|
||||
};
|
||||
|
||||
struct CompressorInnerSplitParams {
|
||||
uint32_t mBaseSize;
|
||||
uint32_t dBaseSize;
|
||||
};
|
||||
|
||||
struct CompressorWorkspaceParams {
|
||||
uint32_t mm1KvResSize;
|
||||
uint32_t mm1ScoreResSize;
|
||||
uint32_t vec1ResSize;
|
||||
uint32_t vec1TailCacheSize;
|
||||
uint32_t dbWorkspaceRatio = 1;
|
||||
};
|
||||
|
||||
struct CompressorTilingData {
|
||||
CompressorBaseParams baseParams;
|
||||
CompressorPageAttentionParams pageAttentionParams;
|
||||
CompressorInnerSplitParams innerSplitParams;
|
||||
CompressorWorkspaceParams workspaceParams;
|
||||
};
|
||||
} // optiling
|
||||
|
||||
#endif // COMPRESSOR_TILING_DATA_H
|
||||
761
csrc/attention/compressor/op_kernel/arch32/compressor_tools.h
Normal file
761
csrc/attention/compressor/op_kernel/arch32/compressor_tools.h
Normal file
@@ -0,0 +1,761 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_tools.h
|
||||
* \brief 放算子都需要、与算子联系紧密、但是又不方便单独独立出来的公共工具
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_TOOLS_H
|
||||
#define COMPRESSOR_TOOLS_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
struct ToolsParams {
|
||||
uint32_t seqSize = 0U;
|
||||
uint32_t cmpRatio = 0U;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorTools {
|
||||
public:
|
||||
__aicore__ inline CompressorTools()
|
||||
{
|
||||
}
|
||||
|
||||
__aicore__ inline void Init(__gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos);
|
||||
|
||||
__aicore__ inline uint32_t GetSeqUsed(uint32_t bIdx);
|
||||
__aicore__ inline uint32_t GetStartPos(uint32_t bIdx);
|
||||
__aicore__ inline uint32_t GetSeqLength(uint32_t bIdx);
|
||||
__aicore__ inline uint32_t GetTIdxByBatch(uint32_t bIdx);
|
||||
|
||||
public:
|
||||
ToolsParams toolParams_{};
|
||||
bool isExistSeqUsed_ = false;
|
||||
|
||||
private:
|
||||
bool isExistStartPos_ = false;
|
||||
GlobalTensor<int32_t> cuSeqlensGm_;
|
||||
GlobalTensor<int32_t> sequsedGm_;
|
||||
GlobalTensor<int32_t> startPosGm_;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorTools<COMP>::Init(__gm__ uint8_t *startPos, __gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *cuSeqlens)
|
||||
{
|
||||
isExistStartPos_ = (startPos != nullptr);
|
||||
if (isExistStartPos_) {
|
||||
startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos);
|
||||
}
|
||||
|
||||
isExistSeqUsed_ = (seqUsed != nullptr);
|
||||
if (isExistSeqUsed_) {
|
||||
sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed);
|
||||
}
|
||||
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorTools<COMP>::GetSeqUsed(uint32_t bIdx)
|
||||
{
|
||||
if (isExistSeqUsed_) {
|
||||
return (uint32_t)sequsedGm_.GetValue(bIdx);
|
||||
} else {
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
return (uint32_t)(cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx));
|
||||
} else {
|
||||
return toolParams_.seqSize;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorTools<COMP>::GetStartPos(uint32_t bIdx)
|
||||
{
|
||||
if (isExistStartPos_) {
|
||||
return (uint32_t)startPosGm_.GetValue(bIdx);
|
||||
} else {
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorTools<COMP>::GetSeqLength(uint32_t bIdx)
|
||||
{
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
return cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx);
|
||||
} else {
|
||||
return toolParams_.seqSize;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorTools<COMP>::GetTIdxByBatch(uint32_t bIdx)
|
||||
{
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
return (uint32_t)(cuSeqlensGm_.GetValue(bIdx));
|
||||
} else {
|
||||
return toolParams_.seqSize * bIdx;
|
||||
}
|
||||
}
|
||||
|
||||
// iterator
|
||||
struct SliceInfo {
|
||||
__aicore__ inline SliceInfo(){};
|
||||
__aicore__ inline SliceInfo(uint32_t bIdx, uint32_t sIdx) : bIdx(bIdx), sIdx(sIdx){};
|
||||
|
||||
uint32_t bIdx = 0U;
|
||||
uint32_t sIdx = 0U;
|
||||
uint32_t bSeqUsed = 0U;
|
||||
uint32_t bStartPos = 0U;
|
||||
|
||||
uint32_t headHolderSeqCnt = 0U;
|
||||
uint32_t validSeqCnt = 0U;
|
||||
uint32_t tailHolderSeqCnt = 0U;
|
||||
|
||||
uint32_t dealSeqCnt = 0;
|
||||
uint32_t dealTcSize = 0U;
|
||||
uint32_t compressTcSize = 0U;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorSliceIterator {
|
||||
public:
|
||||
__aicore__ inline CompressorSliceIterator(CompressorTools<COMP> &tools) : tools_(tools)
|
||||
{
|
||||
}
|
||||
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
|
||||
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
|
||||
__aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt);
|
||||
__aicore__ inline bool IsEnd();
|
||||
__aicore__ inline void IteratorSlice();
|
||||
__aicore__ inline SliceInfo &GetSlice();
|
||||
__aicore__ inline SliceInfo &GetSliceByCmp();
|
||||
|
||||
bool isFirst_ = true;
|
||||
SliceInfo sliceInfo_{};
|
||||
|
||||
private:
|
||||
CompressorTools<COMP> &tools_;
|
||||
|
||||
// iterator
|
||||
uint32_t maxDealSeqCnt_ = 0;
|
||||
uint32_t batch_size_ = 0;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
|
||||
{
|
||||
sliceInfo_.bIdx = bIdx;
|
||||
sliceInfo_.sIdx = sIdx;
|
||||
isFirst_ = true;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
|
||||
{
|
||||
this->batch_size_ = batch_size;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSliceIterator<COMP>::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt)
|
||||
{
|
||||
this->maxDealSeqCnt_ = maxDealSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorSliceIterator<COMP>::IsEnd()
|
||||
{
|
||||
return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSliceIterator<COMP>::IteratorSlice()
|
||||
{
|
||||
bool isUpdateBatchInfo = false;
|
||||
if (!isFirst_) {
|
||||
// 更新剩余未处理的行数
|
||||
maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt;
|
||||
// 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed
|
||||
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) {
|
||||
sliceInfo_.sIdx = 0;
|
||||
sliceInfo_.bIdx++;
|
||||
isUpdateBatchInfo = true;
|
||||
}
|
||||
} else {
|
||||
isUpdateBatchInfo = true;
|
||||
isFirst_ = false;
|
||||
}
|
||||
|
||||
// 更新与bIdx相关的bStartPos和bSeqUsed
|
||||
if (isUpdateBatchInfo) {
|
||||
// SkipInvalidBatch
|
||||
while (sliceInfo_.bIdx < batch_size_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
if (sliceInfo_.bSeqUsed > 0) {
|
||||
break;
|
||||
}
|
||||
sliceInfo_.bIdx++;
|
||||
}
|
||||
if (sliceInfo_.bIdx < batch_size_) {
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline SliceInfo &CompressorSliceIterator<COMP>::GetSliceByCmp()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (isFirst_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
isFirst_ = false;
|
||||
}
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
|
||||
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
sliceInfo_.tailHolderSeqCnt =
|
||||
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
|
||||
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
}
|
||||
|
||||
// 头和尾处理,否则需要处理的seq等于cmpRatio
|
||||
if (sliceInfo_.validSeqCnt < cmpRatio) {
|
||||
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == 0) {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
} else {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio;
|
||||
}
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
|
||||
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
|
||||
|
||||
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
|
||||
if (sliceInfo_.tailHolderSeqCnt > 0) {
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时,其不需要压缩
|
||||
}
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline SliceInfo &CompressorSliceIterator<COMP>::GetSlice()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (isFirst_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
isFirst_ = false;
|
||||
}
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
|
||||
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
sliceInfo_.tailHolderSeqCnt =
|
||||
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
|
||||
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
}
|
||||
|
||||
sliceInfo_.dealSeqCnt = sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt;
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize = sliceInfo_.dealSeqCnt / cmpRatio;
|
||||
|
||||
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
|
||||
if (sliceInfo_.tailHolderSeqCnt > 0) {
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时,其不需要压缩
|
||||
}
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
struct SplitCoreSliceInfo : public SliceInfo {
|
||||
__aicore__ inline SplitCoreSliceInfo(){};
|
||||
__aicore__ inline SplitCoreSliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx){};
|
||||
|
||||
uint32_t preFirstSeqCnt = 0U; // 左边每次迭代基本块的第一个seqCnt大小
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorSplitCoreSliceIterator {
|
||||
public:
|
||||
__aicore__ inline CompressorSplitCoreSliceIterator(CompressorTools<COMP> &tools) : tools_(tools)
|
||||
{
|
||||
}
|
||||
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
|
||||
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
|
||||
__aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt);
|
||||
__aicore__ inline bool IsEnd();
|
||||
__aicore__ inline void IteratorSlice();
|
||||
__aicore__ inline SplitCoreSliceInfo &GetSlice();
|
||||
__aicore__ inline SplitCoreSliceInfo &GetSliceByCmp();
|
||||
__aicore__ inline uint32_t GetBIdx();
|
||||
__aicore__ inline SplitCoreSliceInfo &GetLeftNextCmpSeqCnt();
|
||||
__aicore__ inline SplitCoreSliceInfo &GetRightNextCmpSeqCnt();
|
||||
|
||||
bool isFirst_ = true;
|
||||
bool isLeftFirstBath = false;
|
||||
bool isMaxDealSeqCntFirst = false;
|
||||
|
||||
SplitCoreSliceInfo sliceInfo_{};
|
||||
|
||||
private:
|
||||
CompressorTools<COMP> &tools_;
|
||||
|
||||
// iterator
|
||||
uint32_t maxDealSeqCnt_ = 0;
|
||||
uint32_t batch_size_ = 0;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
|
||||
{
|
||||
sliceInfo_.bIdx = bIdx;
|
||||
sliceInfo_.sIdx = sIdx;
|
||||
isFirst_ = true;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
|
||||
{
|
||||
this->batch_size_ = batch_size;
|
||||
isMaxDealSeqCntFirst = true;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt)
|
||||
{
|
||||
this->maxDealSeqCnt_ = maxDealSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorSplitCoreSliceIterator<COMP>::IsEnd()
|
||||
{
|
||||
return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorSplitCoreSliceIterator<COMP>::GetBIdx()
|
||||
{
|
||||
return sliceInfo_.bIdx;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::IteratorSlice()
|
||||
{
|
||||
bool isUpdateBatchInfo = false;
|
||||
if (isMaxDealSeqCntFirst) {
|
||||
isMaxDealSeqCntFirst = false;
|
||||
}
|
||||
if (!isFirst_) {
|
||||
// 更新剩余未处理的行数
|
||||
maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt;
|
||||
// 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed
|
||||
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) {
|
||||
sliceInfo_.sIdx = 0;
|
||||
// 左边最后一块跳到b=0 s=0处理
|
||||
if (isLeftFirstBath) {
|
||||
isLeftFirstBath = false;
|
||||
} else {
|
||||
sliceInfo_.bIdx++;
|
||||
}
|
||||
isUpdateBatchInfo = true;
|
||||
}
|
||||
} else {
|
||||
isUpdateBatchInfo = true;
|
||||
isFirst_ = false;
|
||||
}
|
||||
|
||||
// 更新与bIdx相关的bStartPos和bSeqUsed
|
||||
if (isUpdateBatchInfo) {
|
||||
// SkipInvalidBatch
|
||||
while (sliceInfo_.bIdx < batch_size_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
if (sliceInfo_.bSeqUsed > 0) {
|
||||
break;
|
||||
}
|
||||
sliceInfo_.bIdx++;
|
||||
}
|
||||
if (sliceInfo_.bIdx < batch_size_) {
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline SplitCoreSliceInfo &CompressorSplitCoreSliceIterator<COMP>::GetLeftNextCmpSeqCnt()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (isFirst_) {
|
||||
// 左边 T轴首次减去T轴最后一块
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(batch_size_ - 1);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(batch_size_ - 1);
|
||||
// 处理最后一块是中间整块或者尾块的情况
|
||||
uint32_t lastSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio == 0 ?
|
||||
cmpRatio :
|
||||
(sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio;
|
||||
// 处理最后一块是头块的情况
|
||||
if (sliceInfo_.bSeqUsed < cmpRatio) {
|
||||
lastSeqCnt = sliceInfo_.bSeqUsed;
|
||||
}
|
||||
|
||||
sliceInfo_.sIdx = sliceInfo_.bSeqUsed - lastSeqCnt;
|
||||
isLeftFirstBath = true;
|
||||
isFirst_ = false;
|
||||
}
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
|
||||
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
sliceInfo_.tailHolderSeqCnt =
|
||||
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
|
||||
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
}
|
||||
|
||||
// 头和尾处理,否则需要处理的seq等于cmpRatio
|
||||
if (sliceInfo_.validSeqCnt < cmpRatio) {
|
||||
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == 0) {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
} else {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio;
|
||||
}
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
|
||||
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
|
||||
|
||||
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
|
||||
if (sliceInfo_.tailHolderSeqCnt > 0) {
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时,其不需要压缩
|
||||
}
|
||||
|
||||
// 记录左边第一个块
|
||||
if (isMaxDealSeqCntFirst) {
|
||||
sliceInfo_.preFirstSeqCnt = sliceInfo_.dealSeqCnt;
|
||||
}
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline SplitCoreSliceInfo &CompressorSplitCoreSliceIterator<COMP>::GetRightNextCmpSeqCnt()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (isFirst_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
isFirst_ = false;
|
||||
}
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
|
||||
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
sliceInfo_.tailHolderSeqCnt =
|
||||
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
|
||||
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
}
|
||||
|
||||
// 头和尾处理,否则需要处理的seq等于cmpRatio
|
||||
if (sliceInfo_.validSeqCnt < cmpRatio) {
|
||||
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == 0) {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
} else {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio;
|
||||
}
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
|
||||
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
|
||||
|
||||
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
|
||||
if (sliceInfo_.tailHolderSeqCnt > 0) {
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时,其不需要压缩
|
||||
}
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
struct Vec1SliceInfo : public SliceInfo {
|
||||
__aicore__ inline Vec1SliceInfo(){};
|
||||
__aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx){};
|
||||
__aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt)
|
||||
: SliceInfo(bIdx, sIdx), dealedSeqCnt(dealedSeqCnt){};
|
||||
|
||||
uint32_t dealedSeqCnt = 0U;
|
||||
uint32_t dealedTcCnt = 0U;
|
||||
uint32_t bSeqLength = 0U;
|
||||
uint32_t compressoredScCnt = 0U;
|
||||
bool isFirst = false;
|
||||
bool isLast = false;
|
||||
};
|
||||
|
||||
struct StatisticInfo {
|
||||
__aicore__ inline StatisticInfo(){};
|
||||
__aicore__ inline StatisticInfo(uint32_t actualTcCnt, uint32_t dealSeqCnt, uint32_t compressorScCnt)
|
||||
: actualTcCnt(actualTcCnt), dealSeqCnt(dealSeqCnt), compressorScCnt(compressorScCnt){};
|
||||
|
||||
uint32_t actualTcCnt = 0U;
|
||||
uint32_t dealSeqCnt = 0U;
|
||||
uint32_t compressorScCnt = 0U;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorVec1SliceIterator {
|
||||
public:
|
||||
__aicore__ inline CompressorVec1SliceIterator(CompressorTools<COMP> &tools) : tools_(tools)
|
||||
{
|
||||
}
|
||||
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt, uint32_t compressoredScCnt);
|
||||
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
|
||||
__aicore__ inline void SetDealedSeqCnt(uint32_t dealedSeqCnt);
|
||||
__aicore__ inline void SetDealedTcCnt(uint32_t dealedTcCnt);
|
||||
__aicore__ inline void SetCompressoredScCnt(uint32_t compressoredScCnt);
|
||||
__aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize);
|
||||
__aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize, uint32_t canDealTcSize);
|
||||
__aicore__ inline uint32_t GetNeedDealTcSize();
|
||||
__aicore__ inline bool IsEnd();
|
||||
template <bool IS_STATISTIC = false>
|
||||
__aicore__ inline void IteratorSlice();
|
||||
__aicore__ inline Vec1SliceInfo &GetSlice();
|
||||
template <bool IS_STATISTIC = false>
|
||||
__aicore__ inline StatisticInfo &FullIteratorSlice();
|
||||
|
||||
private:
|
||||
CompressorTools<COMP> &tools_;
|
||||
|
||||
bool isFirst_ = true;
|
||||
Vec1SliceInfo sliceInfo_{};
|
||||
StatisticInfo statisticInfo_{};
|
||||
uint32_t needDealTcSize_ = 0U;
|
||||
uint32_t batch_size_ = 0U;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
|
||||
{
|
||||
sliceInfo_.bIdx = bIdx;
|
||||
sliceInfo_.sIdx = sIdx;
|
||||
while (tools_.GetSeqLength(sliceInfo_.bIdx) == 0) {
|
||||
sliceInfo_.bIdx++;
|
||||
if (sliceInfo_.bIdx == batch_size_) {
|
||||
sliceInfo_.bIdx = 0;
|
||||
}
|
||||
}
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
|
||||
isFirst_ = true;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt,
|
||||
uint32_t compressoredScCnt)
|
||||
{
|
||||
Reset(bIdx, sIdx);
|
||||
SetDealedSeqCnt(dealedSeqCnt);
|
||||
SetCompressoredScCnt(compressoredScCnt);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
|
||||
{
|
||||
this->batch_size_ = batch_size;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetDealedSeqCnt(uint32_t dealedSeqCnt)
|
||||
{
|
||||
this->sliceInfo_.dealedSeqCnt = dealedSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetCompressoredScCnt(uint32_t compressoredScCnt)
|
||||
{
|
||||
this->sliceInfo_.compressoredScCnt = compressoredScCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetDealedTcCnt(uint32_t dealedTcCnt)
|
||||
{
|
||||
this->sliceInfo_.dealedTcCnt = dealedTcCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetNeedDealTcSize(uint32_t needDealTcSize)
|
||||
{
|
||||
this->needDealTcSize_ = needDealTcSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
template <bool IS_STATISTIC>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::IteratorSlice()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if constexpr (IS_STATISTIC) {
|
||||
statisticInfo_.actualTcCnt += sliceInfo_.dealTcSize;
|
||||
statisticInfo_.compressorScCnt += sliceInfo_.compressTcSize;
|
||||
}
|
||||
needDealTcSize_ -= sliceInfo_.dealTcSize;
|
||||
sliceInfo_.dealedSeqCnt += sliceInfo_.validSeqCnt;
|
||||
sliceInfo_.compressoredScCnt += sliceInfo_.compressTcSize;
|
||||
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx >= sliceInfo_.bSeqUsed) {
|
||||
do {
|
||||
uint32_t seqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
|
||||
if (sliceInfo_.bSeqUsed < seqLength) {
|
||||
uint32_t nextAlignSIdx = Align(sliceInfo_.bStartPos + sliceInfo_.sIdx, cmpRatio) - sliceInfo_.bStartPos;
|
||||
sliceInfo_.dealedSeqCnt += nextAlignSIdx - sliceInfo_.sIdx;
|
||||
uint32_t tcGap = CeilDivT(static_cast<int32_t>(seqLength - nextAlignSIdx),
|
||||
static_cast<int32_t>(cmpRatio));
|
||||
if (sliceInfo_.bSeqUsed == 0 && nextAlignSIdx > sliceInfo_.sIdx) {
|
||||
// 此时bseqused所在压缩块未被纳入计算
|
||||
tcGap++;
|
||||
}
|
||||
sliceInfo_.sIdx = nextAlignSIdx;
|
||||
if (needDealTcSize_ < tcGap) {
|
||||
sliceInfo_.dealedSeqCnt += needDealTcSize_ * cmpRatio;
|
||||
sliceInfo_.sIdx += needDealTcSize_ * cmpRatio;
|
||||
needDealTcSize_ = 0;
|
||||
break;
|
||||
}
|
||||
sliceInfo_.dealedSeqCnt += seqLength - sliceInfo_.sIdx;
|
||||
needDealTcSize_ -= tcGap;
|
||||
}
|
||||
sliceInfo_.bIdx++;
|
||||
if (sliceInfo_.bIdx == batch_size_) {
|
||||
sliceInfo_.bIdx = 0;
|
||||
}
|
||||
sliceInfo_.sIdx = 0;
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
} while (sliceInfo_.bSeqUsed == 0);
|
||||
sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
}
|
||||
if (isFirst_) {
|
||||
isFirst_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorVec1SliceIterator<COMP>::GetNeedDealTcSize()
|
||||
{
|
||||
return needDealTcSize_;
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorVec1SliceIterator<COMP>::IsEnd()
|
||||
{
|
||||
return (needDealTcSize_ == 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline Vec1SliceInfo &CompressorVec1SliceIterator<COMP>::GetSlice()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (sliceInfo_.bSeqUsed < sliceInfo_.sIdx) {
|
||||
sliceInfo_.headHolderSeqCnt = 0;
|
||||
sliceInfo_.validSeqCnt = 0;
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
sliceInfo_.dealTcSize = 0;
|
||||
sliceInfo_.compressTcSize = 0;
|
||||
} else {
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt, cmpRatio) > needDealTcSize_) {
|
||||
sliceInfo_.validSeqCnt = needDealTcSize_ * cmpRatio - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
uint32_t globalTotalSeqCnt = sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt;
|
||||
sliceInfo_.tailHolderSeqCnt = Align(globalTotalSeqCnt, cmpRatio) - globalTotalSeqCnt;
|
||||
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize =
|
||||
(sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt) / cmpRatio;
|
||||
|
||||
sliceInfo_.compressTcSize =
|
||||
(sliceInfo_.headHolderSeqCnt + min(sliceInfo_.validSeqCnt, sliceInfo_.bSeqUsed - sliceInfo_.sIdx)) /
|
||||
cmpRatio;
|
||||
}
|
||||
|
||||
sliceInfo_.isFirst = isFirst_;
|
||||
sliceInfo_.isLast =
|
||||
sliceInfo_.bSeqUsed > sliceInfo_.sIdx &&
|
||||
CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.bSeqUsed - sliceInfo_.sIdx, cmpRatio) >= needDealTcSize_;
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
template <bool IS_STATISTIC>
|
||||
__aicore__ inline StatisticInfo &CompressorVec1SliceIterator<COMP>::FullIteratorSlice()
|
||||
{
|
||||
if constexpr (IS_STATISTIC) {
|
||||
statisticInfo_ = {0U, 0U, 0U};
|
||||
Vec1SliceInfo tempSliceInfo = GetSlice();
|
||||
while (!IsEnd()) {
|
||||
GetSlice();
|
||||
IteratorSlice<IS_STATISTIC>();
|
||||
}
|
||||
Vec1SliceInfo sliceInfo = GetSlice();
|
||||
statisticInfo_.dealSeqCnt = sliceInfo.dealedSeqCnt - tempSliceInfo.dealedSeqCnt;
|
||||
} else {
|
||||
while (!IsEnd()) {
|
||||
GetSlice();
|
||||
IteratorSlice<IS_STATISTIC>();
|
||||
}
|
||||
}
|
||||
return statisticInfo_;
|
||||
}
|
||||
} // namespace Compressor
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,347 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_vector_comm.h
|
||||
* \brief 存放各种vector的公共组件
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_VECTOR_COMM_H
|
||||
#define COMPRESSOR_VECTOR_COMM_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
namespace Compressor {
|
||||
|
||||
|
||||
struct MatRpeatParam {
|
||||
uint32_t row;
|
||||
uint32_t col;
|
||||
uint32_t dtypeMask;
|
||||
uint32_t loopTimes;
|
||||
uint32_t colRemain;
|
||||
uint8_t repeatStride;
|
||||
};
|
||||
|
||||
struct RmsNormParam {
|
||||
float reciprocal;
|
||||
float epsilon;
|
||||
uint32_t row;
|
||||
uint32_t col;
|
||||
};
|
||||
|
||||
/**
|
||||
* @brief ColumnSum 对矩阵按列进行求和
|
||||
* @param dstLocal 输出tensor [1, col],支持和shareTmpUb是同一块空间
|
||||
* @param srcLocal 输入tensor [row, col]
|
||||
* @param shareTmpUb 临时buffer 内部需要的空间为 [ceil(row / 2) * col * sizeof(float)]
|
||||
* @param row 行数
|
||||
* @param col 列数
|
||||
*/
|
||||
__aicore__ inline void ColumnSum(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
|
||||
const LocalTensor<float> &shareTmpUb, uint32_t row, uint32_t col)
|
||||
{
|
||||
// 行数为1时,直接将srcLocal复制到dstLocal
|
||||
if (unlikely(row == 1)) {
|
||||
DataCopy(dstLocal, srcLocal, row * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
return;
|
||||
}
|
||||
for (uint32_t mask = MAX_R << 1; mask > 1; mask >>= 1) {
|
||||
if (row & mask) {
|
||||
// 将输入对半求和后放进临时空间
|
||||
Add(shareTmpUb, srcLocal, srcLocal[mask * col / 2], mask * col / 2); // 2:对矩阵按列做计算
|
||||
PipeBarrier<PIPE_V>();
|
||||
// 将余量加到前一半上
|
||||
if (unlikely(row > mask)) {
|
||||
if ((row - mask) > (mask >> 1)) {
|
||||
Add(shareTmpUb, shareTmpUb, srcLocal[mask * col], mask * col / 2); // 2:对矩阵按列做计算
|
||||
PipeBarrier<PIPE_V>();
|
||||
Add(shareTmpUb, shareTmpUb, srcLocal[(mask + (mask >> 1)) * col], (row - mask - (mask >> 1)) * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
} else {
|
||||
Add(shareTmpUb, shareTmpUb, srcLocal[mask * col], (row - mask) * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
}
|
||||
// 每次将后一半行加到前一半上
|
||||
for (uint32_t i = mask >> 2; i > 1; i >>= 1) {
|
||||
Add(shareTmpUb, shareTmpUb, shareTmpUb[i * col], i * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
if (mask == 2) { // 2:最后一次矩阵运算处理
|
||||
DataCopy(dstLocal, shareTmpUb, col);
|
||||
} else {
|
||||
Add(dstLocal, shareTmpUb, shareTmpUb[col], col);
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief ColumnMax 对矩阵按列进行求最大值
|
||||
* @param dstLocal 输出tensor [1, col],支持和shareTmpUb是同一块空间
|
||||
* @param srcLocal 输入tensor [row, col]
|
||||
* @param shareTmpUb 临时buffer 内部需要的空间为 [ceil(row / 2) * col * sizeof(float)]
|
||||
* @param row 行数
|
||||
* @param col 列数
|
||||
*/
|
||||
__aicore__ inline void ColumnMax(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
|
||||
const LocalTensor<float> &shareTmpUb, uint32_t row, uint32_t col)
|
||||
{
|
||||
// 行数为1时,直接将srcLocal复制到dstLocal
|
||||
if (unlikely(row == 1)) {
|
||||
DataCopy(dstLocal, srcLocal, row * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
return;
|
||||
}
|
||||
for (uint32_t mask = MAX_R << 1; mask > 1; mask >>= 1) {
|
||||
if (row & mask) {
|
||||
// 将输入对半求最大值后放进临时空间
|
||||
Max(shareTmpUb, srcLocal, srcLocal[mask * col / 2], mask * col / 2); // 2:对矩阵按列做计算
|
||||
PipeBarrier<PIPE_V>();
|
||||
// 将余量和前一半求最大值后加到前一半上
|
||||
if (unlikely(row > mask)) {
|
||||
if ((row - mask) > (mask >> 1)) {
|
||||
Max(shareTmpUb, shareTmpUb, srcLocal[mask * col], mask * col / 2); // 2:对矩阵按列做计算
|
||||
PipeBarrier<PIPE_V>();
|
||||
Max(shareTmpUb, shareTmpUb, srcLocal[(mask + (mask >> 1)) * col], (row - mask - (mask >> 1)) * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
} else {
|
||||
Max(shareTmpUb, shareTmpUb, srcLocal[mask * col], (row - mask) * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
}
|
||||
// 每次将后一半行和前一半最大值后加到前一半上
|
||||
for (uint32_t i = mask >> 2; i > 1; i >>= 1) {
|
||||
Max(shareTmpUb, shareTmpUb, shareTmpUb[i * col], i * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
}
|
||||
if (mask == 2) { // 2:最后一次矩阵运算处理
|
||||
DataCopy(dstLocal, shareTmpUb, col);
|
||||
} else {
|
||||
Max(dstLocal, shareTmpUb, shareTmpUb[col], col);
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* @brief MatSubVec 矩阵逐行减向量
|
||||
* @param dstLocal 输出tensor [row, col]
|
||||
* @param src0Local 输入tensor [row, col]
|
||||
* @param src1Local 输入tensor [1, col]
|
||||
* @param repeatParam 描述待处理数据的排布,包括
|
||||
row 行数
|
||||
col 列数
|
||||
dtypeMask 一次迭代参与计算元素数
|
||||
loopTimes 循环次数
|
||||
colRemain 剩余列数
|
||||
repeatStride 循环步长(内存中实际列长度)
|
||||
*/
|
||||
__aicore__ inline void MatSubVec(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
|
||||
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
|
||||
{
|
||||
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
|
||||
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
|
||||
uint32_t offset = 0;
|
||||
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
|
||||
Sub(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
|
||||
repeatParam.dtypeMask, repeatRowTimes,
|
||||
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
|
||||
offset += repeatParam.dtypeMask;
|
||||
}
|
||||
if (repeatParam.colRemain > 0) {
|
||||
Sub(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
|
||||
repeatParam.colRemain, repeatRowTimes,
|
||||
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief MatDivVec 矩阵逐行除以向量
|
||||
* @param dstLocal 输出tensor [row, col]
|
||||
* @param src0Local 输入tensor [row, col]
|
||||
* @param src1Local 输入tensor [1, col]
|
||||
* @param repeatParam 描述待处理数据的排布,包括
|
||||
row 行数
|
||||
col 列数
|
||||
dtypeMask 一次迭代参与计算元素数
|
||||
loopTimes 循环次数
|
||||
colRemain 剩余列数
|
||||
repeatStride 循环步长(内存中实际列长度)
|
||||
*/
|
||||
__aicore__ inline void MatDivVec(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
|
||||
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
|
||||
{
|
||||
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
|
||||
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
|
||||
uint32_t offset = 0;
|
||||
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
|
||||
Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
|
||||
repeatParam.dtypeMask, repeatRowTimes,
|
||||
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
|
||||
offset += repeatParam.dtypeMask;
|
||||
}
|
||||
if (repeatParam.colRemain > 0) {
|
||||
Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
|
||||
repeatParam.colRemain, repeatRowTimes,
|
||||
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief MatMulVec 矩阵逐行乘以向量
|
||||
* @param dstLocal 输出tensor [row, col]
|
||||
* @param src0Local 输入tensor [row, col]
|
||||
* @param src1Local 输入tensor [1, col]
|
||||
* @param repeatParam 描述待处理数据的排布,包括
|
||||
row 行数
|
||||
col 列数
|
||||
dtypeMask 一次迭代参与计算元素数
|
||||
loopTimes 循环次数
|
||||
colRemain 剩余列数
|
||||
repeatStride 循环步长(内存中实际列长度)
|
||||
*/
|
||||
__aicore__ inline void MatMulVec(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
|
||||
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
|
||||
{
|
||||
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
|
||||
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
|
||||
uint32_t offset = 0;
|
||||
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
|
||||
Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
|
||||
repeatParam.dtypeMask, repeatRowTimes,
|
||||
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
|
||||
offset += repeatParam.dtypeMask;
|
||||
}
|
||||
if (repeatParam.colRemain > 0) {
|
||||
Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local[offset],
|
||||
repeatParam.colRemain, repeatRowTimes,
|
||||
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, 0});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief RowSum 矩阵对每行求和
|
||||
* @param dstLocal 输出tensor [1, row]
|
||||
* @param srcLocal 输入tensor [row, col]
|
||||
* @param shareTmpUb 临时buffer 内部需要的空间为 [row, col],支持和srcLocal是同一块空间
|
||||
* @param repeatParam 描述待处理数据的排布,包括
|
||||
row 行数
|
||||
col 列数
|
||||
dtypeMask 一次迭代参与计算元素数
|
||||
loopTimes 循环次数
|
||||
colRemain 剩余列数
|
||||
repeatStride 循环步长(内存中实际列长度)
|
||||
*/
|
||||
__aicore__ inline void RowSum(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
|
||||
const LocalTensor<float> &shareTmpUb, const MatRpeatParam &repeatParam)
|
||||
{
|
||||
uint32_t blockCount = repeatParam.loopTimes;
|
||||
if (blockCount > 0 && repeatParam.colRemain > 0) {
|
||||
Add(shareTmpUb, srcLocal, srcLocal[blockCount * repeatParam.dtypeMask], repeatParam.colRemain,
|
||||
repeatParam.row,
|
||||
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, repeatParam.repeatStride});
|
||||
AscendC::PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
for (uint32_t loopCount = blockCount >> 1; loopCount > 0; loopCount = blockCount >> 1) {
|
||||
blockCount = (blockCount + 1) >> 1;
|
||||
for (uint32_t i = 0; i < loopCount; i++) {
|
||||
Add(shareTmpUb[i * repeatParam.dtypeMask], srcLocal[i * repeatParam.dtypeMask],
|
||||
srcLocal[(i + blockCount) * repeatParam.dtypeMask], repeatParam.dtypeMask, repeatParam.row,
|
||||
{1, 1, 1, repeatParam.repeatStride, repeatParam.repeatStride, repeatParam.repeatStride});
|
||||
}
|
||||
AscendC::PipeBarrier<PIPE_V>();
|
||||
}
|
||||
|
||||
WholeReduceSum(dstLocal, shareTmpUb,
|
||||
(repeatParam.col < repeatParam.dtypeMask) ? repeatParam.col :
|
||||
repeatParam.dtypeMask,
|
||||
repeatParam.row, 1, 1, repeatParam.repeatStride);
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief RowDivs 矩阵每行除以对应元素
|
||||
* @param dstLocal 输出tensor [row, col]
|
||||
* @param src0Local 输入tensor [row, col]
|
||||
* @param src1Local 输入tensor [row, 1],需要扩展到一个datablock中(实际内存需要为[row, FP32_BLOCK_ELEMENT_NUM])
|
||||
* @param repeatParam 描述待处理数据的排布,包括
|
||||
row 行数
|
||||
col 列数
|
||||
dtypeMask 一次迭代参与计算元素数
|
||||
loopTimes 循环次数
|
||||
colRemain 剩余列数
|
||||
repeatStride 循环步长(内存中实际列长度)
|
||||
*/
|
||||
__aicore__ inline void RowDivs(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
|
||||
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
|
||||
{
|
||||
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
|
||||
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
|
||||
uint32_t offset = 0;
|
||||
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
|
||||
Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local,
|
||||
repeatParam.dtypeMask, repeatRowTimes,
|
||||
{1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1});
|
||||
offset += repeatParam.dtypeMask;
|
||||
}
|
||||
if (repeatParam.colRemain > 0) {
|
||||
Div(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local,
|
||||
repeatParam.colRemain, repeatRowTimes,
|
||||
{1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* @brief RowMuls 矩阵每行乘以相同元素
|
||||
* @param dstLocal 输出tensor [row, col]
|
||||
* @param src0Local 输入tensor [row, col]
|
||||
* @param src1Local 输入tensor [row, 1],需要扩展到一个datablock中(实际内存需要为[row, FP32_BLOCK_ELEMENT_NUM])
|
||||
* @param repeatParam 描述待处理数据的排布,包括
|
||||
row 行数
|
||||
col 列数
|
||||
dtypeMask 一次迭代参与计算元素数
|
||||
loopTimes 循环次数
|
||||
colRemain 剩余列数
|
||||
repeatStride 循环步长(内存中实际列长度)
|
||||
*/
|
||||
__aicore__ inline void RowMuls(const LocalTensor<float> &dstLocal, const LocalTensor<float> &src0Local,
|
||||
const LocalTensor<float> &src1Local, const MatRpeatParam &repeatParam)
|
||||
{
|
||||
for (uint32_t row = 0; row < repeatParam.row; row += REPEAT_MAX_NUM) {
|
||||
uint32_t repeatRowTimes = Std::min(repeatParam.row - row, REPEAT_MAX_NUM);
|
||||
uint32_t offset = 0;
|
||||
for (uint32_t i = 0; i < repeatParam.loopTimes; i++) {
|
||||
Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local,
|
||||
repeatParam.dtypeMask, repeatRowTimes,
|
||||
{1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1});
|
||||
offset += repeatParam.dtypeMask;
|
||||
}
|
||||
if (repeatParam.colRemain > 0) {
|
||||
Mul(dstLocal[row * repeatParam.col + offset], src0Local[row * repeatParam.col + offset], src1Local,
|
||||
repeatParam.colRemain, repeatRowTimes,
|
||||
{1, 1, 0, repeatParam.repeatStride, repeatParam.repeatStride, 1});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
#endif // COMPRESSOR_VECTOR_COMM_H
|
||||
87
csrc/attention/compressor/op_kernel/arch32/rms_norm.h
Normal file
87
csrc/attention/compressor/op_kernel/arch32/rms_norm.h
Normal file
@@ -0,0 +1,87 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
|
||||
/*!
|
||||
* \file rms_norm.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef RMS_NORM_H
|
||||
#define RMS_NORM_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_vector_comm.h"
|
||||
|
||||
namespace Compressor {
|
||||
/**
|
||||
* @brief RmsNorm 对矩阵进行rmsnorm
|
||||
* @param dstLocal 输出tensor [row, col],支持和srcLocal是同一块空间
|
||||
* @param srcLocal 输入tensor [row, col]
|
||||
* @param gammaLocal 系数gamma [1, col]
|
||||
* @param shareTmpUb 临时buffer 内部需要的空间为 [(row * col + row) * sizeof(float)]
|
||||
* @param rmsNormParams rms所需系数,包括
|
||||
reciprocal rmsnorm系数reciprocal
|
||||
epsilon rmsnorm系数epsilon
|
||||
row 处理的行数
|
||||
col 列数
|
||||
*/
|
||||
template <typename GammaType>
|
||||
__aicore__ inline void RmsNorm(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
|
||||
const LocalTensor<GammaType> &gammaLocal, const LocalTensor<float> &shareTmpUb,
|
||||
const RmsNormParam &rmsNormParams)
|
||||
{
|
||||
uint64_t cnt = rmsNormParams.row * rmsNormParams.col;
|
||||
LocalTensor<float> temp1Local = shareTmpUb.ReinterpretCast<float>();
|
||||
LocalTensor<float> temp2Local = temp1Local[cnt];
|
||||
|
||||
// temp1Local = srcLocal ^ 2
|
||||
Mul(temp1Local, srcLocal, srcLocal, cnt);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
MatRpeatParam repeatParams = {
|
||||
rmsNormParams.row, // row
|
||||
rmsNormParams.col, // col
|
||||
FP32_REPEAT_ELEMENT_NUM, // dtypeMask
|
||||
rmsNormParams.col / FP32_REPEAT_ELEMENT_NUM, // loopTimes
|
||||
rmsNormParams.col % FP32_REPEAT_ELEMENT_NUM, // colsRemain
|
||||
static_cast<uint8_t>(rmsNormParams.col / FP32_BLOCK_ELEMENT_NUM), // repeatStride
|
||||
};
|
||||
|
||||
// temp2Local[row] = Sum(temp1Local)
|
||||
RowSum(temp2Local, temp1Local, temp1Local, repeatParams);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
|
||||
// temp2Local[row] = temp2Local[row] * reciprocal(1/N)
|
||||
Muls(temp2Local, temp2Local, rmsNormParams.reciprocal, rmsNormParams.row);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
// temp2Local[row] = temp2Local[row] + epsilon
|
||||
Adds(temp2Local, temp2Local, rmsNormParams.epsilon, rmsNormParams.row);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
// temp2Local[row] = Sqrt(temp2Local[row])
|
||||
Sqrt(temp2Local, temp2Local, rmsNormParams.row);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
// temp1Local[row, 8] = brc(temp2Local[row, 1])
|
||||
Brcb(temp1Local, temp2Local, CeilDivT(rmsNormParams.row, BRCB_NUM), {1, 8});
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
// dstLocal = srcLocal / temp1Local(sum)
|
||||
RowDivs(dstLocal, srcLocal, temp1Local, repeatParams);
|
||||
PipeBarrier<PIPE_V>();
|
||||
|
||||
// dstLocal = dstLocal * gammaLocal
|
||||
MatMulVec(dstLocal, dstLocal, gammaLocal, repeatParams);
|
||||
}
|
||||
} // namespace Compressor
|
||||
#endif // MLA_PROLOG_RMS_NORM_H
|
||||
130
csrc/attention/compressor/op_kernel/arch32/rope.h
Normal file
130
csrc/attention/compressor/op_kernel/arch32/rope.h
Normal file
@@ -0,0 +1,130 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file rope.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef ROPE_H
|
||||
#define ROPE_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_vector_comm.h"
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
/**
|
||||
* @brief SetGatherSrcOffset 计算用于interleave模式的offset
|
||||
* @param gatherOffsetLocal 输出tensor [count],数据类型需要为int64_t,使用时要转换
|
||||
* @param count offset的元素个数,一般为列数
|
||||
*/
|
||||
template <typename T>
|
||||
__aicore__ inline void SetGatherSrcOffset(const LocalTensor<int32_t> &gatherOffsetLocal, uint32_t count)
|
||||
{
|
||||
for (uint32_t i = 0; i < 8; i++) {
|
||||
gatherOffsetLocal.SetValue(i, i ^ 1);
|
||||
}
|
||||
|
||||
event_t eventId_S_V = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_V));
|
||||
SetFlag<HardEvent::S_V>(eventId_S_V);
|
||||
WaitFlag<HardEvent::S_V>(eventId_S_V);
|
||||
|
||||
int32_t scalarValue = 8;
|
||||
while (scalarValue < count) {
|
||||
int32_t nextValue = scalarValue * 2;
|
||||
PipeBarrier<PIPE_V>();
|
||||
if (nextValue < count) {
|
||||
Adds(gatherOffsetLocal[scalarValue], gatherOffsetLocal, scalarValue, scalarValue);
|
||||
} else {
|
||||
Adds(gatherOffsetLocal[scalarValue], gatherOffsetLocal, scalarValue, count - scalarValue);
|
||||
break;
|
||||
}
|
||||
scalarValue = nextValue;
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(gatherOffsetLocal, gatherOffsetLocal, static_cast<int32_t>(sizeof(T)), count);
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* @brief RotaryPosEmb 同时做row行的RotaryPosEmb,每一行的元素为col
|
||||
* @param dstLocal 输出tensor [row, actualCol],支持和srcLocal是同一块空间
|
||||
* @param srcLocal 输入tensor [row, actualCol]
|
||||
* @param cosLocal cos系数tensor [row, col]
|
||||
* @param sinLocal sin系数tensor [row, col]
|
||||
* @param shareTmpUb 临时buffer 内部需要的空间为 [row * col * sizeof(float)]
|
||||
* @param gatherOffsetcastLocal 用于interleave模式的offset,数据类型需要为uint64_t
|
||||
* @param row 待处理的行数
|
||||
* @param col 待处理的列数
|
||||
* @param actualCol 实际列数
|
||||
* @param baseAddr 计算基地址
|
||||
*/
|
||||
template <ROTARY_MODE MODE>
|
||||
__aicore__ inline void RotaryPosEmb(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
|
||||
const LocalTensor<float> &cosLocal, const LocalTensor<float> &sinLocal,
|
||||
const LocalTensor<float> &shareTmpUb,
|
||||
const LocalTensor<uint32_t> &gatherOffsetcastLocal, uint32_t row, uint32_t col,
|
||||
uint32_t actualCol, uint64_t baseAddr)
|
||||
{
|
||||
uint64_t cnt = row * col;
|
||||
uint32_t half_col = col >> 1;
|
||||
uint64_t rsvdCnt = 0;
|
||||
LocalTensor<float> reArrLocal = shareTmpUb.ReinterpretCast<float>();
|
||||
if constexpr (MODE == ROTARY_MODE::HALF) {
|
||||
DataCopy(reArrLocal, srcLocal[baseAddr + half_col],
|
||||
{static_cast<uint16_t>(row), static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM)),
|
||||
static_cast<uint16_t>(CeilDivT(actualCol - half_col, FP32_BLOCK_ELEMENT_NUM)),
|
||||
static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM))});
|
||||
DataCopy(reArrLocal[half_col], srcLocal[baseAddr],
|
||||
{static_cast<uint16_t>(row), static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM)),
|
||||
static_cast<uint16_t>(CeilDivT(actualCol - half_col, FP32_BLOCK_ELEMENT_NUM)),
|
||||
static_cast<uint16_t>(CeilDivT(half_col, FP32_BLOCK_ELEMENT_NUM))});
|
||||
PipeBarrier<PIPE_V>();
|
||||
Muls(reArrLocal, reArrLocal, float(-1), half_col, row,
|
||||
{1, 1, static_cast<uint8_t>(CeilDivT(static_cast<uint32_t>(col), FP32_BLOCK_ELEMENT_NUM)),
|
||||
static_cast<uint8_t>(CeilDivT(static_cast<uint32_t>(col), FP32_BLOCK_ELEMENT_NUM))});
|
||||
} else if constexpr (MODE == ROTARY_MODE::INTERLEAVE) {
|
||||
for (uint32_t i = 0; i < row; i++) {
|
||||
Gather(reArrLocal[i * col], srcLocal[i * actualCol + baseAddr], gatherOffsetcastLocal, 0, col);
|
||||
}
|
||||
PipeBarrier<PIPE_V>();
|
||||
uint32_t repeatTimes = cnt / FP32_REPEAT_ELEMENT_NUM;
|
||||
uint32_t remainder = cnt % FP32_REPEAT_ELEMENT_NUM;
|
||||
uint64_t fullMask = 0x5555555555555555;
|
||||
uint64_t partialMask = 0x55;
|
||||
SetVectorMask<float, MaskMode::NORMAL>(0, fullMask);
|
||||
Muls<float, false>(reArrLocal, reArrLocal, float(-1), MASK_PLACEHOLDER, repeatTimes,
|
||||
{1, 1, FP32_BLOCK_ELEMENT_NUM, FP32_BLOCK_ELEMENT_NUM});
|
||||
|
||||
if (unlikely(remainder > 0)) {
|
||||
SetVectorMask<float, MaskMode::NORMAL>(0, partialMask);
|
||||
Muls<float, false>(reArrLocal[repeatTimes * FP32_REPEAT_ELEMENT_NUM],
|
||||
reArrLocal[repeatTimes * FP32_REPEAT_ELEMENT_NUM], float(-1), MASK_PLACEHOLDER,
|
||||
remainder / FP32_BLOCK_ELEMENT_NUM, {1, 1, 1, 1});
|
||||
}
|
||||
ResetMask();
|
||||
}
|
||||
|
||||
PipeBarrier<PIPE_V>();
|
||||
BinaryRepeatParams computeParams{1,
|
||||
1,
|
||||
1,
|
||||
static_cast<uint8_t>(CeilDivT(actualCol, FP32_BLOCK_ELEMENT_NUM)),
|
||||
static_cast<uint8_t>(CeilDivT(actualCol, FP32_BLOCK_ELEMENT_NUM)),
|
||||
static_cast<uint8_t>(CeilDivT(col, FP32_BLOCK_ELEMENT_NUM))};
|
||||
Mul(dstLocal[baseAddr], srcLocal[baseAddr], cosLocal, col, row, computeParams);
|
||||
Mul(reArrLocal, reArrLocal, sinLocal, cnt);
|
||||
PipeBarrier<PIPE_V>();
|
||||
Add(dstLocal[baseAddr], dstLocal[baseAddr], reArrLocal, col, row, computeParams);
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
52
csrc/attention/compressor/op_kernel/arch32/soft_max.h
Normal file
52
csrc/attention/compressor/op_kernel/arch32/soft_max.h
Normal file
@@ -0,0 +1,52 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
|
||||
/*!
|
||||
* \file soft_max.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef SOFT_MAX_H
|
||||
#define SOFT_MAX_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_vector_comm.h"
|
||||
|
||||
namespace Compressor {
|
||||
/**
|
||||
* @brief ColumnSoftMax 对矩阵按列进行SoftMax
|
||||
* @param dstLocal 输出tensor [row, col],支持和srcLocal是同一块空间
|
||||
* @param srcLocal 输入tensor [row, col]
|
||||
* @param shareTmpUb 临时buffer 内部需要的空间为 [floor(row / 2) * col * sizeof(float)]
|
||||
* @param row 行数
|
||||
* @param col 列数
|
||||
*/
|
||||
__aicore__ inline void ColumnSoftMax(const LocalTensor<float> &dstLocal, const LocalTensor<float> &srcLocal,
|
||||
const LocalTensor<float> &shareTmpUb, uint32_t row, uint32_t col)
|
||||
{
|
||||
uint32_t dtypeMask = FP32_REPEAT_ELEMENT_NUM;
|
||||
uint32_t dLoop = col / dtypeMask;
|
||||
uint32_t dRemain = col % dtypeMask;
|
||||
uint8_t repeatStride = col / FP32_BLOCK_ELEMENT_NUM;
|
||||
ColumnMax(shareTmpUb, srcLocal, shareTmpUb, row, col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
MatSubVec(dstLocal, srcLocal, shareTmpUb, {row, col, dtypeMask, dLoop, dRemain, repeatStride});
|
||||
PipeBarrier<PIPE_V>();
|
||||
Exp(dstLocal, dstLocal, row * col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
ColumnSum(shareTmpUb, dstLocal, shareTmpUb, row, col);
|
||||
PipeBarrier<PIPE_V>();
|
||||
MatDivVec(dstLocal, dstLocal, shareTmpUb, {row, col, dtypeMask, dLoop, dRemain, repeatStride});
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,439 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_block_cube.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_BLOCK_CUBE_H
|
||||
#define COMPRESSOR_BLOCK_CUBE_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_tools.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
template<typename COMP> class CompressorBlockCube {
|
||||
using MM1_OUT_T = float;
|
||||
public:
|
||||
__aicore__ inline CompressorBlockCube(){};
|
||||
__aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools<COMP> &tools);
|
||||
__aicore__ inline void Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut);
|
||||
__aicore__ inline void InitBuffers(TPipe *pipe);
|
||||
__aicore__ inline void InitGlobalBuffers(const GlobalTensor<MM1_OUT_T>& kvMm1ResGm, const GlobalTensor<MM1_OUT_T>& scoreMm1ResGm);
|
||||
__aicore__ inline void AllocEventID(TPipe *pipe);
|
||||
__aicore__ inline void FreeEventID(TPipe *pipe);
|
||||
__aicore__ inline void ComputeMm1(const RunInfo &info);
|
||||
|
||||
private:
|
||||
using T = float;
|
||||
using X_T = typename AscendC::Conditional<COMP::xDtype == X_DTYPE::BF16, bfloat16_t, half>::type;
|
||||
|
||||
__aicore__ inline uint32_t GetMSize(const RunInfo &info, uint32_t coffId);
|
||||
__aicore__ inline void CopyXGmToL1(const RunInfo &info, LocalTensor<X_T> xL1Tensor, uint32_t hIdx, uint32_t kBase);
|
||||
__aicore__ inline void CopyWeightGmToL1(LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase, uint32_t coffId);
|
||||
__aicore__ inline void LoadAToL0(const RunInfo &info, LocalTensor<X_T> aL0Tensor, LocalTensor<X_T> xL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize);
|
||||
__aicore__ inline void LoadBToL0(const RunInfo &info, LocalTensor<X_T> bL0Tensor, LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase, uint32_t nStart, uint32_t nDealSize);
|
||||
__aicore__ inline void MatrixMmad(LocalTensor<T> cL0Tensor, LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C);
|
||||
__aicore__ inline void CopyOutMm1Res(const RunInfo &info, LocalTensor<T> cL0Tensor,
|
||||
uint32_t coffId, uint32_t mStart, uint32_t mDealSize, uint32_t nStart, uint32_t nDealSize);
|
||||
|
||||
ConstInfo constInfo_ = {};
|
||||
CompressorTools<COMP> tools_;
|
||||
|
||||
// GM
|
||||
GlobalTensor<X_T> xGm_;
|
||||
GlobalTensor<X_T> wkvGm_;
|
||||
GlobalTensor<X_T> wgateGm_;
|
||||
GlobalTensor<MM1_OUT_T>kvMm1ResGm;
|
||||
GlobalTensor<MM1_OUT_T>scoreMm1ResGm;
|
||||
GlobalTensor<int32_t> cuSeqlensGm_;
|
||||
GlobalTensor<int32_t> sequsedGm_;
|
||||
GlobalTensor<int32_t> startPosGm_;
|
||||
bool isExistSeqUsed = false;
|
||||
|
||||
// =================================L1 Buffer=================================
|
||||
static constexpr uint32_t L1_X_SIZE = 128 * 1024;
|
||||
static constexpr uint32_t L1_W_SIZE = 128 * 1024;
|
||||
// L1 Buffer
|
||||
TBuf<TPosition::A1> xBufL1;
|
||||
TBuf<TPosition::A1> wBufL1;
|
||||
// =================================L0 Buffer=================================
|
||||
// L0 buffer size
|
||||
static constexpr uint32_t L0A_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k
|
||||
static constexpr uint32_t L0B_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k
|
||||
static constexpr uint32_t L0C_PP_SIZE = 64 * 1024; // (128 * 2) * 64 * 4 = 64k
|
||||
// L0_A
|
||||
TBuf<TPosition::A2> tmpBufL0A;
|
||||
// L0_B
|
||||
TBuf<TPosition::B2> tmpBufL0B;
|
||||
// L0_C
|
||||
TBuf<TPosition::CO1> tmpBufL0C;
|
||||
// =================================Event&Buffer ID===========================
|
||||
// mte2 <> mte1 EventID
|
||||
static constexpr uint32_t X_EVENT0 = EVENT_ID0;
|
||||
static constexpr uint32_t X_EVENT1 = EVENT_ID1;
|
||||
uint32_t xBufId = 0; // 用于DB计数
|
||||
static constexpr uint32_t W_EVENT0 = EVENT_ID4;
|
||||
static constexpr uint32_t W_EVENT1 = EVENT_ID5;
|
||||
static constexpr uint32_t W_EVENT2 = EVENT_ID6;
|
||||
static constexpr uint32_t W_EVENT3 = EVENT_ID7;
|
||||
uint32_t wBufId = 0; // 用于DB计数
|
||||
// mte1 <> mmad EventID
|
||||
static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3;
|
||||
static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4;
|
||||
uint32_t l0abBufId = 0;
|
||||
// mmad <> fixpipe EventID
|
||||
static constexpr uint32_t L0C_EVENT0 = EVENT_ID0; // 每块L0C单独分配EVENT_ID
|
||||
static constexpr uint32_t L0C_EVENT1 = EVENT_ID1;
|
||||
static constexpr uint32_t L0C_EVENT2 = EVENT_ID2;
|
||||
static constexpr uint32_t L0C_EVENT3 = EVENT_ID3;
|
||||
uint32_t l0cBufId = 0;
|
||||
|
||||
// =================================Loop======================================
|
||||
uint32_t curBIdx_ = 0;
|
||||
uint32_t curSIdx_ = 0;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::InitParams(const ConstInfo &constInfo, const CompressorTools<COMP> &tools)
|
||||
{
|
||||
this->constInfo_ = constInfo;
|
||||
this->tools_ = tools;
|
||||
}
|
||||
|
||||
template <typename COMP> __aicore__ inline void CompressorBlockCube<COMP>::Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut)
|
||||
{
|
||||
xGm_.SetGlobalBuffer((__gm__ X_T *)x);
|
||||
wkvGm_.SetGlobalBuffer((__gm__ X_T *)wKv);
|
||||
wgateGm_.SetGlobalBuffer((__gm__ X_T *)wGate);
|
||||
startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos);
|
||||
isExistSeqUsed = (seqUsed != nullptr);
|
||||
if (isExistSeqUsed) {
|
||||
sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed);
|
||||
}
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::InitBuffers(TPipe *pipe)
|
||||
{
|
||||
// L1
|
||||
// 1. coff=1时, mBase=256, kL1=256, X单次拷贝到L1的数据量最大为mBase*kL1*sizeof(BF16/FP16)=256*256*2=128K
|
||||
// 2. coff=2时, mBase=128, kL1=256, r最大为128, X单次拷贝到L1的最大数据量为(128+r)*kL1*sizeof(BF16/FP16)<=128K
|
||||
pipe->InitBuffer(xBufL1, L1_X_SIZE * 2);
|
||||
// dBaseSize<=64, wkv和wgate各一份, kL1=256, 右矩阵为dBaseSize*2*sizeof(BF16/FP16)<=64K
|
||||
// cur和pre循环使用, 2份buffer就足够
|
||||
pipe->InitBuffer(wBufL1, L1_W_SIZE * 2);
|
||||
|
||||
// L0
|
||||
pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2);
|
||||
pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2);
|
||||
pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 4);
|
||||
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::InitGlobalBuffers(const GlobalTensor<MM1_OUT_T>& kvMm1ResGm, const GlobalTensor<MM1_OUT_T>& scoreMm1ResGm)
|
||||
{
|
||||
this->kvMm1ResGm = kvMm1ResGm;
|
||||
this->scoreMm1ResGm = scoreMm1ResGm;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::AllocEventID(TPipe *pipe)
|
||||
{
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT0);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT1);
|
||||
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT0);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT1);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT2);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT3);
|
||||
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT0);
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT1);
|
||||
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT0);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT1);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT2);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT3);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::FreeEventID(TPipe *pipe)
|
||||
{
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT0);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT1);
|
||||
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT0);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT1);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT2);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT3);
|
||||
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT0);
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT1);
|
||||
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT0);
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT1);
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT2);
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT3);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::CopyXGmToL1(const RunInfo &info, LocalTensor<X_T> xL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase)
|
||||
{
|
||||
uint32_t tStart = tools_.GetTIdxByBatch(info.bStart) + info.sStart; // 此基本块在整个序列中的位置
|
||||
uint32_t copySeqCnt = info.dealSeqCnt; // 此基本块处理的长度
|
||||
|
||||
uint32_t xL1Offset = 0 * (32 / sizeof(X_T));
|
||||
uint64_t sIdx = tStart; // 起始s在整个T的起始点
|
||||
uint64_t gmOffset = sIdx * constInfo_.hSize + hIdx;
|
||||
uint32_t nValue = copySeqCnt;
|
||||
uint32_t dValue = kBase; // 拷贝的列数kBase
|
||||
uint32_t srcDValue = constInfo_.hSize;
|
||||
uint32_t dstNzC0Stride = (copySeqCnt + 15) / 16 * 16; // 1行变2行的行方向的偏移,需要16对齐
|
||||
CopySingleMatrixNDToNZ(xL1Tensor[xL1Offset], xGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::CopyWeightGmToL1(LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase, uint32_t coffId)
|
||||
{
|
||||
// coffId=0, 搬运左矩阵的数据; coffId=1, 搬运右矩阵的数据
|
||||
uint64_t gmOffset = coffId * constInfo_.headDim * constInfo_.hSize + constInfo_.dIdx * constInfo_.dBaseSize * constInfo_.hSize + hIdx;
|
||||
uint32_t wkvL1Offset = 0;
|
||||
uint32_t wgateL1Offset = constInfo_.dBaseSize * (32 / sizeof(X_T)); // wgate与wkv的起始点相隔dBaseSize个32B
|
||||
uint32_t nValue = constInfo_.dBaseSize;
|
||||
uint32_t dValue = kBase;
|
||||
uint32_t srcDValue = constInfo_.hSize;
|
||||
uint32_t dstNzC0Stride = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
|
||||
CopySingleMatrixNDToNZ(wL1Tensor[wkvL1Offset], wkvGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
CopySingleMatrixNDToNZ(wL1Tensor[wgateL1Offset], wgateGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::LoadAToL0(const RunInfo &info, LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> xL1Tensor, uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize)
|
||||
{
|
||||
uint32_t mSize = info.dealSeqCnt;
|
||||
|
||||
uint32_t mSizeAlign = Align(mSize, 16U);
|
||||
uint32_t xTensorOffset = kStart * mSizeAlign + mStart * (32 / sizeof(X_T));
|
||||
uint32_t mDealSizeAlign = Align(mDealSize, 16U);
|
||||
|
||||
LoadData2DParamsV2 loadData2DParamsV2;
|
||||
loadData2DParamsV2.mStartPosition = 0;
|
||||
loadData2DParamsV2.kStartPosition = 0;
|
||||
loadData2DParamsV2.mStep = mDealSizeAlign / 16;
|
||||
loadData2DParamsV2.kStep = kBase / (32 / sizeof(X_T));
|
||||
loadData2DParamsV2.srcStride = mSizeAlign / 16;
|
||||
loadData2DParamsV2.dstStride = loadData2DParamsV2.mStep;
|
||||
loadData2DParamsV2.ifTranspose = false;
|
||||
LoadData(aL0Tensor, xL1Tensor[xTensorOffset], loadData2DParamsV2);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::LoadBToL0(const RunInfo &info, LocalTensor<X_T> bL0Tensor, LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase, uint32_t nStart, uint32_t nDealSize)
|
||||
{
|
||||
uint32_t nSize = 2 * constInfo_.dBaseSize;
|
||||
|
||||
uint32_t nSizeAlign = Align(nSize, 16U);
|
||||
uint64_t wTensorOffset = nSizeAlign * kStart + nStart * (32 / sizeof(X_T));
|
||||
uint32_t nDealSizeAlign = Align(nDealSize, 16U);
|
||||
|
||||
LoadData2DParamsV2 loadData2DParamsV2;
|
||||
loadData2DParamsV2.mStartPosition = 0;
|
||||
loadData2DParamsV2.kStartPosition = 0;
|
||||
loadData2DParamsV2.mStep = nDealSizeAlign / 16;
|
||||
loadData2DParamsV2.kStep = kBase / (32 / sizeof(X_T));
|
||||
loadData2DParamsV2.srcStride = nSizeAlign / 16;
|
||||
loadData2DParamsV2.dstStride = loadData2DParamsV2.mStep;
|
||||
loadData2DParamsV2.ifTranspose = false;
|
||||
LoadData(bL0Tensor, wL1Tensor[wTensorOffset], loadData2DParamsV2);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::MatrixMmad(LocalTensor<T> cL0Tensor, LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C)
|
||||
{
|
||||
MmadParams mmadParams;
|
||||
mmadParams.m = mActSize < 16 ? 16 : mActSize;
|
||||
mmadParams.n = nDealSize;
|
||||
mmadParams.k = kActSize;
|
||||
mmadParams.cmatrixInitVal = isInitL0C;
|
||||
mmadParams.cmatrixSource = false;
|
||||
Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::CopyOutMm1Res(const RunInfo &info, LocalTensor<T> cL0Tensor,
|
||||
uint32_t coffId, uint32_t mStart, uint32_t mDealSize, uint32_t nStart, uint32_t nDealSize)
|
||||
{
|
||||
// coffId=0, 存左矩阵的数据; coffId=1, 存右矩阵的数据
|
||||
FixpipeParamsV220 fixParams;
|
||||
fixParams.mSize = mDealSize;
|
||||
fixParams.srcStride = (mDealSize + 15) / 16 * 16; // 需要16对齐
|
||||
fixParams.dstStride = (uint32_t)COMP::coff * constInfo_.headDim;
|
||||
fixParams.ndNum = 1;
|
||||
|
||||
uint64_t dbOffset = info.cubeDbIdx * constInfo_.dbSize;
|
||||
uint64_t gmOffset = constInfo_.dIdx * constInfo_.dBaseSize + coffId * constInfo_.headDim + mStart * fixParams.dstStride + dbOffset;
|
||||
uint32_t kvOffset = (mDealSize + 15) / 16 * 16 * nStart;
|
||||
uint32_t scoreOffset = (mDealSize + 15) / 16 * 16 * ((nStart + constInfo_.dBaseSize) % (2 * constInfo_.dBaseSize));
|
||||
if (nStart < constInfo_.dBaseSize) {
|
||||
fixParams.nSize = min(constInfo_.dBaseSize - nStart, nDealSize);
|
||||
Fixpipe(kvMm1ResGm[gmOffset], cL0Tensor[kvOffset], fixParams);
|
||||
}
|
||||
if (nStart + nDealSize > constInfo_.dBaseSize) {
|
||||
fixParams.nSize = min(nStart + nDealSize - constInfo_.dBaseSize, nDealSize);
|
||||
Fixpipe(scoreMm1ResGm[gmOffset], cL0Tensor[scoreOffset], fixParams);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorBlockCube<COMP>::GetMSize(const RunInfo &info, uint32_t coffId)
|
||||
{
|
||||
return info.dealSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCube<COMP>::ComputeMm1(const RunInfo &info)
|
||||
{
|
||||
static constexpr uint32_t K_L1_BASE = 256;
|
||||
static constexpr uint32_t M_L0_BASE = 128;
|
||||
static constexpr uint32_t K_L0_BASE = 128;
|
||||
static constexpr uint32_t N_L0_BASE = 128;
|
||||
uint32_t nCoff = (uint32_t)COMP::coff;
|
||||
|
||||
// hSize为K_SIZE=512的倍数
|
||||
uint32_t hStart = info.hStart;
|
||||
uint32_t hSize = info.dealKSize;
|
||||
uint32_t hIdxStart = (constInfo_.aiCoreIdx % constInfo_.dBasicBlockNum) * K_L1_BASE; // 每组核内的h循环起始不同
|
||||
uint32_t kSize = K_L1_BASE;
|
||||
for (uint32_t h = 0; h < hSize; h += K_L1_BASE) {
|
||||
// h方向错位搬运
|
||||
uint32_t hIdx = (h + hIdxStart) % (CeilDivT(hSize, K_L1_BASE) * K_L1_BASE);
|
||||
if (hIdx + K_L1_BASE > hSize) {
|
||||
kSize = hSize - hIdx;
|
||||
} else {
|
||||
kSize = K_L1_BASE;
|
||||
}
|
||||
bool isFirst = (h == 0);
|
||||
bool isLast = (h + K_L1_BASE >= hSize);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT0 + xBufId);
|
||||
LocalTensor<X_T> xL1Tensor = xBufL1.GetWithOffset<X_T>(L1_X_SIZE / sizeof(X_T), xBufId * L1_X_SIZE);
|
||||
CopyXGmToL1(info, xL1Tensor, hStart + hIdx, kSize);
|
||||
SetFlag<HardEvent::MTE2_MTE1>(X_EVENT0 + xBufId);
|
||||
WaitFlag<HardEvent::MTE2_MTE1>(X_EVENT0 + xBufId);
|
||||
for (uint32_t i = nCoff; i > 0; i--) {
|
||||
// coffId=0, 计算pre数据; coffId=1, 计算cur数据
|
||||
uint32_t coffId = i - 1;
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT0 + wBufId);
|
||||
LocalTensor<X_T> wL1Tensor = wBufL1.GetWithOffset<X_T>(L1_W_SIZE / sizeof(X_T), wBufId * L1_W_SIZE);
|
||||
CopyWeightGmToL1(wL1Tensor, hStart + hIdx, kSize, coffId);
|
||||
SetFlag<HardEvent::MTE2_MTE1>(W_EVENT0 + wBufId);
|
||||
WaitFlag<HardEvent::MTE2_MTE1>(W_EVENT0 + wBufId);
|
||||
|
||||
uint32_t mSize = GetMSize(info, coffId);
|
||||
uint32_t actMDealSize = M_L0_BASE;
|
||||
for (uint32_t mL0 = 0; mL0 < mSize; mL0 += M_L0_BASE) {
|
||||
if (mL0 + M_L0_BASE > mSize) {
|
||||
actMDealSize = mSize - mL0;
|
||||
}
|
||||
uint32_t nDealSize = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
|
||||
uint32_t actNDealSize = N_L0_BASE;
|
||||
for (uint32_t nL0 = 0; nL0 < nDealSize; nL0 += N_L0_BASE) {
|
||||
if (nL0 + N_L0_BASE > nDealSize) {
|
||||
actNDealSize = nDealSize - nL0;
|
||||
}
|
||||
l0cBufId = (mL0 / M_L0_BASE) * 2 + (nL0 / N_L0_BASE) + coffId;
|
||||
LocalTensor<T> cL0Tensor =
|
||||
tmpBufL0C.GetWithOffset<T>((L0C_PP_SIZE / sizeof(T)), l0cBufId * L0C_PP_SIZE);
|
||||
if (isFirst) {
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT0 + l0cBufId);
|
||||
}
|
||||
uint32_t actKDealSize = K_L0_BASE;
|
||||
for (uint32_t kL0 = 0; kL0 < kSize; kL0 += K_L0_BASE) {
|
||||
if (kL0 + K_L0_BASE > kSize) {
|
||||
actKDealSize = kSize - kL0;
|
||||
}
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT0 + l0abBufId);
|
||||
LocalTensor<X_T> aL0Tensor =
|
||||
tmpBufL0A.GetWithOffset<X_T>(L0A_PP_SIZE / sizeof(X_T), l0abBufId * L0A_PP_SIZE);
|
||||
LocalTensor<X_T> bL0Tensor =
|
||||
tmpBufL0B.GetWithOffset<X_T>(L0B_PP_SIZE / sizeof(X_T), l0abBufId * L0B_PP_SIZE);
|
||||
LoadAToL0(info, aL0Tensor, xL1Tensor, kL0, actKDealSize, mL0, actMDealSize);
|
||||
LoadBToL0(info, bL0Tensor, wL1Tensor, kL0, actKDealSize, nL0, actNDealSize);
|
||||
SetFlag<HardEvent::MTE1_M>(L0AB_EVENT0 + l0abBufId);
|
||||
WaitFlag<HardEvent::MTE1_M>(L0AB_EVENT0 + l0abBufId);
|
||||
bool isInitL0C = isFirst && (kL0 == 0);
|
||||
MatrixMmad(cL0Tensor, aL0Tensor, bL0Tensor, actMDealSize, actNDealSize, actKDealSize,
|
||||
isInitL0C);
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT0 + l0abBufId);
|
||||
l0abBufId = (l0abBufId + 1) % 2;
|
||||
}
|
||||
if (isLast) {
|
||||
SetFlag<HardEvent::M_FIX>(L0C_EVENT0 + l0cBufId);
|
||||
WaitFlag<HardEvent::M_FIX>(L0C_EVENT0 + l0cBufId);
|
||||
CopyOutMm1Res(info, cL0Tensor, coffId, mL0, actMDealSize, nL0, actNDealSize);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT0 + l0cBufId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT0 + wBufId);
|
||||
wBufId = (wBufId + 1) % 2;
|
||||
}
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT0 + xBufId);
|
||||
xBufId = (xBufId + 1) % 2;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif // COMPRESSOR_BLOCK_CUBE_H
|
||||
@@ -0,0 +1,437 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_block_cube_full_load.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_BLOCK_CUBE_FULL_LOAD_H
|
||||
#define COMPRESSOR_BLOCK_CUBE_FULL_LOAD_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_tools.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
template<typename COMP> class CompressorBlockCubeFullLoad {
|
||||
using MM1_OUT_T = float;
|
||||
public:
|
||||
__aicore__ inline CompressorBlockCubeFullLoad(){};
|
||||
__aicore__ inline void InitParams(const ConstInfo &constInfo, const CompressorTools<COMP> &tools);
|
||||
__aicore__ inline void Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut);
|
||||
__aicore__ inline void InitBuffers(TPipe *pipe);
|
||||
__aicore__ inline void InitGlobalBuffers(const GlobalTensor<MM1_OUT_T>& kvMm1ResGm, const GlobalTensor<MM1_OUT_T>& scoreMm1ResGm);
|
||||
__aicore__ inline void AllocEventID(TPipe *pipe);
|
||||
__aicore__ inline void FreeEventID(TPipe *pipe);
|
||||
__aicore__ inline void ComputeMm1(const RunInfo &info);
|
||||
|
||||
private:
|
||||
using T = float;
|
||||
using X_T = typename AscendC::Conditional<COMP::xDtype == X_DTYPE::BF16, bfloat16_t, half>::type;
|
||||
|
||||
__aicore__ inline uint32_t GetMSize(const RunInfo &info, uint32_t coffId);
|
||||
__aicore__ inline void CopyXGmToL1(LocalTensor<X_T> xL1Tensor, uint32_t hIdx, uint32_t kBase);
|
||||
__aicore__ inline void CopyWeightGmToL1(LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase, uint32_t coffId);
|
||||
__aicore__ inline void LoadAToL0(LocalTensor<X_T> aL0Tensor, LocalTensor<X_T> xL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize);
|
||||
__aicore__ inline void LoadBToL0(LocalTensor<X_T> bL0Tensor, LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase, uint32_t nStart, uint32_t nDealSize);
|
||||
__aicore__ inline void MatrixMmad(LocalTensor<T> cL0Tensor, LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C);
|
||||
__aicore__ inline void CopyOutMm1Res(const RunInfo &info, LocalTensor<T> cL0Tensor,
|
||||
uint32_t coffId, uint32_t mStart, uint32_t mDealSize, uint32_t nStart, uint32_t nDealSize);
|
||||
|
||||
ConstInfo constInfo_ = {};
|
||||
CompressorTools<COMP> tools_;
|
||||
|
||||
// GM
|
||||
GlobalTensor<X_T> xGm_;
|
||||
GlobalTensor<X_T> wkvGm_;
|
||||
GlobalTensor<X_T> wgateGm_;
|
||||
GlobalTensor<MM1_OUT_T>kvMm1ResGm;
|
||||
GlobalTensor<MM1_OUT_T>scoreMm1ResGm;
|
||||
GlobalTensor<int32_t> cuSeqlensGm_;
|
||||
GlobalTensor<int32_t> sequsedGm_;
|
||||
GlobalTensor<int32_t> startPosGm_;
|
||||
bool isExistSeqUsed = false;
|
||||
|
||||
// =================================L1 Buffer=================================
|
||||
static constexpr uint32_t L1_X_SIZE = 128 * 1024;
|
||||
static constexpr uint32_t L1_W_SIZE = 128 * 1024;
|
||||
// L1 Buffer
|
||||
TBuf<TPosition::A1> xBufL1;
|
||||
TBuf<TPosition::A1> wBufL1;
|
||||
// =================================L0 Buffer=================================
|
||||
// L0 buffer size
|
||||
static constexpr uint32_t L0A_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k
|
||||
static constexpr uint32_t L0B_PP_SIZE = 32 * 1024; // 128 * 128 * 2 = 32k
|
||||
static constexpr uint32_t L0C_PP_SIZE = 64 * 1024; // (128 * 2) * 64 * 4 = 64k
|
||||
// L0_A
|
||||
TBuf<TPosition::A2> tmpBufL0A;
|
||||
// L0_B
|
||||
TBuf<TPosition::B2> tmpBufL0B;
|
||||
// L0_C
|
||||
TBuf<TPosition::CO1> tmpBufL0C;
|
||||
// =================================Event&Buffer ID===========================
|
||||
// mte2 <> mte1 EventID
|
||||
static constexpr uint32_t X_EVENT0 = EVENT_ID0;
|
||||
static constexpr uint32_t X_EVENT1 = EVENT_ID1;
|
||||
uint32_t xBufId = 0; // 用于DB计数
|
||||
static constexpr uint32_t W_EVENT0 = EVENT_ID4;
|
||||
static constexpr uint32_t W_EVENT1 = EVENT_ID5;
|
||||
static constexpr uint32_t W_EVENT2 = EVENT_ID6;
|
||||
static constexpr uint32_t W_EVENT3 = EVENT_ID7;
|
||||
uint32_t wBufId = 0; // 用于DB计数
|
||||
// mte1 <> mmad EventID
|
||||
static constexpr uint32_t L0AB_EVENT0 = EVENT_ID3;
|
||||
static constexpr uint32_t L0AB_EVENT1 = EVENT_ID4;
|
||||
uint32_t l0abBufId = 0;
|
||||
// mmad <> fixpipe EventID
|
||||
static constexpr uint32_t L0C_EVENT0 = EVENT_ID0; // 每块L0C单独分配EVENT_ID
|
||||
static constexpr uint32_t L0C_EVENT1 = EVENT_ID1;
|
||||
static constexpr uint32_t L0C_EVENT2 = EVENT_ID2;
|
||||
static constexpr uint32_t L0C_EVENT3 = EVENT_ID3;
|
||||
uint32_t l0cBufId = 0;
|
||||
|
||||
// =================================Loop======================================
|
||||
uint32_t curBIdx_ = 0;
|
||||
uint32_t curSIdx_ = 0;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::InitParams(const ConstInfo &constInfo, const CompressorTools<COMP> &tools)
|
||||
{
|
||||
this->constInfo_ = constInfo;
|
||||
this->tools_ = tools;
|
||||
}
|
||||
|
||||
template <typename COMP> __aicore__ inline void CompressorBlockCubeFullLoad<COMP>::Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut)
|
||||
{
|
||||
xGm_.SetGlobalBuffer((__gm__ X_T *)x);
|
||||
wkvGm_.SetGlobalBuffer((__gm__ X_T *)wKv);
|
||||
wgateGm_.SetGlobalBuffer((__gm__ X_T *)wGate);
|
||||
startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos);
|
||||
isExistSeqUsed = (seqUsed != nullptr);
|
||||
if (isExistSeqUsed) {
|
||||
sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed);
|
||||
}
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::InitBuffers(TPipe *pipe)
|
||||
{
|
||||
// L1
|
||||
// 1. coff=1时, mBase=256, kL1=256, X单次拷贝到L1的数据量最大为mBase*kL1*sizeof(BF16/FP16)=256*256*2=128K
|
||||
// 2. coff=2时, mBase=128, kL1=256, r最大为128, X单次拷贝到L1的最大数据量为(128+r)*kL1*sizeof(BF16/FP16)<=128K
|
||||
pipe->InitBuffer(xBufL1, L1_X_SIZE * 2);
|
||||
// dBaseSize<=64, wkv和wgate各一份, kL1=256, 右矩阵为dBaseSize*2*sizeof(BF16/FP16)<=64K
|
||||
// cur和pre循环使用, 2份buffer就足够
|
||||
pipe->InitBuffer(wBufL1, L1_W_SIZE * 2);
|
||||
|
||||
// L0
|
||||
pipe->InitBuffer(tmpBufL0A, L0A_PP_SIZE * 2);
|
||||
pipe->InitBuffer(tmpBufL0B, L0B_PP_SIZE * 2);
|
||||
pipe->InitBuffer(tmpBufL0C, L0C_PP_SIZE * 4);
|
||||
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::InitGlobalBuffers(const GlobalTensor<MM1_OUT_T>& kvMm1ResGm, const GlobalTensor<MM1_OUT_T>& scoreMm1ResGm)
|
||||
{
|
||||
this->kvMm1ResGm = kvMm1ResGm;
|
||||
this->scoreMm1ResGm = scoreMm1ResGm;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::AllocEventID(TPipe *pipe)
|
||||
{
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT0);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT1);
|
||||
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT0);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT1);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT2);
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT3);
|
||||
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT0);
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT1);
|
||||
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT0);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT1);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT2);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT3);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::FreeEventID(TPipe *pipe)
|
||||
{
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT0);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT1);
|
||||
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT0);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT1);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT2);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT3);
|
||||
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT0);
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT1);
|
||||
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT0);
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT1);
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT2);
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT3);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::CopyXGmToL1(LocalTensor<X_T> xL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase)
|
||||
{
|
||||
uint32_t copySeqCnt = constInfo_.mEnd - constInfo_.mStart; //info.dealSeqCnt; // 此基本块处理的长度
|
||||
|
||||
uint32_t xL1Offset = 0 * (32 / sizeof(X_T));
|
||||
uint64_t sIdx = constInfo_.mStart; // 起始s在整个T的起始点
|
||||
uint64_t gmOffset = sIdx * constInfo_.hSize + hIdx;
|
||||
uint32_t nValue = copySeqCnt;
|
||||
uint32_t dValue = kBase; // 拷贝的列数kBase
|
||||
uint32_t srcDValue = constInfo_.hSize;
|
||||
uint32_t dstNzC0Stride = (copySeqCnt + 15) / 16 * 16; // 1行变2行的行方向的偏移,需要16对齐
|
||||
CopySingleMatrixNDToNZ(xL1Tensor[xL1Offset], xGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::CopyWeightGmToL1(LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t hIdx, uint32_t kBase, uint32_t coffId)
|
||||
{
|
||||
uint64_t gmOffset = coffId * constInfo_.headDim * constInfo_.hSize + constInfo_.nStart * constInfo_.hSize + hIdx;
|
||||
uint32_t wkvL1Offset = 0;
|
||||
uint32_t wgateL1Offset = constInfo_.dBaseSize * (32 / sizeof(X_T)); // wgate与wkv的起始点相隔dBaseSize个32B
|
||||
uint32_t nValue = constInfo_.dBaseSize;
|
||||
uint32_t dValue = kBase;
|
||||
uint32_t srcDValue = constInfo_.hSize;
|
||||
uint32_t dstNzC0Stride = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
|
||||
CopySingleMatrixNDToNZ(wL1Tensor[wkvL1Offset], wkvGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
CopySingleMatrixNDToNZ(wL1Tensor[wgateL1Offset], wgateGm_[gmOffset], nValue, dValue, srcDValue, dstNzC0Stride);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::LoadAToL0(LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> xL1Tensor, uint32_t kStart, uint32_t kBase, uint32_t mStart, uint32_t mDealSize)
|
||||
{
|
||||
uint32_t mSize = constInfo_.mEnd - constInfo_.mStart;
|
||||
|
||||
uint32_t mSizeAlign = Align(mSize, 16U);
|
||||
uint32_t xTensorOffset = kStart * mSizeAlign + mStart * (32 / sizeof(X_T));
|
||||
uint32_t mDealSizeAlign = Align(mDealSize, 16U);
|
||||
|
||||
LoadData2DParamsV2 loadData2DParamsV2;
|
||||
loadData2DParamsV2.mStartPosition = 0;
|
||||
loadData2DParamsV2.kStartPosition = 0;
|
||||
loadData2DParamsV2.mStep = mDealSizeAlign / 16;
|
||||
loadData2DParamsV2.kStep = kBase / (32 / sizeof(X_T));
|
||||
loadData2DParamsV2.srcStride = mSizeAlign / 16;
|
||||
loadData2DParamsV2.dstStride = loadData2DParamsV2.mStep;
|
||||
loadData2DParamsV2.ifTranspose = false;
|
||||
LoadData(aL0Tensor, xL1Tensor[xTensorOffset], loadData2DParamsV2);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::LoadBToL0(LocalTensor<X_T> bL0Tensor, LocalTensor<X_T> wL1Tensor,
|
||||
uint32_t kStart, uint32_t kBase, uint32_t nStart, uint32_t nDealSize)
|
||||
{
|
||||
uint32_t nSize = 2 * constInfo_.dBaseSize;
|
||||
|
||||
uint32_t nSizeAlign = Align(nSize, 16U);
|
||||
uint64_t wTensorOffset = nSizeAlign * kStart + nStart * (32 / sizeof(X_T));
|
||||
uint32_t nDealSizeAlign = Align(nDealSize, 16U);
|
||||
|
||||
LoadData2DParamsV2 loadData2DParamsV2;
|
||||
loadData2DParamsV2.mStartPosition = 0;
|
||||
loadData2DParamsV2.kStartPosition = 0;
|
||||
loadData2DParamsV2.mStep = nDealSizeAlign / 16;
|
||||
loadData2DParamsV2.kStep = kBase / (32 / sizeof(X_T));
|
||||
loadData2DParamsV2.srcStride = nSizeAlign / 16;
|
||||
loadData2DParamsV2.dstStride = loadData2DParamsV2.mStep;
|
||||
loadData2DParamsV2.ifTranspose = false;
|
||||
LoadData(bL0Tensor, wL1Tensor[wTensorOffset], loadData2DParamsV2);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::MatrixMmad(LocalTensor<T> cL0Tensor, LocalTensor<X_T> aL0Tensor,
|
||||
LocalTensor<X_T> bL0Tensor, uint32_t mActSize, uint32_t nDealSize, uint32_t kActSize, bool isInitL0C)
|
||||
{
|
||||
MmadParams mmadParams;
|
||||
mmadParams.m = mActSize < 16 ? 16 : mActSize;
|
||||
mmadParams.n = nDealSize;
|
||||
mmadParams.k = kActSize;
|
||||
mmadParams.cmatrixInitVal = isInitL0C;
|
||||
mmadParams.cmatrixSource = false;
|
||||
Mmad(cL0Tensor, aL0Tensor, bL0Tensor, mmadParams);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::CopyOutMm1Res(const RunInfo &info, LocalTensor<T> cL0Tensor,
|
||||
uint32_t coffId, uint32_t mStart, uint32_t mDealSize, uint32_t nStart, uint32_t nDealSize)
|
||||
{
|
||||
// coffId=0, 存左矩阵的数据; coffId=1, 存右矩阵的数据
|
||||
FixpipeParamsV220 fixParams;
|
||||
fixParams.mSize = mDealSize;
|
||||
fixParams.srcStride = (mDealSize + 15) / 16 * 16; // 需要16对齐
|
||||
fixParams.dstStride = (uint32_t)COMP::coff * constInfo_.headDim;
|
||||
fixParams.ndNum = 1;
|
||||
|
||||
uint64_t dbOffset = info.cubeDbIdx * constInfo_.dbSize;
|
||||
uint64_t gmOffset = constInfo_.nStart + coffId * constInfo_.headDim + mStart * fixParams.dstStride + dbOffset;
|
||||
uint32_t kvOffset = (mDealSize + 15) / 16 * 16 * nStart;
|
||||
uint32_t scoreOffset = (mDealSize + 15) / 16 * 16 * ((nStart + constInfo_.dBaseSize) % (2 * constInfo_.dBaseSize));
|
||||
|
||||
if (nStart < constInfo_.dBaseSize) {
|
||||
fixParams.nSize = min(constInfo_.dBaseSize - nStart, nDealSize);
|
||||
Fixpipe(kvMm1ResGm[gmOffset], cL0Tensor[kvOffset], fixParams);
|
||||
}
|
||||
if (nStart + nDealSize > constInfo_.dBaseSize) {
|
||||
fixParams.nSize = min(nStart + nDealSize - constInfo_.dBaseSize, nDealSize);
|
||||
Fixpipe(scoreMm1ResGm[gmOffset], cL0Tensor[scoreOffset], fixParams);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorBlockCubeFullLoad<COMP>::GetMSize(const RunInfo &info, uint32_t coffId)
|
||||
{
|
||||
return info.dealSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorBlockCubeFullLoad<COMP>::ComputeMm1(const RunInfo &info)
|
||||
{
|
||||
uint32_t mSize = info.dealSeqCnt;
|
||||
if (mSize == 0) {
|
||||
return;
|
||||
}
|
||||
static constexpr uint32_t K_L1_BASE = 128;
|
||||
static constexpr uint32_t M_L0_BASE = 128;
|
||||
static constexpr uint32_t K_L0_BASE = 128;
|
||||
static constexpr uint32_t N_L0_BASE = 128;
|
||||
uint32_t nCoff = (uint32_t)COMP::coff;
|
||||
|
||||
// hSize为K_SIZE=512的倍数
|
||||
uint32_t hStart = constInfo_.kStart;
|
||||
uint32_t hSize = constInfo_.kEnd - constInfo_.kStart;
|
||||
uint32_t hIdxStart = (constInfo_.aiCoreIdx % constInfo_.dBasicBlockNum) * K_L1_BASE; // 每组核内的h循环起始不同
|
||||
uint32_t kSize = K_L1_BASE;
|
||||
for (uint32_t h = 0; h < hSize; h += K_L1_BASE) {
|
||||
// h方向错位搬运
|
||||
uint32_t hIdx = (h + hIdxStart) % (CeilDivT(hSize, K_L1_BASE) * K_L1_BASE);
|
||||
if (hIdx + K_L1_BASE > hSize) {
|
||||
kSize = hSize - hIdx;
|
||||
} else {
|
||||
kSize = K_L1_BASE;
|
||||
}
|
||||
bool isFirst = (h == 0);
|
||||
bool isLast = (h + K_L1_BASE >= hSize);
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(X_EVENT0 + xBufId);
|
||||
LocalTensor<X_T> xL1Tensor = xBufL1.GetWithOffset<X_T>(L1_X_SIZE / sizeof(X_T), xBufId * L1_X_SIZE);
|
||||
CopyXGmToL1(xL1Tensor, hStart + hIdx, kSize);
|
||||
SetFlag<HardEvent::MTE2_MTE1>(X_EVENT0 + xBufId);
|
||||
WaitFlag<HardEvent::MTE2_MTE1>(X_EVENT0 + xBufId);
|
||||
for (uint32_t i = nCoff; i > 0; i--) {
|
||||
// coffId=0, 计算pre数据; coffId=1, 计算cur数据
|
||||
uint32_t coffId = i - 1;
|
||||
WaitFlag<HardEvent::MTE1_MTE2>(W_EVENT0 + wBufId);
|
||||
LocalTensor<X_T> wL1Tensor = wBufL1.GetWithOffset<X_T>(L1_W_SIZE / sizeof(X_T), wBufId * L1_W_SIZE);
|
||||
CopyWeightGmToL1(wL1Tensor, hStart + hIdx, kSize, coffId);
|
||||
SetFlag<HardEvent::MTE2_MTE1>(W_EVENT0 + wBufId);
|
||||
WaitFlag<HardEvent::MTE2_MTE1>(W_EVENT0 + wBufId);
|
||||
|
||||
uint32_t actMDealSize = M_L0_BASE;
|
||||
for (uint32_t mL0 = 0; mL0 < mSize; mL0 += M_L0_BASE) {
|
||||
if (mL0 + M_L0_BASE > mSize) {
|
||||
actMDealSize = mSize - mL0;
|
||||
}
|
||||
uint32_t nDealSize = 2 * constInfo_.dBaseSize; // 2: wkv和wgate各搬运dBaseSize行, dBaseSize需保证8的倍数
|
||||
uint32_t actNDealSize = N_L0_BASE;
|
||||
for (uint32_t nL0 = 0; nL0 < nDealSize; nL0 += N_L0_BASE) {
|
||||
if (nL0 + N_L0_BASE > nDealSize) {
|
||||
actNDealSize = nDealSize - nL0;
|
||||
}
|
||||
l0cBufId = (mL0 / M_L0_BASE) * 2 + (nL0 / N_L0_BASE) + coffId;
|
||||
LocalTensor<T> cL0Tensor = tmpBufL0C.GetWithOffset<T>((L0C_PP_SIZE / sizeof(T)), l0cBufId * L0C_PP_SIZE);
|
||||
if (isFirst) {
|
||||
WaitFlag<HardEvent::FIX_M>(L0C_EVENT0 + l0cBufId);
|
||||
}
|
||||
uint32_t actKDealSize = K_L0_BASE;
|
||||
for (uint32_t kL0 = 0; kL0 < kSize; kL0 += K_L0_BASE) {
|
||||
if (kL0 + K_L0_BASE > kSize) {
|
||||
actKDealSize = kSize - kL0;
|
||||
}
|
||||
WaitFlag<HardEvent::M_MTE1>(L0AB_EVENT0 + l0abBufId);
|
||||
LocalTensor<X_T> aL0Tensor = tmpBufL0A.GetWithOffset<X_T>(L0A_PP_SIZE / sizeof(X_T), l0abBufId * L0A_PP_SIZE);
|
||||
LocalTensor<X_T> bL0Tensor = tmpBufL0B.GetWithOffset<X_T>(L0B_PP_SIZE / sizeof(X_T), l0abBufId * L0B_PP_SIZE);
|
||||
LoadAToL0(aL0Tensor, xL1Tensor, kL0, actKDealSize, mL0, actMDealSize);
|
||||
LoadBToL0(bL0Tensor, wL1Tensor, kL0, actKDealSize, nL0, actNDealSize);
|
||||
SetFlag<HardEvent::MTE1_M>(L0AB_EVENT0 + l0abBufId);
|
||||
WaitFlag<HardEvent::MTE1_M>(L0AB_EVENT0 + l0abBufId);
|
||||
bool isInitL0C = isFirst && (kL0 == 0);
|
||||
MatrixMmad(cL0Tensor, aL0Tensor, bL0Tensor, actMDealSize, actNDealSize, actKDealSize, isInitL0C);
|
||||
SetFlag<HardEvent::M_MTE1>(L0AB_EVENT0 + l0abBufId);
|
||||
l0abBufId = (l0abBufId + 1) % 2;
|
||||
}
|
||||
if (isLast) {
|
||||
SetFlag<HardEvent::M_FIX>(L0C_EVENT0 + l0cBufId);
|
||||
WaitFlag<HardEvent::M_FIX>(L0C_EVENT0 + l0cBufId);
|
||||
CopyOutMm1Res(info, cL0Tensor, coffId, mL0, actMDealSize, nL0, actNDealSize);
|
||||
SetFlag<HardEvent::FIX_M>(L0C_EVENT0 + l0cBufId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
SetFlag<HardEvent::MTE1_MTE2>(W_EVENT0 + wBufId);
|
||||
wBufId = (wBufId + 1) % 2;
|
||||
}
|
||||
SetFlag<HardEvent::MTE1_MTE2>(X_EVENT0 + xBufId);
|
||||
xBufId = (xBufId + 1) % 2;
|
||||
}
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif // COMPRESSOR_BLOCK_CUBE_FULL_LOAD_H
|
||||
1369
csrc/attention/compressor/op_kernel/arch35/compressor_block_vec.h
Normal file
1369
csrc/attention/compressor/op_kernel/arch35/compressor_block_vec.h
Normal file
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
445
csrc/attention/compressor/op_kernel/arch35/compressor_comm.h
Normal file
445
csrc/attention/compressor/op_kernel/arch35/compressor_comm.h
Normal file
@@ -0,0 +1,445 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_comm.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_COMM_H
|
||||
#define COMPRESSOR_COMM_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include "kernel_operator_list_tensor_intf.h"
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
#include "lib/matmul_intf.h"
|
||||
#include "lib/matrix/matmul/tiling.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
template <typename T>
|
||||
__aicore__ inline T CeilDivT(T num1, T num2)
|
||||
{
|
||||
if (num2 == 0) {
|
||||
return static_cast<T>(0);
|
||||
}
|
||||
return (num1 + num2 - 1) / num2;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Align(T num, T rnd)
|
||||
{
|
||||
return (((rnd) == 0) ? 0 : (((num) + (rnd)-1) / (rnd) * (rnd)));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T Trunc(T num, T rnd)
|
||||
{
|
||||
return ((rnd) == 0) ? 0 : (((num) / (rnd) * (rnd)));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T FloorPow2(T num)
|
||||
{
|
||||
if (num == 0)
|
||||
return 1;
|
||||
for (uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) {
|
||||
num |= (num >> i);
|
||||
}
|
||||
return num - (num >> 1);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline T CeilPow2(T num)
|
||||
{
|
||||
if (num <= 1)
|
||||
return 1;
|
||||
num--;
|
||||
for (uint32_t i = 1; i < sizeof(T) * 8; i <<= 1) {
|
||||
num |= (num >> i);
|
||||
}
|
||||
num++;
|
||||
return num;
|
||||
}
|
||||
|
||||
enum class X_LAYOUT : std::uint8_t {
|
||||
BSH = static_cast<std::uint8_t>(0),
|
||||
TH = static_cast<std::uint8_t>(1)
|
||||
};
|
||||
|
||||
enum class X_DTYPE : std::uint8_t {
|
||||
BF16 = static_cast<std::uint8_t>(0),
|
||||
FP16 = static_cast<std::uint8_t>(1)
|
||||
};
|
||||
|
||||
enum class COFF : std::uint8_t {
|
||||
DISABLE = static_cast<std::uint8_t>(1),
|
||||
OVERLAP = static_cast<std::uint8_t>(2)
|
||||
};
|
||||
|
||||
enum class ROTARY_MODE : std::uint8_t {
|
||||
HALF = static_cast<std::uint8_t>(1),
|
||||
INTERLEAVE = static_cast<std::uint8_t>(2)
|
||||
};
|
||||
|
||||
enum class CACHE_MODE : std::uint8_t {
|
||||
CONTINUOUS = static_cast<std::uint8_t>(1),
|
||||
CYCLE = static_cast<std::uint8_t>(2)
|
||||
};
|
||||
|
||||
enum class TEMPLATE_ID : uint8_t {
|
||||
NORMAL = 0,
|
||||
EMPTY_X = 1,
|
||||
FULL_LOAD = 2
|
||||
};
|
||||
|
||||
template <X_LAYOUT X_L, X_DTYPE X_T, COFF C, ROTARY_MODE Rotary_Mode, CACHE_MODE Cache_Mode, typename... Args>
|
||||
struct COMPType {
|
||||
static constexpr X_LAYOUT xLayout = X_L;
|
||||
static constexpr X_DTYPE xDtype = X_T;
|
||||
static constexpr COFF coff = C;
|
||||
static constexpr ROTARY_MODE rotaryMode = Rotary_Mode;
|
||||
static constexpr CACHE_MODE cacheMode = Cache_Mode;
|
||||
};
|
||||
|
||||
struct CmpBlockInfo {
|
||||
__aicore__ inline CmpBlockInfo(){};
|
||||
__aicore__ inline CmpBlockInfo(uint32_t bIdx, uint32_t sIdx, bool needReset = false)
|
||||
: bIdx(bIdx), sIdx(sIdx), needReset(needReset){};
|
||||
|
||||
uint32_t bIdx = 0U;
|
||||
uint32_t sIdx = 0U;
|
||||
uint32_t bSeqUsed = 0U;
|
||||
uint32_t bStartPos = 0U;
|
||||
bool needReset = false;
|
||||
bool isFirst = true;
|
||||
|
||||
uint32_t headSeqCnt = 0U;
|
||||
uint32_t validSeqCnt = 0U;
|
||||
uint32_t tailSeqCnt = 0U;
|
||||
bool isCompress = 0U;
|
||||
};
|
||||
|
||||
struct BasicBlockInfo {
|
||||
uint32_t bIdx = 0;
|
||||
uint32_t sIdx = 0;
|
||||
uint32_t compressedTcNum = 0;
|
||||
uint32_t dealSeqCnt = 0;
|
||||
uint32_t dealTcNum = 0;
|
||||
};
|
||||
|
||||
struct BatchInfo {
|
||||
uint32_t tcNum = 0;
|
||||
uint32_t compressedTcNum = 0;
|
||||
uint32_t remSeqCnt = 0;
|
||||
uint32_t seqCnt = 0;
|
||||
uint32_t seqUsedCnt = 0;
|
||||
uint32_t headHolderSeq = 0;
|
||||
uint32_t bStartPos = 0;
|
||||
uint32_t bIdx = 0;
|
||||
uint32_t sIdx = 0;
|
||||
};
|
||||
|
||||
struct ConstInfo {
|
||||
// 整个AICORE的任务信息, 左闭右开区间[ (bStart, s2Start), (bEnd, s2End) )
|
||||
uint32_t bStart = 0U;
|
||||
uint32_t sStart = 0U;
|
||||
uint32_t bEnd = 0U;
|
||||
uint32_t sEnd = 0U;
|
||||
|
||||
// 分核相关
|
||||
uint32_t usedCoreNum = 0;
|
||||
uint32_t dBaseSize = 0;
|
||||
uint32_t mBaseSize = 0;
|
||||
uint32_t kBaseSize = 0;
|
||||
uint32_t kBaseNum = 0;
|
||||
uint32_t tcSize = 0;
|
||||
uint32_t tcBaseSize = 0;
|
||||
uint32_t tcBasicBlockNum = 0;
|
||||
uint32_t dBasicBlockNum = 0;
|
||||
uint32_t coreGroupNum = 0;
|
||||
uint32_t singleCoreDealTcBasicNum = 0;
|
||||
uint32_t dIdx = 0;
|
||||
uint32_t mStart = 0;
|
||||
uint32_t mEnd = 0;
|
||||
uint32_t nStart = 0;
|
||||
uint32_t nEnd = 0;
|
||||
uint32_t kStart = 0;
|
||||
uint32_t kEnd = 0;
|
||||
uint32_t mLoopNum = 0;
|
||||
uint32_t bIdxOfLastTc = 0;
|
||||
uint32_t sIdxOfLastTc = 0;
|
||||
uint32_t mGroupNum = 0;
|
||||
uint32_t mCurGroupIdx = 0;
|
||||
|
||||
// shape及参数
|
||||
uint32_t batchSize = 0;
|
||||
uint32_t hSize = 0;
|
||||
uint32_t sSize = 0;
|
||||
uint32_t headDim = 0;
|
||||
uint32_t ropeHeadDim = 0;
|
||||
uint32_t cmpRatio = 0;
|
||||
float normEps = 1e-6;
|
||||
float reciprocalD = 0;
|
||||
uint64_t stateCacheStrideDim0 = 0;
|
||||
|
||||
uint32_t curGroupIdx = 0;
|
||||
uint32_t tailGroupIdx = 0;
|
||||
uint32_t tailBasicBlockNum = 0;
|
||||
uint32_t realDealBasicBlockNum = 0;
|
||||
|
||||
// pageAttention
|
||||
uint32_t blockNum = 0;
|
||||
uint32_t blockSize = 0;
|
||||
uint32_t maxBlockNumPerBatch = 0;
|
||||
|
||||
// workSpace
|
||||
uint32_t dbWorkspaceRatio = 1;
|
||||
uint32_t mm1KvResSize = 0;
|
||||
uint32_t mm1ScoreResSize = 0;
|
||||
uint32_t vec1TailCacheSize = 0;
|
||||
uint32_t vec1ResSize = 0;
|
||||
uint32_t mm1ResSize = 0; // 所有cube输出kv/score结果的总大小
|
||||
|
||||
uint32_t aiCoreIdx = 0;
|
||||
uint32_t nSize = 0;
|
||||
|
||||
uint32_t dbSize = 0;
|
||||
};
|
||||
|
||||
struct RunInfo {
|
||||
bool isValid = false;
|
||||
uint32_t cubeDbIdx = 0; // kernel主循环索引
|
||||
|
||||
// 增加字段
|
||||
uint32_t dealTcNum = 0;
|
||||
// 右边相关信息
|
||||
uint32_t bStart = 0;
|
||||
uint32_t sStart = 0;
|
||||
uint32_t dealSeqCnt = 0;
|
||||
// 左边相关信息
|
||||
uint32_t preBStart = 0;
|
||||
uint32_t preSStart = 0;
|
||||
uint32_t preDealSeqCnt = 0; // 左边需要处理的s大小
|
||||
uint32_t preFirstSeqCnt = 0; // 左边首块大小
|
||||
|
||||
uint32_t kStartIdx = 0;
|
||||
uint32_t dealKSize = 0;
|
||||
uint32_t hStart = 0;
|
||||
|
||||
uint32_t bEnd = 0;
|
||||
uint32_t sEnd = 0;
|
||||
uint32_t bStartSeqIdx = 0;
|
||||
uint32_t bEndSeqIdx = 0;
|
||||
|
||||
// v2分核信息 sc是左闭右开
|
||||
uint32_t scStart = 0;
|
||||
uint32_t scEnd = 0;
|
||||
uint32_t dealScSize = 0;
|
||||
|
||||
// vec1Res offset
|
||||
uint64_t vec1ResOffset = 0;
|
||||
};
|
||||
|
||||
struct Vec1RunInfo {
|
||||
// vec相关信息,一次syncAll需处理数据的起始索引
|
||||
bool resetResFlag = false; // v1积攒N轮 是否是N轮的起始轮
|
||||
uint32_t c1v1DbIdx = 0; // vec1 doubleBuffer索引
|
||||
uint32_t v1v2DbIdx = 0; // v1v2 doubleBuffer索引
|
||||
uint32_t bStart = 0;
|
||||
uint32_t sStart = 0;
|
||||
uint32_t dealTcNum = 0;
|
||||
uint32_t dealScSize = 0;
|
||||
};
|
||||
|
||||
struct Vec2RunInfo {
|
||||
// uint32_t bStart = 0;
|
||||
uint32_t v2DbIdx = 0; // v2 doubleBuffer索引
|
||||
uint32_t sStart = 0;
|
||||
uint32_t bEnd = 0;
|
||||
uint32_t sEnd = 0;
|
||||
// v2分核信息 sc是左闭右开
|
||||
uint32_t scStart = 0;
|
||||
uint32_t scEnd = 0;
|
||||
|
||||
// 增加字段
|
||||
uint32_t bStart = 0;
|
||||
uint32_t compressedId = 0;
|
||||
uint32_t bCompressedId = 0;
|
||||
uint32_t dealScSize = 0;
|
||||
};
|
||||
|
||||
struct MSplitInfo {
|
||||
uint32_t vecStartB = 0U;
|
||||
uint32_t vecStartS = 0U;
|
||||
uint32_t vecEndB = 0U;
|
||||
uint32_t vecEndS = 0U;
|
||||
uint32_t dealTcNum = 0U;
|
||||
// vec1Res offset
|
||||
uint64_t vec1StartOffset = 0;
|
||||
uint64_t vec1ResOffset = 0;
|
||||
};
|
||||
|
||||
struct BlockInfo {
|
||||
__aicore__ inline BlockInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealSeqSize)
|
||||
: bIdx(bIdx), sIdx(sIdx), dealSeqSize(dealSeqSize){};
|
||||
uint32_t bIdx = 0U;
|
||||
uint32_t sIdx = 0U;
|
||||
uint32_t dealSeqSize = 0;
|
||||
|
||||
uint32_t isFirst = true;
|
||||
uint32_t bSeqUsed = 0U;
|
||||
uint32_t bStartPos = 0U;
|
||||
uint32_t headHolderSeqCnt = 0U;
|
||||
uint32_t validSeqCnt = 0U;
|
||||
uint32_t tailHolderSeqCnt = 0U;
|
||||
uint32_t dealTcSize = 0U;
|
||||
uint32_t tailValidSeqCnt = 0U;
|
||||
uint32_t compressTcSize = 0U;
|
||||
};
|
||||
|
||||
struct LoopInfo {
|
||||
uint32_t groupSize = 0U;
|
||||
uint32_t groupNum = 0U;
|
||||
uint32_t coreRowIdx = 0U;
|
||||
uint32_t coreColIdx = 0U;
|
||||
uint32_t dLoopIdx = 0U;
|
||||
bool isCoreRowFirst = false;
|
||||
bool isCoreRowLast = false;
|
||||
bool isCoreLoopFirst = false;
|
||||
bool isCoreLoopLast = false;
|
||||
};
|
||||
|
||||
struct Vec1SplitInfo {
|
||||
uint32_t dealSeqStartIdx = 0;
|
||||
uint32_t dealSeqCnt = 0;
|
||||
uint32_t dBaseSize = 0;
|
||||
uint32_t vec1GroupSize = 0;
|
||||
uint32_t vec1GroupNum = 0;
|
||||
uint32_t dealTcSize = 0;
|
||||
uint32_t dealTcNum = 0;
|
||||
uint32_t dealBatchNum = 0;
|
||||
uint32_t preDealTcSize = 0;
|
||||
uint32_t preDealBatchNum = 0;
|
||||
uint32_t curBStart = 0;
|
||||
uint32_t curSStart = 0;
|
||||
uint32_t curCompressedCnt = 0;
|
||||
uint32_t preCompressedCnt = 0;
|
||||
uint32_t totalCompressedCnt = 0;
|
||||
uint32_t tcSplitSize = 0;
|
||||
uint32_t dSplitSize = 0;
|
||||
uint32_t dLoopCount = 0;
|
||||
};
|
||||
|
||||
struct Vec2SplitInfo {
|
||||
uint32_t dealedScCnt = 0;
|
||||
uint32_t preScCnt = 0;
|
||||
uint32_t dealScNum = 0;
|
||||
uint32_t curBStart = 0;
|
||||
uint32_t curScStart = 0;
|
||||
};
|
||||
|
||||
// BUFFER的字节数
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_32B = 32;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_64B = 64;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_256B = 256;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_512B = 512;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_1K = 1024;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_2K = 2048;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_4K = 4096;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_8K = 8192;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_16K = 16384;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_32K = 32768;
|
||||
inline constexpr uint32_t BUFFER_SIZE_BYTE_64K = 65536;
|
||||
|
||||
// BLOCK和REPEAT的字节数
|
||||
inline constexpr uint64_t BYTE_BLOCK = 32UL;
|
||||
inline constexpr uint32_t REPEAT_BLOCK_BYTE = 256U;
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline constexpr T BlockElementNum()
|
||||
{
|
||||
return BYTE_BLOCK / sizeof(T);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline constexpr T RepeatElementNum()
|
||||
{
|
||||
return REPEAT_BLOCK_BYTE / sizeof(T);
|
||||
}
|
||||
// BLOCK和REPEAT的FP32元素数
|
||||
inline constexpr uint32_t FP32_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(float); // 8
|
||||
inline constexpr uint32_t FP16_BLOCK_ELEMENT_NUM = BYTE_BLOCK / sizeof(bfloat16_t); // 16
|
||||
inline constexpr uint32_t FP32_REPEAT_ELEMENT_NUM = REPEAT_BLOCK_BYTE / sizeof(float); // 64
|
||||
inline constexpr uint32_t REPEAT_STRIDE_NUM = REPEAT_BLOCK_BYTE / BYTE_BLOCK; // 8
|
||||
inline constexpr uint32_t REPEAT_MAX_NUM = 255;
|
||||
inline constexpr uint32_t BRCB_NUM = 8;
|
||||
inline constexpr uint32_t MAX_R = 256;
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void CopySingleMatrixNDToNZ(LocalTensor<T> l1Tensor, const GlobalTensor<T> gmTensor, uint32_t nValue,
|
||||
uint32_t dValue, uint32_t srcDValue, uint32_t dstNzC0Stride)
|
||||
{
|
||||
Nd2NzParams nd2nzPara;
|
||||
nd2nzPara.ndNum = 1;
|
||||
nd2nzPara.nValue = nValue; // nd矩阵的行数
|
||||
if constexpr (IsSameType<T, int4b_t>::value) {
|
||||
constexpr uint32_t HALF_SIZE_DIVISOR = 2;
|
||||
nd2nzPara.dValue = dValue / HALF_SIZE_DIVISOR;
|
||||
nd2nzPara.srcDValue = srcDValue / HALF_SIZE_DIVISOR;
|
||||
} else {
|
||||
nd2nzPara.dValue = dValue; // nd矩阵的列数
|
||||
nd2nzPara.srcDValue = srcDValue; // 同一nd矩阵相邻行起始地址间的偏移
|
||||
}
|
||||
nd2nzPara.dstNzC0Stride = dstNzC0Stride;
|
||||
nd2nzPara.dstNzNStride = 1;
|
||||
nd2nzPara.srcNdMatrixStride = 0;
|
||||
nd2nzPara.dstNzMatrixStride = 0;
|
||||
DataCopy(l1Tensor, gmTensor, nd2nzPara);
|
||||
}
|
||||
template <typename T>
|
||||
__aicore__ inline void DumpTensorForDim2(GlobalTensor<T> tensor, uint32_t desc, uint32_t dumpSize, uint32_t row,
|
||||
uint32_t col)
|
||||
{
|
||||
uint32_t array2[] = {static_cast<uint32_t>(row), static_cast<uint32_t>(col)};
|
||||
AscendC::ShapeInfo shapeInfo(2, array2);
|
||||
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DumpTensorForDim2(LocalTensor<T> tensor, uint32_t desc, uint32_t dumpSize, uint32_t row,
|
||||
uint32_t col)
|
||||
{
|
||||
uint32_t array2[] = {static_cast<uint32_t>(row), static_cast<uint32_t>(col)};
|
||||
AscendC::ShapeInfo shapeInfo(2, array2);
|
||||
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DumpTensorForDim2(LocalTensor<T> tensor, uint32_t desc, uint32_t dumpSize)
|
||||
{
|
||||
uint32_t col = 32 / sizeof(T);
|
||||
uint32_t array2[] = {static_cast<uint32_t>(dumpSize / col), static_cast<uint32_t>(col)};
|
||||
AscendC::ShapeInfo shapeInfo(2, array2);
|
||||
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void DumpTensorForDim2(GlobalTensor<T> tensor, uint32_t desc, uint32_t dumpSize)
|
||||
{
|
||||
uint32_t col = 32 / sizeof(T);
|
||||
uint32_t array2[] = {static_cast<uint32_t>(dumpSize / col), static_cast<uint32_t>(col)};
|
||||
AscendC::ShapeInfo shapeInfo(2, array2);
|
||||
// AscendC::DumpTensor(tensor, desc, dumpSize, shapeInfo);
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
#endif
|
||||
628
csrc/attention/compressor/op_kernel/arch35/compressor_kernel.h
Normal file
628
csrc/attention/compressor/op_kernel/arch35/compressor_kernel.h
Normal file
@@ -0,0 +1,628 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_kernel.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_KERNEL_H
|
||||
#define COMPRESSOR_KERNEL_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_template_tiling_key.h"
|
||||
#include "compressor_tiling_data.h"
|
||||
#include "compressor_tools.h"
|
||||
#include "compressor_block_cube.h"
|
||||
#include "compressor_block_vec.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorKernel {
|
||||
public:
|
||||
__aicore__ inline CompressorKernel(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData)
|
||||
: pipe_(pipe), tilingData_(tilingData) {}
|
||||
|
||||
__aicore__ inline void Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *workspace);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
// ================================Init functions==================================
|
||||
__aicore__ inline void InitWorkspace(__gm__ uint8_t *workspace);
|
||||
// ================================Process functions================================
|
||||
__aicore__ inline void InitTilingData();
|
||||
__aicore__ inline void SplitK();
|
||||
// 获取基本块数量
|
||||
__aicore__ inline uint32_t GetLoopTimes();
|
||||
__aicore__ inline void SkipInvalidBatch(BatchInfo &batchInfo);
|
||||
__aicore__ inline void UpdateCurGroup(BasicBlockInfo &basicBlockInfo, BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq);
|
||||
__aicore__ inline BasicBlockInfo SkipOneLoop(BatchInfo &batchInfo);
|
||||
// 计算分核基本信息
|
||||
__aicore__ inline void CalcSplitCoreInfo();
|
||||
|
||||
__aicore__ inline void AllocEventID();
|
||||
__aicore__ inline void FreeEventID();
|
||||
__aicore__ inline void ComputeMm1(const RunInfo &info, bool isNeedExcute);
|
||||
__aicore__ inline void ComputeVec1(const Vec1RunInfo &info);
|
||||
__aicore__ inline void ComputeVec2(const Vec2RunInfo &info);
|
||||
|
||||
__aicore__ inline bool IsNeedExcuteC1(RunInfo info);
|
||||
__aicore__ inline bool IsNeedSyncAll(uint32_t curBasicBlockIdx);
|
||||
__aicore__ inline void CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx);
|
||||
__aicore__ inline void UpdateVec2Info(Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info);
|
||||
__aicore__ inline bool IsNeedExcuteV2(Vec2RunInfo &vec2Info);
|
||||
|
||||
using X_T = typename AscendC::Conditional<COMP::xDtype == X_DTYPE::BF16, bfloat16_t, half>::type;
|
||||
using T = float;
|
||||
using MM1_OUT_T = T;
|
||||
using VEC1_OUT_T = T;
|
||||
|
||||
// 常量
|
||||
static constexpr uint64_t SYNC_MODE0 = 0;
|
||||
static constexpr uint64_t SYNC_MODE2 = 2;
|
||||
static constexpr uint32_t SYNC_C1_FLAG = 3;
|
||||
static constexpr uint32_t SYNC_V1_FLAG = 4;
|
||||
static constexpr uint32_t SYNC_V1_FLAG2 = 5;
|
||||
static constexpr uint32_t SYNC_C1_V1_FLAG = 7;
|
||||
static constexpr uint32_t SYNC_V1_C1_FLAG = 9;
|
||||
|
||||
// ==============================TilingData&TPipe==============================
|
||||
TPipe* pipe_;
|
||||
const optiling::CompressorTilingData* __restrict tilingData_;
|
||||
// ===========================Workspace Global Tensor===========================
|
||||
GlobalTensor<MM1_OUT_T> mm1KvResGm;
|
||||
GlobalTensor<MM1_OUT_T> mm1ScoreResGm;
|
||||
GlobalTensor<MM1_OUT_T> vec1KvCacheGm;
|
||||
GlobalTensor<MM1_OUT_T> vec1ScoreCacheGm;
|
||||
GlobalTensor<MM1_OUT_T> Vec1InputKvGm;
|
||||
GlobalTensor<MM1_OUT_T> Vec1InputScoreGm;
|
||||
GlobalTensor<VEC1_OUT_T> vec1ResGm;
|
||||
GlobalTensor<VEC1_OUT_T> vec2InputGm;
|
||||
// ================================Task Info====================================
|
||||
CompressorTools<COMP> tools_;
|
||||
ConstInfo constInfo{};
|
||||
uint32_t aiCoreIdx = 0;
|
||||
|
||||
// ==============================Service Define==============================
|
||||
CompressorBlockCube<COMP> blockCube_;
|
||||
CompressorBlockVector<COMP> blockVec_;
|
||||
|
||||
uint32_t allCompressedTcNum_ = 0;
|
||||
uint32_t curCompressedTcNum_ = 0;
|
||||
uint32_t accDealSize = 0;
|
||||
uint32_t loopTimes = 0;
|
||||
uint32_t cubeLoop = 0;
|
||||
uint32_t vec1Loop = 0;
|
||||
uint32_t vec2Loop = 0;
|
||||
uint32_t kStartIdx_ = 0;
|
||||
uint32_t dealKSize_ = 0;
|
||||
uint32_t hStart_ = 0;
|
||||
bool isFirstUpdateCurGroup = true;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *workspace)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
constInfo.aiCoreIdx = GetBlockIdx() / 2;
|
||||
} else {
|
||||
constInfo.aiCoreIdx = GetBlockIdx();
|
||||
}
|
||||
InitTilingData();
|
||||
// init tools
|
||||
tools_.toolParams_.seqSize = tilingData_->baseParams.seqSize;
|
||||
tools_.toolParams_.cmpRatio = tilingData_->baseParams.cmpRatio;
|
||||
tools_.Init(startPos, seqUsed, cuSeqlens);
|
||||
|
||||
// 剔除尾部的无效batch
|
||||
for (; constInfo.batchSize > 0; --constInfo.batchSize) {
|
||||
uint32_t bSeqUsed = tools_.GetSeqLength(constInfo.batchSize - 1);
|
||||
if (bSeqUsed > 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 所有batch的有效序列都为0时, 直接退出
|
||||
if (constInfo.batchSize == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// 0. 计算最后一个Tc块的起始位置
|
||||
constInfo.bIdxOfLastTc = constInfo.batchSize - 1;
|
||||
// 1. 计算head_dim的切分大小, 构建ConstInfo的其他信息
|
||||
CalcSplitCoreInfo();
|
||||
SplitK(); // K轴切分
|
||||
// 2. 计算循环次数
|
||||
loopTimes = GetLoopTimes();
|
||||
// 3. 初始化workspace
|
||||
InitWorkspace(workspace);
|
||||
// 4. 初始化block层
|
||||
if ASCEND_IS_AIC {
|
||||
blockCube_.InitParams(constInfo, tools_);
|
||||
blockCube_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos,
|
||||
stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut);
|
||||
blockCube_.InitBuffers(pipe_);
|
||||
blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);
|
||||
} else {
|
||||
blockVec_.InitParams(constInfo, tools_);
|
||||
blockVec_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable,
|
||||
cuSeqlens, seqUsed, startPos, cmpKvOut);
|
||||
blockVec_.InitBuffers(pipe_);
|
||||
blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::InitTilingData() {
|
||||
constInfo.cmpRatio = tilingData_->baseParams.cmpRatio;
|
||||
constInfo.batchSize = tilingData_->baseParams.batchSize;
|
||||
constInfo.mBaseSize = tilingData_->innerSplitParams.mBaseSize;
|
||||
constInfo.dBaseSize = tilingData_->innerSplitParams.dBaseSize;
|
||||
constInfo.kBaseSize = tilingData_->baseParams.hiddenSize;
|
||||
constInfo.kBaseNum = 1;
|
||||
constInfo.headDim = tilingData_->baseParams.headDim;
|
||||
constInfo.hSize = tilingData_->baseParams.hiddenSize;
|
||||
constInfo.sSize = tilingData_->baseParams.seqSize;
|
||||
constInfo.ropeHeadDim = tilingData_->baseParams.ropeHeadDim;
|
||||
constInfo.normEps = tilingData_->baseParams.normEps;
|
||||
constInfo.stateCacheStrideDim0 = tilingData_->baseParams.stateCacheStrideDim0;
|
||||
constInfo.reciprocalD = tilingData_->baseParams.reciprocalD;
|
||||
constInfo.usedCoreNum = tilingData_->baseParams.usedCoreNum;
|
||||
|
||||
constInfo.blockNum = tilingData_->pageAttentionParams.blockNum;
|
||||
constInfo.blockSize = tilingData_->pageAttentionParams.blockSize;
|
||||
constInfo.maxBlockNumPerBatch = tilingData_->pageAttentionParams.maxBlockNumPerBatch;
|
||||
|
||||
constInfo.nSize = tilingData_->baseParams.nSize;
|
||||
constInfo.vec1TailCacheSize = tilingData_->workspaceParams.vec1TailCacheSize;
|
||||
constInfo.dbWorkspaceRatio = tilingData_->workspaceParams.dbWorkspaceRatio;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::SplitK()
|
||||
{
|
||||
uint32_t mSize = 0;
|
||||
for (uint32_t i = 0; i < constInfo.batchSize; i++) {
|
||||
uint32_t bSeqUsed = tools_.GetSeqLength(i);
|
||||
// 获取m大小
|
||||
mSize += bSeqUsed;
|
||||
}
|
||||
|
||||
uint32_t mBaseNum = CeilDivT(mSize, constInfo.mBaseSize);
|
||||
if (constInfo.dBasicBlockNum * mBaseNum < constInfo.usedCoreNum) {
|
||||
constInfo.kBaseNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum;
|
||||
uint32_t kAlignSize =
|
||||
CeilDivT(Align(constInfo.hSize, static_cast<uint32_t>(BUFFER_SIZE_BYTE_32B / sizeof(X_T))),
|
||||
constInfo.kBaseNum);
|
||||
constInfo.kBaseSize = Trunc(kAlignSize, static_cast<uint32_t>(BUFFER_SIZE_BYTE_32B / sizeof(X_T)));
|
||||
// 当切m轴无法满足开满核时,不切m轴(切m处理有点复杂)
|
||||
constInfo.mGroupNum = 1; // 在m轴处理上所有核当一个组
|
||||
constInfo.mCurGroupIdx = 0; // 只有一个组
|
||||
}
|
||||
// 每轮固定不变,预计算后主循环直接复用
|
||||
if (constInfo.kBaseNum > 1) {
|
||||
kStartIdx_ = constInfo.aiCoreIdx / constInfo.dBasicBlockNum;
|
||||
if (constInfo.curGroupIdx + 1 < constInfo.coreGroupNum) {
|
||||
dealKSize_ = constInfo.kBaseSize;
|
||||
hStart_ = kStartIdx_ * dealKSize_;
|
||||
} else {
|
||||
dealKSize_ = kStartIdx_ < constInfo.coreGroupNum ?
|
||||
constInfo.hSize - kStartIdx_ * constInfo.kBaseSize : 0;
|
||||
hStart_ = kStartIdx_ * constInfo.kBaseSize;
|
||||
}
|
||||
} else {
|
||||
kStartIdx_ = 0;
|
||||
dealKSize_ = constInfo.hSize;
|
||||
hStart_ = kStartIdx_ * dealKSize_;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::SkipInvalidBatch(BatchInfo &batchInfo)
|
||||
{
|
||||
for (; batchInfo.bIdx < constInfo.batchSize; ++batchInfo.bIdx) {
|
||||
batchInfo.seqCnt = tools_.GetSeqLength(batchInfo.bIdx);
|
||||
if (batchInfo.seqCnt > 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
batchInfo.remSeqCnt = batchInfo.seqCnt;
|
||||
if (tools_.isExistSeqUsed_) {
|
||||
batchInfo.seqUsedCnt = tools_.GetSeqUsed(batchInfo.bIdx);
|
||||
} else {
|
||||
batchInfo.seqUsedCnt = batchInfo.seqCnt;
|
||||
}
|
||||
if (batchInfo.bIdx < constInfo.batchSize) {
|
||||
batchInfo.bStartPos = tools_.GetStartPos(batchInfo.bIdx);
|
||||
batchInfo.sIdx = 0;
|
||||
batchInfo.headHolderSeq = batchInfo.bStartPos & (constInfo.cmpRatio - 1);
|
||||
batchInfo.tcNum = (batchInfo.bStartPos + batchInfo.seqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio;
|
||||
batchInfo.compressedTcNum = (batchInfo.bStartPos + batchInfo.seqUsedCnt) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio;
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::UpdateCurGroup(BasicBlockInfo &basicBlockInfo,
|
||||
BatchInfo batchInfo, uint32_t &curGroupQuota, uint32_t curDealSeq)
|
||||
{
|
||||
// 更新当前组的信息
|
||||
if (curGroupQuota == 0 && !isFirstUpdateCurGroup) {
|
||||
return;
|
||||
}
|
||||
isFirstUpdateCurGroup = false;
|
||||
basicBlockInfo.bIdx = batchInfo.bIdx;
|
||||
uint32_t curGroupDealSeq = curGroupQuota < curDealSeq ? curGroupQuota : curDealSeq;
|
||||
basicBlockInfo.sIdx = batchInfo.sIdx + curGroupDealSeq;
|
||||
basicBlockInfo.dealSeqCnt += curGroupDealSeq;
|
||||
curGroupQuota -= curGroupDealSeq;
|
||||
// 结尾需要跳batch,需要考虑在当前组起始为末尾,或者当前组起始大于整个M轴
|
||||
if ((curGroupQuota == 0 || basicBlockInfo.bIdx == constInfo.batchSize - 1) && basicBlockInfo.sIdx == batchInfo.seqCnt) {
|
||||
basicBlockInfo.sIdx = 0;
|
||||
for (basicBlockInfo.bIdx++; basicBlockInfo.bIdx < constInfo.batchSize; ++basicBlockInfo.bIdx) {
|
||||
uint32_t seqCnt = tools_.GetSeqLength(basicBlockInfo.bIdx);
|
||||
if (seqCnt > 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline BasicBlockInfo CompressorKernel<COMP>::SkipOneLoop(BatchInfo &batchInfo)
|
||||
{
|
||||
BasicBlockInfo basicBlockInfo{};
|
||||
isFirstUpdateCurGroup = true;
|
||||
uint32_t curGroupQuota = constInfo.mBaseSize * constInfo.mCurGroupIdx; // m轴当前组起始
|
||||
bool curGroupStartFlag = false;
|
||||
uint32_t quota = constInfo.mGroupNum * constInfo.mBaseSize;
|
||||
|
||||
for (; batchInfo.bIdx < constInfo.batchSize;) {
|
||||
uint32_t curDealSeq = 0;
|
||||
uint32_t curDealTcNum = 0;
|
||||
uint32_t curDealCompressedTcNum = 0;
|
||||
// 无法处理完当前整个batch
|
||||
if (quota < batchInfo.remSeqCnt) {
|
||||
// 向下对齐r,
|
||||
uint32_t alignSeq = constInfo.cmpRatio;
|
||||
if (batchInfo.bIdx == 0) {
|
||||
alignSeq = constInfo.cmpRatio - batchInfo.headHolderSeq;
|
||||
}
|
||||
if (quota > alignSeq) {
|
||||
uint32_t delta = (batchInfo.bStartPos + batchInfo.sIdx + quota) & (constInfo.cmpRatio - 1); // 超出对齐的部分
|
||||
curDealSeq = quota - delta;
|
||||
quota -= curDealSeq;
|
||||
curDealTcNum = (curDealSeq + constInfo.cmpRatio - 1) / constInfo.cmpRatio;
|
||||
curDealCompressedTcNum = min(curDealTcNum, batchInfo.compressedTcNum);
|
||||
// 更新当前组所需信息
|
||||
UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq);
|
||||
// 更新batch信息
|
||||
batchInfo.remSeqCnt = batchInfo.remSeqCnt - curDealSeq;
|
||||
batchInfo.sIdx = batchInfo.sIdx + curDealSeq;
|
||||
batchInfo.compressedTcNum -= curDealCompressedTcNum;
|
||||
batchInfo.tcNum -= curDealTcNum;
|
||||
// 更新loop信息
|
||||
basicBlockInfo.dealTcNum += curDealTcNum;
|
||||
basicBlockInfo.compressedTcNum += curDealCompressedTcNum;
|
||||
}
|
||||
break;
|
||||
} else {
|
||||
// 处理整个batch
|
||||
quota -= batchInfo.remSeqCnt;
|
||||
curDealSeq = batchInfo.remSeqCnt;
|
||||
curDealTcNum = batchInfo.tcNum;
|
||||
// 更新当前组所需信息
|
||||
UpdateCurGroup(basicBlockInfo, batchInfo, curGroupQuota, curDealSeq);
|
||||
// 更新batch和loop信息
|
||||
batchInfo.remSeqCnt = 0;
|
||||
basicBlockInfo.dealTcNum += batchInfo.tcNum;
|
||||
basicBlockInfo.compressedTcNum += batchInfo.compressedTcNum;
|
||||
batchInfo.bIdx++;
|
||||
SkipInvalidBatch(batchInfo);
|
||||
}
|
||||
}
|
||||
uint32_t totalDataSize = constInfo.mGroupNum * constInfo.mBaseSize - quota;
|
||||
// 2. 当前组的起始偏移
|
||||
uint32_t currentGroupStart = constInfo.mCurGroupIdx * constInfo.mBaseSize;
|
||||
|
||||
// 3. 安全判断
|
||||
if (currentGroupStart >= totalDataSize) {
|
||||
// 超出尾块
|
||||
basicBlockInfo.dealSeqCnt = 0;
|
||||
} else {
|
||||
// 还在有效范围内,计算剩余量
|
||||
uint32_t remaining = totalDataSize - currentGroupStart;
|
||||
basicBlockInfo.dealSeqCnt = (remaining < constInfo.mBaseSize) ? remaining : constInfo.mBaseSize;
|
||||
}
|
||||
return basicBlockInfo;
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorKernel<COMP>::GetLoopTimes()
|
||||
{
|
||||
// 计算主循环次数
|
||||
uint32_t loopTimes = 0;
|
||||
BatchInfo batchInfo{};
|
||||
SkipInvalidBatch(batchInfo);
|
||||
for (;batchInfo.bIdx < constInfo.batchSize; ++loopTimes) {
|
||||
SkipOneLoop(batchInfo);
|
||||
}
|
||||
return loopTimes;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::CalcSplitCoreInfo()
|
||||
{
|
||||
// D方向的基本块数量
|
||||
constInfo.dBasicBlockNum = constInfo.headDim / constInfo.dBaseSize;
|
||||
// 核的组数
|
||||
constInfo.coreGroupNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum;
|
||||
// 每个核处理的d方向的索引
|
||||
constInfo.dIdx = constInfo.aiCoreIdx % constInfo.dBasicBlockNum;
|
||||
// 当前组id
|
||||
constInfo.curGroupIdx = constInfo.aiCoreIdx / constInfo.dBasicBlockNum;
|
||||
constInfo.mGroupNum = constInfo.coreGroupNum;
|
||||
constInfo.mCurGroupIdx = constInfo.curGroupIdx;
|
||||
|
||||
constInfo.mm1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.coreGroupNum;
|
||||
|
||||
uint32_t coff = (uint32_t)COMP::coff;
|
||||
constInfo.mm1KvResSize = constInfo.mBaseSize * constInfo.headDim * coff;
|
||||
constInfo.mm1ScoreResSize = constInfo.mBaseSize * constInfo.headDim * coff;
|
||||
constInfo.vec1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.nSize;
|
||||
|
||||
constInfo.dbSize = constInfo.coreGroupNum * constInfo.mm1KvResSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::InitWorkspace(__gm__ uint8_t *workspace) {
|
||||
uint64_t offset = 0;
|
||||
uint64_t mm1KvResStartOffset = offset;
|
||||
// mm1KvResGm
|
||||
mm1KvResGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + offset +
|
||||
constInfo.curGroupIdx * constInfo.mm1KvResSize * sizeof(MM1_OUT_T)));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1KvResSize * sizeof(MM1_OUT_T);
|
||||
|
||||
uint64_t mm1ScoreResStartOffset = offset;
|
||||
// mm1ScoreResGm
|
||||
mm1ScoreResGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + offset +
|
||||
constInfo.curGroupIdx * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T)));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T);
|
||||
|
||||
Vec1InputKvGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + mm1KvResStartOffset));
|
||||
|
||||
Vec1InputScoreGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + mm1ScoreResStartOffset));
|
||||
|
||||
vec1KvCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T);
|
||||
|
||||
vec1ScoreCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T);
|
||||
|
||||
uint64_t beforeVecOffset = offset;
|
||||
|
||||
// vec1Res
|
||||
vec1ResGm.SetGlobalBuffer(
|
||||
(__gm__ VEC1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.vec1ResSize * sizeof(VEC1_OUT_T);
|
||||
// vec2Input
|
||||
vec2InputGm.SetGlobalBuffer(
|
||||
(__gm__ VEC1_OUT_T *)(workspace + beforeVecOffset));
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::ComputeMm1(const RunInfo &info, bool isNeedExcute) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_FIX>(SYNC_V1_C1_FLAG + info.cubeDbIdx);
|
||||
if (isNeedExcute) {
|
||||
blockCube_.ComputeMm1(info);
|
||||
}
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_FIX>(SYNC_C1_FLAG);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_FIX>(SYNC_C1_FLAG);
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_FIX>(SYNC_C1_V1_FLAG + info.cubeDbIdx);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::ComputeVec1(const Vec1RunInfo &info) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_C1_V1_FLAG + info.c1v1DbIdx);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG2 + info.c1v1DbIdx);
|
||||
blockVec_.ComputeVec1(info);
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG);
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_V1_C1_FLAG + info.c1v1DbIdx);
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE3>(SYNC_V1_FLAG2 + (info.c1v1DbIdx + 1) % constInfo.dbWorkspaceRatio);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::ComputeVec2(const Vec2RunInfo &info) {
|
||||
blockVec_.ComputeVec2(info);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::AllocEventID()
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
blockCube_.AllocEventID(pipe_);
|
||||
} else {
|
||||
blockVec_.AllocEventID();
|
||||
for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) {
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_V1_C1_FLAG + i);
|
||||
}
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE3>(SYNC_V1_FLAG2);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::FreeEventID()
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_FIX>(SYNC_V1_C1_FLAG + i);
|
||||
}
|
||||
blockCube_.FreeEventID(pipe_);
|
||||
} else {
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG2 + loopTimes % constInfo.dbWorkspaceRatio);
|
||||
blockVec_.FreeEventID();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernel<COMP>::IsNeedExcuteC1(RunInfo info)
|
||||
{
|
||||
// B超出范围则cube不执行
|
||||
return info.bStart < constInfo.batchSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::CalcC1V1Params(RunInfo &info, Vec1RunInfo &vec1Info, BatchInfo &batchInfo, uint32_t loopIdx)
|
||||
{
|
||||
vec1Info.bStart = batchInfo.bIdx;
|
||||
vec1Info.sStart = batchInfo.sIdx;
|
||||
vec1Info.resetResFlag = (loopIdx & (constInfo.nSize - 1)) == 0;
|
||||
vec1Info.c1v1DbIdx = (vec1Loop++ & (constInfo.dbWorkspaceRatio - 1));
|
||||
vec1Info.v1v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1));
|
||||
BasicBlockInfo basicBlockInfo = SkipOneLoop(batchInfo);
|
||||
info.cubeDbIdx = (cubeLoop++ & (constInfo.dbWorkspaceRatio - 1));
|
||||
info.dealSeqCnt = basicBlockInfo.dealSeqCnt;
|
||||
info.dealTcNum = basicBlockInfo.dealTcNum;
|
||||
info.bStart = basicBlockInfo.bIdx;
|
||||
info.sStart = basicBlockInfo.sIdx;
|
||||
info.kStartIdx = kStartIdx_;
|
||||
info.dealKSize = dealKSize_;
|
||||
info.hStart = hStart_;
|
||||
vec1Info.dealTcNum = basicBlockInfo.dealTcNum;
|
||||
vec1Info.dealScSize = basicBlockInfo.compressedTcNum;
|
||||
allCompressedTcNum_ += basicBlockInfo.compressedTcNum;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernel<COMP>::IsNeedExcuteV2(Vec2RunInfo &vec2Info)
|
||||
{
|
||||
return (vec2Info.dealScSize > 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernel<COMP>::IsNeedSyncAll(uint32_t curBasicBlockIdx)
|
||||
{
|
||||
if (allCompressedTcNum_ == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint32_t cnt = curBasicBlockIdx + 1;
|
||||
if ((cnt == loopTimes) || (cnt % constInfo.nSize == 0)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::UpdateVec2Info(
|
||||
Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info)
|
||||
{
|
||||
// nSize轮起始先重置v2Info信息
|
||||
if (curBasicBlockIdx % constInfo.nSize == 0) {
|
||||
vec2Info.v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1));
|
||||
vec2Info.bStart = info.bStart;
|
||||
vec2Info.sStart = info.sStart;
|
||||
// 将sStart转成bCompressedId
|
||||
uint32_t startPos = tools_.GetStartPos(info.bStart);
|
||||
if (tools_.isExistSeqUsed_) {
|
||||
uint32_t seqUsed = tools_.GetSeqUsed(info.bStart);
|
||||
if (vec2Info.sStart >= seqUsed) {
|
||||
vec2Info.bStart++;
|
||||
vec2Info.sStart = 0;
|
||||
}
|
||||
}
|
||||
vec2Info.bCompressedId = (startPos + vec2Info.sStart) / constInfo.cmpRatio - startPos / constInfo.cmpRatio;
|
||||
|
||||
vec2Info.dealScSize = 0;
|
||||
} else if ((curBasicBlockIdx + 1) % constInfo.nSize == 0) {
|
||||
vec2Loop++;
|
||||
}
|
||||
vec2Info.dealScSize += info.dealScSize;
|
||||
vec2Info.compressedId += info.dealScSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernel<COMP>::Process()
|
||||
{
|
||||
// 所有batch的有效序列都为0时, 直接退出
|
||||
if (constInfo.batchSize == 0) {
|
||||
return;
|
||||
}
|
||||
AllocEventID();
|
||||
|
||||
BatchInfo batchInfo{};
|
||||
|
||||
RunInfo extraInfo[1];
|
||||
Vec1RunInfo vec1Info{};
|
||||
Vec2RunInfo vec2Info{};
|
||||
SkipInvalidBatch(batchInfo);
|
||||
for (uint32_t i = 0; i < loopTimes; ++i) {
|
||||
RunInfo &extraInfo0 = extraInfo[0];
|
||||
CalcC1V1Params(extraInfo0, vec1Info, batchInfo, i);
|
||||
bool isNeedExcuteC1 = IsNeedExcuteC1(extraInfo0);
|
||||
|
||||
if ASCEND_IS_AIC {
|
||||
ComputeMm1(extraInfo0, isNeedExcuteC1);
|
||||
} else {
|
||||
ComputeVec1(vec1Info);
|
||||
UpdateVec2Info(vec2Info, i, vec1Info);
|
||||
|
||||
if (IsNeedSyncAll(i)) {
|
||||
SyncAll();
|
||||
if (IsNeedExcuteV2(vec2Info)) {
|
||||
ComputeVec2(vec2Info);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
FreeEventID();
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif // COMPRESSOR_KERNEL_H
|
||||
@@ -0,0 +1,474 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_kernel_full_load.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_KERNEL_FULL_LOAD_H
|
||||
#define COMPRESSOR_KERNEL_FULL_LOAD_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
#include "compressor_template_tiling_key.h"
|
||||
#include "compressor_tiling_data.h"
|
||||
#include "compressor_tools.h"
|
||||
#include "compressor_block_cube_full_load.h"
|
||||
#include "compressor_block_vec_full_load.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
template <typename COMP>
|
||||
class CompressorKernelFullLoad {
|
||||
public:
|
||||
__aicore__ inline CompressorKernelFullLoad(TPipe* pipe, const optiling::CompressorTilingData* __restrict tilingData)
|
||||
: pipe_(pipe), tilingData_(tilingData) {}
|
||||
|
||||
__aicore__ inline void Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *workspace);
|
||||
__aicore__ inline void Process();
|
||||
|
||||
private:
|
||||
// ================================Init functions==================================
|
||||
__aicore__ inline void InitWorkspace(__gm__ uint8_t *workspace);
|
||||
// ================================Process functions================================
|
||||
__aicore__ inline void InitTilingData();
|
||||
// 获取基本块数量
|
||||
__aicore__ inline void SkipInvalidBatch(BatchInfo &batchInfo);
|
||||
// 计算分核基本信息
|
||||
__aicore__ inline void CalcSplitCoreInfo();
|
||||
|
||||
__aicore__ inline void AllocEventID();
|
||||
__aicore__ inline void FreeEventID();
|
||||
__aicore__ inline void ComputeMm1(const RunInfo &info);
|
||||
__aicore__ inline void ComputeVec1(const Vec1RunInfo &info);
|
||||
__aicore__ inline void ComputeVec2(const Vec2RunInfo &info);
|
||||
|
||||
__aicore__ inline bool IsNeedExcuteC1(RunInfo info);
|
||||
__aicore__ inline bool IsNeedSyncAll(uint32_t curBasicBlockIdx);
|
||||
__aicore__ inline void UpdateVec2Info(Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info);
|
||||
__aicore__ inline bool IsNeedExcuteV2(Vec2RunInfo &vec2Info);
|
||||
__aicore__ inline void CalcCubeParams(RunInfo &info);
|
||||
__aicore__ inline void CalcV1Params(Vec1RunInfo &vec1Info);
|
||||
|
||||
using X_T = typename AscendC::Conditional<COMP::xDtype == X_DTYPE::BF16, bfloat16_t, half>::type;
|
||||
using T = float;
|
||||
using MM1_OUT_T = T;
|
||||
using VEC1_OUT_T = T;
|
||||
|
||||
// 常量
|
||||
static constexpr uint64_t SYNC_MODE0 = 0;
|
||||
static constexpr uint64_t SYNC_MODE2 = 2;
|
||||
static constexpr uint32_t SYNC_C1_FLAG = 3;
|
||||
static constexpr uint32_t SYNC_V1_FLAG = 4;
|
||||
static constexpr uint32_t SYNC_V1_FLAG2 = 5;
|
||||
static constexpr uint32_t SYNC_C1_V1_FLAG = 7;
|
||||
static constexpr uint32_t SYNC_V1_C1_FLAG = 9;
|
||||
|
||||
// ==============================TilingData&TPipe==============================
|
||||
TPipe* pipe_;
|
||||
const optiling::CompressorTilingData* __restrict tilingData_;
|
||||
// ===========================Workspace Global Tensor===========================
|
||||
GlobalTensor<MM1_OUT_T> mm1KvResGm;
|
||||
GlobalTensor<MM1_OUT_T> mm1ScoreResGm;
|
||||
GlobalTensor<MM1_OUT_T> vec1KvCacheGm;
|
||||
GlobalTensor<MM1_OUT_T> vec1ScoreCacheGm;
|
||||
GlobalTensor<MM1_OUT_T> Vec1InputKvGm;
|
||||
GlobalTensor<MM1_OUT_T> Vec1InputScoreGm;
|
||||
GlobalTensor<VEC1_OUT_T> vec1ResGm;
|
||||
GlobalTensor<VEC1_OUT_T> vec2InputGm;
|
||||
// ================================Task Info====================================
|
||||
CompressorTools<COMP> tools_;
|
||||
ConstInfo constInfo{};
|
||||
uint32_t aiCoreIdx = 0;
|
||||
|
||||
// ==============================Service Define==============================
|
||||
CompressorBlockCubeFullLoad<COMP> blockCube_;
|
||||
CompressorBlockVectorFullLoad<COMP> blockVec_;
|
||||
|
||||
uint32_t allCompressedTcNum_ = 0;
|
||||
uint32_t curCompressedTcNum_ = 0;
|
||||
uint32_t accDealSize = 0;
|
||||
uint32_t loopTimes = 0;
|
||||
uint32_t cubeLoop = 0;
|
||||
uint32_t vec1Loop = 0;
|
||||
uint32_t vec2Loop = 0;
|
||||
uint32_t kStartIdx_ = 0;
|
||||
uint32_t dealKSize_ = 0;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::Init(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *workspace)
|
||||
{
|
||||
if ASCEND_IS_AIV {
|
||||
constInfo.aiCoreIdx = GetBlockIdx() / 2;
|
||||
} else {
|
||||
constInfo.aiCoreIdx = GetBlockIdx();
|
||||
}
|
||||
|
||||
InitTilingData();
|
||||
// init tools
|
||||
tools_.toolParams_.seqSize = tilingData_->baseParams.seqSize;
|
||||
tools_.toolParams_.cmpRatio = tilingData_->baseParams.cmpRatio;
|
||||
tools_.Init(startPos, seqUsed, cuSeqlens);
|
||||
|
||||
// 剔除尾部的无效batch
|
||||
for (; constInfo.batchSize > 0; --constInfo.batchSize) {
|
||||
uint32_t bSeqUsed = tools_.GetSeqLength(constInfo.batchSize - 1);
|
||||
if (bSeqUsed > 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// 所有batch的有效序列都为0时, 直接退出
|
||||
if (constInfo.batchSize == 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// 0. 计算最后一个Tc块的起始位置
|
||||
constInfo.bIdxOfLastTc = constInfo.batchSize - 1;
|
||||
// 1. 计算head_dim的切分大小, 构建ConstInfo的其他信息
|
||||
CalcSplitCoreInfo();
|
||||
// 2. 计算循环次数
|
||||
loopTimes = 1;
|
||||
// 3. 初始化workspace
|
||||
InitWorkspace(workspace);
|
||||
// 4. 初始化block层
|
||||
if ASCEND_IS_AIC {
|
||||
blockCube_.InitParams(constInfo, tools_);
|
||||
blockCube_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos,
|
||||
stateBlockTable, cuSeqlens, seqUsed, startPos, cmpKvOut);
|
||||
blockCube_.InitBuffers(pipe_);
|
||||
blockCube_.InitGlobalBuffers(mm1KvResGm, mm1ScoreResGm);
|
||||
} else {
|
||||
blockVec_.InitParams(constInfo, tools_);
|
||||
blockVec_.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable,
|
||||
cuSeqlens, seqUsed, startPos, cmpKvOut);
|
||||
blockVec_.InitBuffers(pipe_);
|
||||
blockVec_.InitVec1GlobalTensor(Vec1InputKvGm, Vec1InputScoreGm, vec1KvCacheGm, vec1ScoreCacheGm, vec1ResGm, vec2InputGm);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::InitTilingData() {
|
||||
constInfo.cmpRatio = tilingData_->baseParams.cmpRatio;
|
||||
constInfo.batchSize = tilingData_->baseParams.batchSize;
|
||||
constInfo.mBaseSize = tilingData_->innerSplitParams.mBaseSize;
|
||||
constInfo.dBaseSize = tilingData_->innerSplitParams.dBaseSize;
|
||||
constInfo.headDim = tilingData_->baseParams.headDim;
|
||||
constInfo.hSize = tilingData_->baseParams.hiddenSize;
|
||||
constInfo.sSize = tilingData_->baseParams.seqSize;
|
||||
constInfo.ropeHeadDim = tilingData_->baseParams.ropeHeadDim;
|
||||
constInfo.normEps = tilingData_->baseParams.normEps;
|
||||
constInfo.stateCacheStrideDim0 = tilingData_->baseParams.stateCacheStrideDim0;
|
||||
constInfo.reciprocalD = tilingData_->baseParams.reciprocalD;
|
||||
constInfo.usedCoreNum = tilingData_->baseParams.usedCoreNum;
|
||||
|
||||
constInfo.blockNum = tilingData_->pageAttentionParams.blockNum;
|
||||
constInfo.blockSize = tilingData_->pageAttentionParams.blockSize;
|
||||
constInfo.maxBlockNumPerBatch = tilingData_->pageAttentionParams.maxBlockNumPerBatch;
|
||||
|
||||
constInfo.nSize = tilingData_->baseParams.nSize;
|
||||
constInfo.vec1TailCacheSize = tilingData_->workspaceParams.vec1TailCacheSize;
|
||||
constInfo.dbWorkspaceRatio = tilingData_->workspaceParams.dbWorkspaceRatio;
|
||||
|
||||
constInfo.mStart = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].mStart;
|
||||
constInfo.mEnd = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].mEnd;
|
||||
constInfo.nStart = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].nStart;
|
||||
constInfo.nEnd = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].nEnd;
|
||||
constInfo.kStart = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].kStart;
|
||||
constInfo.kEnd = tilingData_->baseParams.splitCoreParam[constInfo.aiCoreIdx].kEnd;
|
||||
constInfo.mLoopNum = tilingData_->baseParams.mLoopNum;
|
||||
constInfo.kBaseNum = tilingData_->baseParams.kBaseNum;
|
||||
constInfo.kBaseSize = tilingData_->baseParams.kBaseSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::SkipInvalidBatch(BatchInfo &batchInfo)
|
||||
{
|
||||
for (; batchInfo.bIdx < constInfo.batchSize; ++batchInfo.bIdx) {
|
||||
batchInfo.seqCnt = tools_.GetSeqLength(batchInfo.bIdx);
|
||||
if (batchInfo.seqCnt > 0) {
|
||||
break;
|
||||
}
|
||||
}
|
||||
batchInfo.remSeqCnt = batchInfo.seqCnt;
|
||||
if (tools_.isExistSeqUsed_) {
|
||||
batchInfo.seqUsedCnt = tools_.GetSeqUsed(batchInfo.bIdx);
|
||||
} else {
|
||||
batchInfo.seqUsedCnt = batchInfo.seqCnt;
|
||||
}
|
||||
if (batchInfo.bIdx < constInfo.batchSize) {
|
||||
batchInfo.bStartPos = tools_.GetStartPos(batchInfo.bIdx);
|
||||
batchInfo.sIdx = 0;
|
||||
batchInfo.headHolderSeq = batchInfo.bStartPos & (constInfo.cmpRatio - 1);
|
||||
batchInfo.tcNum = (batchInfo.bStartPos + batchInfo.seqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio;
|
||||
batchInfo.compressedTcNum = (batchInfo.bStartPos + batchInfo.seqUsedCnt) / constInfo.cmpRatio - batchInfo.bStartPos / constInfo.cmpRatio;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::CalcSplitCoreInfo()
|
||||
{
|
||||
// D方向的基本块数量
|
||||
constInfo.dBasicBlockNum = constInfo.headDim / constInfo.dBaseSize;
|
||||
// 核的组数
|
||||
constInfo.coreGroupNum = constInfo.usedCoreNum / constInfo.dBasicBlockNum;
|
||||
// 当前组id
|
||||
constInfo.curGroupIdx = constInfo.aiCoreIdx / constInfo.dBasicBlockNum;
|
||||
constInfo.mGroupNum = constInfo.coreGroupNum;
|
||||
constInfo.mCurGroupIdx = constInfo.curGroupIdx;
|
||||
|
||||
uint32_t coff = (uint32_t)COMP::coff;
|
||||
constInfo.mm1KvResSize = constInfo.mBaseSize * constInfo.headDim * coff;
|
||||
constInfo.mm1ScoreResSize = constInfo.mBaseSize * constInfo.headDim * coff;
|
||||
constInfo.vec1ResSize = constInfo.mBaseSize * constInfo.headDim * constInfo.nSize;
|
||||
|
||||
constInfo.dbSize = constInfo.coreGroupNum * constInfo.mm1KvResSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::InitWorkspace(__gm__ uint8_t *workspace) {
|
||||
uint64_t offset = 0;
|
||||
uint64_t mm1KvResStartOffset = offset;
|
||||
// mm1KvResGm
|
||||
mm1KvResGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + offset +
|
||||
constInfo.curGroupIdx * constInfo.mm1KvResSize * sizeof(MM1_OUT_T)));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1KvResSize * sizeof(MM1_OUT_T);
|
||||
|
||||
uint64_t mm1ScoreResStartOffset = offset;
|
||||
// mm1ScoreResGm
|
||||
mm1ScoreResGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + offset +
|
||||
constInfo.curGroupIdx * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T)));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.mm1ScoreResSize * sizeof(MM1_OUT_T);
|
||||
|
||||
Vec1InputKvGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + mm1KvResStartOffset));
|
||||
|
||||
Vec1InputScoreGm.SetGlobalBuffer(
|
||||
(__gm__ MM1_OUT_T *)(workspace + mm1ScoreResStartOffset));
|
||||
|
||||
vec1KvCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T);
|
||||
|
||||
vec1ScoreCacheGm.SetGlobalBuffer((__gm__ MM1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.vec1TailCacheSize * sizeof(MM1_OUT_T);
|
||||
|
||||
uint64_t beforeVecOffset = offset;
|
||||
|
||||
// vec1Res
|
||||
vec1ResGm.SetGlobalBuffer(
|
||||
(__gm__ VEC1_OUT_T *)(workspace + offset));
|
||||
offset += constInfo.dbWorkspaceRatio * constInfo.coreGroupNum * constInfo.vec1ResSize * sizeof(VEC1_OUT_T);
|
||||
// vec2Input
|
||||
vec2InputGm.SetGlobalBuffer(
|
||||
(__gm__ VEC1_OUT_T *)(workspace + beforeVecOffset));
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::ComputeMm1(const RunInfo &info) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_FIX>(SYNC_V1_C1_FLAG + info.cubeDbIdx);
|
||||
blockCube_.ComputeMm1(info);
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_FIX>(SYNC_C1_FLAG);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_FIX>(SYNC_C1_FLAG);
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_FIX>(SYNC_C1_V1_FLAG + info.cubeDbIdx);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::ComputeVec1(const Vec1RunInfo &info) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_C1_V1_FLAG + info.c1v1DbIdx);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG2 + info.c1v1DbIdx);
|
||||
blockVec_.ComputeVec1();
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG);
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG);
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_V1_C1_FLAG + info.c1v1DbIdx);
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE3>(SYNC_V1_FLAG2 + (info.c1v1DbIdx + 1) % constInfo.dbWorkspaceRatio);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::ComputeVec2(const Vec2RunInfo &info) {
|
||||
blockVec_.ComputeVec2();
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::AllocEventID()
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
blockCube_.AllocEventID(pipe_);
|
||||
} else {
|
||||
blockVec_.AllocEventID();
|
||||
for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) {
|
||||
CrossCoreSetFlag<SYNC_MODE2, PIPE_MTE2>(SYNC_V1_C1_FLAG + i);
|
||||
}
|
||||
CrossCoreSetFlag<SYNC_MODE0, PIPE_MTE3>(SYNC_V1_FLAG2);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::FreeEventID()
|
||||
{
|
||||
if ASCEND_IS_AIC {
|
||||
for (int i = 0; i < constInfo.dbWorkspaceRatio; ++i) {
|
||||
CrossCoreWaitFlag<SYNC_MODE2, PIPE_FIX>(SYNC_V1_C1_FLAG + i);
|
||||
}
|
||||
blockCube_.FreeEventID(pipe_);
|
||||
} else {
|
||||
CrossCoreWaitFlag<SYNC_MODE0, PIPE_MTE2>(SYNC_V1_FLAG2 + loopTimes % constInfo.dbWorkspaceRatio);
|
||||
blockVec_.FreeEventID();
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernelFullLoad<COMP>::IsNeedExcuteC1(RunInfo info)
|
||||
{
|
||||
// B超出范围则cube不执行
|
||||
return info.bStart < constInfo.batchSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::CalcCubeParams(RunInfo &info)
|
||||
{
|
||||
info.cubeDbIdx = (cubeLoop++ & (constInfo.dbWorkspaceRatio - 1));
|
||||
info.dealSeqCnt = constInfo.mEnd - constInfo.mStart;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::CalcV1Params(Vec1RunInfo &vec1Info)
|
||||
{
|
||||
vec1Info.bStart = 0;
|
||||
vec1Info.sStart = 0;
|
||||
vec1Info.resetResFlag = false;
|
||||
vec1Info.c1v1DbIdx = (vec1Loop++ & (constInfo.dbWorkspaceRatio - 1));
|
||||
vec1Info.v1v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1));
|
||||
for (uint32_t bIdx = 0; bIdx < constInfo.batchSize; ++bIdx) {
|
||||
uint64_t bSeqCnt = tools_.GetSeqLength(bIdx);
|
||||
if (bSeqCnt == 0) {
|
||||
continue;
|
||||
}
|
||||
uint64_t bStartPos = tools_.GetStartPos(bIdx);
|
||||
vec1Info.dealTcNum += (bStartPos + bSeqCnt + constInfo.cmpRatio - 1) / constInfo.cmpRatio - bStartPos / constInfo.cmpRatio;
|
||||
vec1Info.dealScSize += (bStartPos + bSeqCnt) / constInfo.cmpRatio - bStartPos / constInfo.cmpRatio;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernelFullLoad<COMP>::IsNeedExcuteV2(Vec2RunInfo &vec2Info)
|
||||
{
|
||||
return (vec2Info.dealScSize > 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorKernelFullLoad<COMP>::IsNeedSyncAll(uint32_t curBasicBlockIdx)
|
||||
{
|
||||
if (allCompressedTcNum_ == 0) {
|
||||
return false;
|
||||
}
|
||||
|
||||
uint32_t cnt = curBasicBlockIdx + 1;
|
||||
if ((cnt == loopTimes) || (cnt % constInfo.nSize == 0)) {
|
||||
return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::UpdateVec2Info(
|
||||
Vec2RunInfo &vec2Info, uint32_t curBasicBlockIdx, const Vec1RunInfo &info)
|
||||
{
|
||||
// nSize轮起始先重置v2Info信息
|
||||
if (curBasicBlockIdx % constInfo.nSize == 0) {
|
||||
vec2Info.v2DbIdx = (vec2Loop & (constInfo.dbWorkspaceRatio - 1));
|
||||
vec2Info.bStart = info.bStart;
|
||||
vec2Info.sStart = info.sStart;
|
||||
// 将sStart转成bCompressedId
|
||||
uint32_t startPos = tools_.GetStartPos(info.bStart);
|
||||
if (tools_.isExistSeqUsed_) {
|
||||
uint32_t seqUsed = tools_.GetSeqUsed(info.bStart);
|
||||
if (vec2Info.sStart >= seqUsed) {
|
||||
vec2Info.bStart++;
|
||||
vec2Info.sStart = 0;
|
||||
}
|
||||
}
|
||||
vec2Info.bCompressedId = (startPos + vec2Info.sStart) / constInfo.cmpRatio - startPos / constInfo.cmpRatio;
|
||||
|
||||
vec2Info.dealScSize = 0;
|
||||
} else if ((curBasicBlockIdx + 1) % constInfo.nSize == 0) {
|
||||
vec2Loop++;
|
||||
}
|
||||
vec2Info.dealScSize += info.dealScSize;
|
||||
vec2Info.compressedId += info.dealScSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorKernelFullLoad<COMP>::Process()
|
||||
{
|
||||
// 所有batch的有效序列都为0时, 直接退出
|
||||
if (constInfo.batchSize == 0) {
|
||||
return;
|
||||
}
|
||||
AllocEventID();
|
||||
|
||||
RunInfo extraInfo[1];
|
||||
Vec1RunInfo vec1Info{};
|
||||
Vec2RunInfo vec2Info{};
|
||||
for (uint32_t i = 0; i < loopTimes; ++i) {
|
||||
RunInfo &extraInfo0 = extraInfo[0];
|
||||
if ASCEND_IS_AIV {
|
||||
CalcV1Params(vec1Info);
|
||||
} else {
|
||||
CalcCubeParams(extraInfo0);
|
||||
}
|
||||
if ASCEND_IS_AIC {
|
||||
ComputeMm1(extraInfo0);
|
||||
} else {
|
||||
ComputeVec1(vec1Info);
|
||||
UpdateVec2Info(vec2Info, i, vec1Info);
|
||||
SyncAll();
|
||||
if (IsNeedExcuteV2(vec2Info)) {
|
||||
ComputeVec2(vec2Info);
|
||||
}
|
||||
}
|
||||
}
|
||||
FreeEventID();
|
||||
}
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif // COMPRESSOR_KERNEL_FULL_LOAD_H
|
||||
@@ -0,0 +1,54 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file COMPRESSOR_template_tiling_key.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_TEMPLATE_TILING_KEY_H
|
||||
#define COMPRESSOR_TEMPLATE_TILING_KEY_H
|
||||
|
||||
#include "ascendc/host_api/tiling/template_argument.h"
|
||||
|
||||
#define ASCENDC_TPL_1_BW 1 // 每个参数占用1个bit位
|
||||
#define ASCENDC_TPL_2_BW 2 // 每个参数占用2个bit位
|
||||
#define ASCENDC_TPL_4_BW 4 // 每个参数占用4个bit位
|
||||
|
||||
// 可表示的tilingkey范围为64bit,注意不可超过限制
|
||||
ASCENDC_TPL_ARGS_DECL(compressor, // 算子唯一标识,与opType保持一致
|
||||
// 可能需要切分之后的headdim
|
||||
// bit:0 LAYOUT 0:BSH 1:TH
|
||||
ASCENDC_TPL_UINT_DECL(X_LAYOUT, ASCENDC_TPL_1_BW, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
// bit:1-4 x的dtype 0:BF16 1:FP16
|
||||
ASCENDC_TPL_UINT_DECL(X_DTYPE, ASCENDC_TPL_4_BW, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
// bit:5-6 coff 1:无需overlap 2:需要overlap
|
||||
ASCENDC_TPL_UINT_DECL(COFF, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
// bit:7-8 rotary_mode 1:half 2:interleave
|
||||
ASCENDC_TPL_UINT_DECL(ROTARY_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
// bit:9-10 cache_mode 1:CONTINUOUS 2:cycle
|
||||
ASCENDC_TPL_UINT_DECL(CACHE_MODE, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
// bit:11-12 template_id 0:empty_tensor 1:normal 2:full load
|
||||
ASCENDC_TPL_UINT_DECL(TEMPLATE_ID, ASCENDC_TPL_2_BW, ASCENDC_TPL_UI_LIST, 0, 1, 2),
|
||||
|
||||
);
|
||||
|
||||
ASCENDC_TPL_SEL(
|
||||
|
||||
ASCENDC_TPL_ARGS_SEL(ASCENDC_TPL_UINT_SEL(X_LAYOUT, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
ASCENDC_TPL_UINT_SEL(X_DTYPE, ASCENDC_TPL_UI_LIST, 0, 1),
|
||||
ASCENDC_TPL_UINT_SEL(COFF, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
ASCENDC_TPL_UINT_SEL(ROTARY_MODE, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
ASCENDC_TPL_UINT_SEL(CACHE_MODE, ASCENDC_TPL_UI_LIST, 1, 2),
|
||||
ASCENDC_TPL_UINT_SEL(TEMPLATE_ID, ASCENDC_TPL_UI_LIST, 0, 1, 2),
|
||||
ASCENDC_TPL_TILING_STRUCT_SEL(optiling::CompressorTilingData)),
|
||||
);
|
||||
|
||||
#endif // COMPRESSOR_TEMPLATE_TILING_KEY_H
|
||||
@@ -0,0 +1,83 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file COMPRESSOR_tiling_datay.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_TILING_DATA_H
|
||||
#define COMPRESSOR_TILING_DATA_H
|
||||
#include <cstdint>
|
||||
#include "kernel_tiling/kernel_tiling.h"
|
||||
|
||||
const uint32_t CMP_MAX_AIC_CORE_NUM = 36;
|
||||
|
||||
namespace optiling {
|
||||
struct CompressorSplitCoreParams {
|
||||
uint32_t mStart;
|
||||
uint32_t mEnd;
|
||||
uint32_t nStart;
|
||||
uint32_t nEnd;
|
||||
uint32_t kStart;
|
||||
uint32_t kEnd;
|
||||
};
|
||||
|
||||
// 1. 基础参数结构体
|
||||
struct CompressorBaseParams {
|
||||
uint32_t batchSize = 0; // bastch size(批大小)
|
||||
uint32_t seqSize = 0; // sequence size(kvs大小)
|
||||
uint32_t hiddenSize = 0; // hidden size(隐藏层大小)
|
||||
uint32_t tokenSize = 0; // token size = batchSize * seqSize(token总数:批大小x序列1长度)
|
||||
uint32_t headDim = 0; // head size of kv
|
||||
uint32_t ropeHeadDim = 64; // dim size per rope head 64(单个带RoPE头的维度)
|
||||
uint32_t csSize = 0; // Compress sequence len
|
||||
uint32_t cmpRatio = 4; // Compress ratio
|
||||
uint32_t cgSize = 0; // Compress group size
|
||||
float normEps = 1e-6; // RMSNorm eps
|
||||
float reciprocalD = 0; // 1分之D
|
||||
uint32_t usedCoreNum = 0; // 使用核数
|
||||
uint32_t nSize = 0; // 控制v2积攒的轮数
|
||||
uint64_t stateCacheStrideDim0 = 0; // stateCache第0维的stride
|
||||
uint32_t kBaseNum = 0;
|
||||
uint32_t kBaseSize = 0;
|
||||
uint32_t coreGroupNum = 0;
|
||||
uint32_t mLoopNum = 0;
|
||||
CompressorSplitCoreParams splitCoreParam[CMP_MAX_AIC_CORE_NUM];
|
||||
};
|
||||
|
||||
struct CompressorPageAttentionParams {
|
||||
uint32_t blockNum = 0;
|
||||
uint32_t blockSize = 1;
|
||||
uint32_t maxBlockNumPerBatch = 1;
|
||||
};
|
||||
|
||||
struct CompressorInnerSplitParams {
|
||||
uint32_t mBaseSize;
|
||||
uint32_t dBaseSize;
|
||||
};
|
||||
|
||||
struct CompressorWorkspaceParams {
|
||||
uint32_t mm1KvResSize;
|
||||
uint32_t mm1ScoreResSize;
|
||||
uint32_t vec1ResSize;
|
||||
uint32_t vec1TailCacheSize;
|
||||
uint32_t dbWorkspaceRatio = 1;
|
||||
};
|
||||
|
||||
struct CompressorTilingData {
|
||||
CompressorBaseParams baseParams;
|
||||
CompressorPageAttentionParams pageAttentionParams;
|
||||
CompressorInnerSplitParams innerSplitParams;
|
||||
CompressorWorkspaceParams workspaceParams;
|
||||
};
|
||||
} // optiling
|
||||
|
||||
#endif // COMPRESSOR_TILING_DATA_H
|
||||
890
csrc/attention/compressor/op_kernel/arch35/compressor_tools.h
Normal file
890
csrc/attention/compressor/op_kernel/arch35/compressor_tools.h
Normal file
@@ -0,0 +1,890 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor_tools.h
|
||||
* \brief 放算子都需要、与算子联系紧密、但是又不方便单独独立出来的公共工具
|
||||
*/
|
||||
|
||||
#ifndef COMPRESSOR_TOOLS_H
|
||||
#define COMPRESSOR_TOOLS_H
|
||||
|
||||
#include "compressor_comm.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
namespace Compressor {
|
||||
|
||||
struct ToolsParams {
|
||||
uint32_t seqSize = 0U;
|
||||
uint32_t cmpRatio = 0U;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorTools {
|
||||
public:
|
||||
__aicore__ inline CompressorTools() {}
|
||||
|
||||
__aicore__ inline void Init(__gm__ uint8_t *cuSeqlens, __gm__ uint8_t *seqUsed, __gm__ uint8_t *startPos);
|
||||
|
||||
__aicore__ inline uint32_t GetSeqUsed(uint32_t bIdx);
|
||||
__aicore__ inline uint32_t GetStartPos(uint32_t bIdx);
|
||||
__aicore__ inline uint32_t GetSeqLength(uint32_t bIdx);
|
||||
__aicore__ inline uint32_t GetTIdxByBatch(uint32_t bIdx);
|
||||
|
||||
public:
|
||||
ToolsParams toolParams_ {};
|
||||
bool isExistSeqUsed_ = false;
|
||||
|
||||
private:
|
||||
bool isExistStartPos_ = false;
|
||||
GlobalTensor<int32_t> cuSeqlensGm_;
|
||||
GlobalTensor<int32_t> sequsedGm_;
|
||||
GlobalTensor<int32_t> startPosGm_;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorTools<COMP>::Init(__gm__ uint8_t *startPos, __gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *cuSeqlens)
|
||||
{
|
||||
isExistStartPos_ = (startPos != nullptr);
|
||||
if (isExistStartPos_) {
|
||||
startPosGm_.SetGlobalBuffer((__gm__ int32_t *)startPos);
|
||||
}
|
||||
|
||||
isExistSeqUsed_ = (seqUsed != nullptr);
|
||||
if (isExistSeqUsed_) {
|
||||
sequsedGm_.SetGlobalBuffer((__gm__ int32_t *)seqUsed);
|
||||
}
|
||||
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
cuSeqlensGm_.SetGlobalBuffer((__gm__ int32_t *)cuSeqlens);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorTools<COMP>::GetSeqUsed(uint32_t bIdx)
|
||||
{
|
||||
if (isExistSeqUsed_) {
|
||||
return (uint32_t)sequsedGm_.GetValue(bIdx);
|
||||
} else {
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
return (uint32_t)(cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx));
|
||||
} else {
|
||||
return toolParams_.seqSize;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorTools<COMP>::GetStartPos(uint32_t bIdx)
|
||||
{
|
||||
if (isExistStartPos_) {
|
||||
return (uint32_t)startPosGm_.GetValue(bIdx);
|
||||
} else {
|
||||
return 0;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorTools<COMP>::GetSeqLength(uint32_t bIdx)
|
||||
{
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
return cuSeqlensGm_.GetValue(bIdx + 1) - cuSeqlensGm_.GetValue(bIdx);
|
||||
} else {
|
||||
return toolParams_.seqSize;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorTools<COMP>::GetTIdxByBatch(uint32_t bIdx)
|
||||
{
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
return (uint32_t)(cuSeqlensGm_.GetValue(bIdx));
|
||||
} else {
|
||||
return toolParams_.seqSize * bIdx;
|
||||
}
|
||||
}
|
||||
|
||||
// iterator
|
||||
struct SliceInfo {
|
||||
__aicore__ inline SliceInfo(){};
|
||||
__aicore__ inline SliceInfo(uint32_t bIdx, uint32_t sIdx) : bIdx(bIdx), sIdx(sIdx) {};
|
||||
|
||||
uint32_t bIdx = 0U;
|
||||
uint32_t sIdx = 0U;
|
||||
uint32_t bSeqUsed = 0U;
|
||||
uint32_t bStartPos = 0U;
|
||||
|
||||
uint32_t headHolderSeqCnt = 0U;
|
||||
uint32_t validSeqCnt = 0U;
|
||||
uint32_t tailHolderSeqCnt = 0U;
|
||||
|
||||
uint32_t dealSeqCnt = 0;
|
||||
uint32_t dealTcSize = 0U;
|
||||
uint32_t compressTcSize = 0U;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorSliceIterator {
|
||||
public:
|
||||
__aicore__ inline CompressorSliceIterator(CompressorTools<COMP> &tools) : tools_(tools) {}
|
||||
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
|
||||
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
|
||||
__aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt);
|
||||
__aicore__ inline bool IsEnd();
|
||||
__aicore__ inline void IteratorSlice();
|
||||
__aicore__ inline SliceInfo& GetSlice();
|
||||
__aicore__ inline SliceInfo& GetSliceByCmp();
|
||||
|
||||
bool isFirst_ = true;
|
||||
SliceInfo sliceInfo_{};
|
||||
|
||||
private:
|
||||
CompressorTools<COMP> &tools_;
|
||||
|
||||
// iterator
|
||||
uint32_t maxDealSeqCnt_ = 0;
|
||||
uint32_t batch_size_ = 0;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
|
||||
{
|
||||
sliceInfo_.bIdx = bIdx;
|
||||
sliceInfo_.sIdx = sIdx;
|
||||
isFirst_ = true;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
|
||||
{
|
||||
this->batch_size_ = batch_size;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSliceIterator<COMP>::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt)
|
||||
{
|
||||
this->maxDealSeqCnt_ = maxDealSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorSliceIterator<COMP>::IsEnd()
|
||||
{
|
||||
return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSliceIterator<COMP>::IteratorSlice()
|
||||
{
|
||||
bool isUpdateBatchInfo = false;
|
||||
if (!isFirst_) {
|
||||
// 更新剩余未处理的行数
|
||||
maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt;
|
||||
// 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed
|
||||
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) {
|
||||
sliceInfo_.sIdx = 0;
|
||||
sliceInfo_.bIdx++;
|
||||
isUpdateBatchInfo = true;
|
||||
}
|
||||
} else {
|
||||
isUpdateBatchInfo = true;
|
||||
isFirst_ = false;
|
||||
}
|
||||
|
||||
// 更新与bIdx相关的bStartPos和bSeqUsed
|
||||
if (isUpdateBatchInfo) {
|
||||
// SkipInvalidBatch
|
||||
while (sliceInfo_.bIdx < batch_size_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
if (sliceInfo_.bSeqUsed > 0) {
|
||||
break;
|
||||
}
|
||||
sliceInfo_.bIdx++;
|
||||
}
|
||||
if (sliceInfo_.bIdx < batch_size_) {
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline SliceInfo& CompressorSliceIterator<COMP>::GetSliceByCmp()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (isFirst_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
isFirst_ = false;
|
||||
}
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
|
||||
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
sliceInfo_.tailHolderSeqCnt =
|
||||
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
|
||||
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
}
|
||||
|
||||
// 头和尾处理,否则需要处理的seq等于cmpRatio
|
||||
if (sliceInfo_.validSeqCnt < cmpRatio) {
|
||||
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == 0) {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
} else {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio;
|
||||
}
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
|
||||
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
|
||||
|
||||
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
|
||||
if (sliceInfo_.tailHolderSeqCnt > 0) {
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时,其不需要压缩
|
||||
}
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline SliceInfo& CompressorSliceIterator<COMP>::GetSlice()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (isFirst_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
isFirst_ = false;
|
||||
}
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
|
||||
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
sliceInfo_.tailHolderSeqCnt =
|
||||
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
|
||||
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
}
|
||||
|
||||
sliceInfo_.dealSeqCnt = sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt;
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize = sliceInfo_.dealSeqCnt / cmpRatio;
|
||||
|
||||
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
|
||||
if (sliceInfo_.tailHolderSeqCnt > 0) {
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时,其不需要压缩
|
||||
}
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
struct SplitCoreSliceInfo : public SliceInfo {
|
||||
__aicore__ inline SplitCoreSliceInfo() {};
|
||||
__aicore__ inline SplitCoreSliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx) {};
|
||||
|
||||
uint32_t preFirstSeqCnt = 0U; // 左边每次迭代基本块的第一个seqCnt大小
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorSplitCoreSliceIterator {
|
||||
public:
|
||||
__aicore__ inline CompressorSplitCoreSliceIterator(CompressorTools<COMP> &tools) : tools_(tools) {}
|
||||
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
|
||||
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
|
||||
__aicore__ inline void SetMaxDealSeqCnt(uint32_t maxDealSeqCnt);
|
||||
__aicore__ inline bool IsEnd();
|
||||
__aicore__ inline void IteratorSlice();
|
||||
__aicore__ inline SplitCoreSliceInfo& GetSlice();
|
||||
__aicore__ inline SplitCoreSliceInfo& GetSliceByCmp();
|
||||
__aicore__ inline uint32_t GetBIdx();
|
||||
__aicore__ inline SplitCoreSliceInfo& GetLeftNextCmpSeqCnt();
|
||||
__aicore__ inline SplitCoreSliceInfo& GetRightNextCmpSeqCnt();
|
||||
|
||||
bool isFirst_ = true;
|
||||
bool isLeftFirstBath = false;
|
||||
bool isMaxDealSeqCntFirst = false;
|
||||
|
||||
SplitCoreSliceInfo sliceInfo_{};
|
||||
|
||||
private:
|
||||
CompressorTools<COMP> &tools_;
|
||||
|
||||
// iterator
|
||||
uint32_t maxDealSeqCnt_ = 0;
|
||||
uint32_t batch_size_ = 0;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
|
||||
{
|
||||
sliceInfo_.bIdx = bIdx;
|
||||
sliceInfo_.sIdx = sIdx;
|
||||
isFirst_ = true;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
|
||||
{
|
||||
this->batch_size_ = batch_size;
|
||||
isMaxDealSeqCntFirst = true;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::SetMaxDealSeqCnt(uint32_t maxDealSeqCnt)
|
||||
{
|
||||
this->maxDealSeqCnt_ = maxDealSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorSplitCoreSliceIterator<COMP>::IsEnd()
|
||||
{
|
||||
return (sliceInfo_.bIdx >= batch_size_) || (maxDealSeqCnt_ == 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorSplitCoreSliceIterator<COMP>::GetBIdx()
|
||||
{
|
||||
return sliceInfo_.bIdx;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorSplitCoreSliceIterator<COMP>::IteratorSlice()
|
||||
{
|
||||
bool isUpdateBatchInfo = false;
|
||||
if (isMaxDealSeqCntFirst) {
|
||||
isMaxDealSeqCntFirst = false;
|
||||
}
|
||||
if (!isFirst_) {
|
||||
// 更新剩余未处理的行数
|
||||
maxDealSeqCnt_ -= sliceInfo_.dealSeqCnt;
|
||||
// 更新sIdx和bIdx、以及与bIdx相关的bStartPos和bSeqUsed
|
||||
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == sliceInfo_.bSeqUsed) {
|
||||
sliceInfo_.sIdx = 0;
|
||||
// 左边最后一块跳到b=0 s=0处理
|
||||
if (isLeftFirstBath) {
|
||||
isLeftFirstBath = false;
|
||||
} else {
|
||||
sliceInfo_.bIdx++;
|
||||
}
|
||||
isUpdateBatchInfo = true;
|
||||
}
|
||||
} else {
|
||||
isUpdateBatchInfo = true;
|
||||
isFirst_ = false;
|
||||
}
|
||||
|
||||
// 更新与bIdx相关的bStartPos和bSeqUsed
|
||||
if (isUpdateBatchInfo) {
|
||||
// SkipInvalidBatch
|
||||
while (sliceInfo_.bIdx < batch_size_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
if (sliceInfo_.bSeqUsed > 0) {
|
||||
break;
|
||||
}
|
||||
sliceInfo_.bIdx++;
|
||||
}
|
||||
if (sliceInfo_.bIdx < batch_size_) {
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline SplitCoreSliceInfo& CompressorSplitCoreSliceIterator<COMP>::GetLeftNextCmpSeqCnt()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (isFirst_) {
|
||||
// 左边 T轴首次减去T轴最后一块
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(batch_size_ - 1);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(batch_size_ - 1);
|
||||
// 处理最后一块是中间整块或者尾块的情况
|
||||
uint32_t lastSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio == 0 ?
|
||||
cmpRatio :
|
||||
(sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) % cmpRatio;
|
||||
// 处理最后一块是头块的情况
|
||||
if (sliceInfo_.bSeqUsed < cmpRatio) {
|
||||
lastSeqCnt = sliceInfo_.bSeqUsed;
|
||||
}
|
||||
|
||||
sliceInfo_.sIdx = sliceInfo_.bSeqUsed - lastSeqCnt;
|
||||
isLeftFirstBath = true;
|
||||
isFirst_ = false;
|
||||
}
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
|
||||
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
sliceInfo_.tailHolderSeqCnt =
|
||||
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
|
||||
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
}
|
||||
|
||||
// 头和尾处理,否则需要处理的seq等于cmpRatio
|
||||
if (sliceInfo_.validSeqCnt < cmpRatio) {
|
||||
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == 0) {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
} else {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio;
|
||||
}
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
|
||||
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
|
||||
|
||||
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
|
||||
if (sliceInfo_.tailHolderSeqCnt > 0) {
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时,其不需要压缩
|
||||
}
|
||||
|
||||
// 记录左边第一个块
|
||||
if (isMaxDealSeqCntFirst) {
|
||||
sliceInfo_.preFirstSeqCnt = sliceInfo_.dealSeqCnt;
|
||||
}
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline SplitCoreSliceInfo& CompressorSplitCoreSliceIterator<COMP>::GetRightNextCmpSeqCnt()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (isFirst_) {
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
isFirst_ = false;
|
||||
}
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt > maxDealSeqCnt_) {
|
||||
sliceInfo_.validSeqCnt = maxDealSeqCnt_ - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
sliceInfo_.tailHolderSeqCnt =
|
||||
cmpRatio - (sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt) % cmpRatio;
|
||||
if (sliceInfo_.tailHolderSeqCnt == cmpRatio) {
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
}
|
||||
|
||||
// 头和尾处理,否则需要处理的seq等于cmpRatio
|
||||
if (sliceInfo_.validSeqCnt < cmpRatio) {
|
||||
sliceInfo_.dealSeqCnt = sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx == 0) {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
} else {
|
||||
sliceInfo_.dealSeqCnt = cmpRatio;
|
||||
}
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.dealSeqCnt;
|
||||
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize = (sliceInfo_.dealSeqCnt + cmpRatio - 1) / cmpRatio;
|
||||
|
||||
// 因为是一个batch的数据, 只有最后一个压缩块才可能不需要压缩, 此时sliceInfo_.tailHolderSeqCnt > 0
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize;
|
||||
if (sliceInfo_.tailHolderSeqCnt > 0) {
|
||||
sliceInfo_.compressTcSize = sliceInfo_.dealTcSize - 1; // 最后一个压缩块不满时,其不需要压缩
|
||||
}
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
struct Vec1SliceInfo : public SliceInfo {
|
||||
__aicore__ inline Vec1SliceInfo() {};
|
||||
__aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx) : SliceInfo(bIdx, sIdx) {};
|
||||
__aicore__ inline Vec1SliceInfo(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt)
|
||||
: SliceInfo(bIdx, sIdx), dealedSeqCnt(dealedSeqCnt) {};
|
||||
|
||||
uint32_t dealedSeqCnt = 0U;
|
||||
uint32_t dealedTcCnt = 0U;
|
||||
uint32_t bSeqLength = 0U;
|
||||
uint32_t compressoredScCnt = 0U;
|
||||
bool isFirst = false;
|
||||
bool isLast = false;
|
||||
};
|
||||
|
||||
struct StatisticInfo {
|
||||
__aicore__ inline StatisticInfo() {};
|
||||
__aicore__ inline StatisticInfo(uint32_t actualTcCnt, uint32_t dealSeqCnt, uint32_t compressorScCnt)
|
||||
: actualTcCnt(actualTcCnt), dealSeqCnt(dealSeqCnt), compressorScCnt(compressorScCnt) {};
|
||||
|
||||
uint32_t actualTcCnt = 0U;
|
||||
uint32_t dealSeqCnt = 0U;
|
||||
uint32_t compressorScCnt = 0U;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorVec1SliceIterator {
|
||||
public:
|
||||
__aicore__ inline CompressorVec1SliceIterator(CompressorTools<COMP> &tools) : tools_(tools) {}
|
||||
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx);
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt, uint32_t compressoredScCnt);
|
||||
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
|
||||
__aicore__ inline void SetDealedSeqCnt(uint32_t dealedSeqCnt);
|
||||
__aicore__ inline void SetDealedTcCnt(uint32_t dealedTcCnt);
|
||||
__aicore__ inline void SetCompressoredScCnt(uint32_t compressoredScCnt);
|
||||
__aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize);
|
||||
__aicore__ inline void SetNeedDealTcSize(uint32_t needDealTcSize, uint32_t canDealTcSize);
|
||||
__aicore__ inline uint32_t GetNeedDealTcSize();
|
||||
__aicore__ inline bool IsEnd();
|
||||
template <bool IS_STATISTIC = false>
|
||||
__aicore__ inline void IteratorSlice();
|
||||
__aicore__ inline Vec1SliceInfo &GetSlice();
|
||||
template <bool IS_STATISTIC = false>
|
||||
__aicore__ inline StatisticInfo &FullIteratorSlice();
|
||||
|
||||
private:
|
||||
CompressorTools<COMP> &tools_;
|
||||
|
||||
bool isFirst_ = true;
|
||||
Vec1SliceInfo sliceInfo_{};
|
||||
StatisticInfo statisticInfo_{};
|
||||
uint32_t needDealTcSize_ = 0U;
|
||||
uint32_t batch_size_ = 0U;
|
||||
};
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx)
|
||||
{
|
||||
sliceInfo_.bIdx = bIdx;
|
||||
sliceInfo_.sIdx = sIdx;
|
||||
while (tools_.GetSeqLength(sliceInfo_.bIdx) == 0) {
|
||||
sliceInfo_.bIdx++;
|
||||
if (sliceInfo_.bIdx == batch_size_) {
|
||||
sliceInfo_.bIdx = 0;
|
||||
}
|
||||
}
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
|
||||
isFirst_ = true;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t sIdx, uint32_t dealedSeqCnt,
|
||||
uint32_t compressoredScCnt)
|
||||
{
|
||||
Reset(bIdx, sIdx);
|
||||
SetDealedSeqCnt(dealedSeqCnt);
|
||||
SetCompressoredScCnt(compressoredScCnt);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
|
||||
{
|
||||
this->batch_size_ = batch_size;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetDealedSeqCnt(uint32_t dealedSeqCnt)
|
||||
{
|
||||
this->sliceInfo_.dealedSeqCnt = dealedSeqCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetCompressoredScCnt(uint32_t compressoredScCnt)
|
||||
{
|
||||
this->sliceInfo_.compressoredScCnt = compressoredScCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetDealedTcCnt(uint32_t dealedTcCnt)
|
||||
{
|
||||
this->sliceInfo_.dealedTcCnt = dealedTcCnt;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::SetNeedDealTcSize(uint32_t needDealTcSize)
|
||||
{
|
||||
this->needDealTcSize_ = needDealTcSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
template <bool IS_STATISTIC>
|
||||
__aicore__ inline void CompressorVec1SliceIterator<COMP>::IteratorSlice()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if constexpr (IS_STATISTIC) {
|
||||
statisticInfo_.actualTcCnt += sliceInfo_.dealTcSize;
|
||||
statisticInfo_.compressorScCnt += sliceInfo_.compressTcSize;
|
||||
}
|
||||
needDealTcSize_ -= sliceInfo_.dealTcSize;
|
||||
sliceInfo_.dealedSeqCnt += sliceInfo_.validSeqCnt;
|
||||
sliceInfo_.compressoredScCnt += sliceInfo_.compressTcSize;
|
||||
sliceInfo_.sIdx += sliceInfo_.validSeqCnt;
|
||||
if (sliceInfo_.sIdx >= sliceInfo_.bSeqUsed) {
|
||||
do {
|
||||
uint32_t seqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
|
||||
if (sliceInfo_.bSeqUsed < seqLength) {
|
||||
uint32_t nextAlignSIdx = Align(sliceInfo_.bStartPos + sliceInfo_.sIdx, cmpRatio) - sliceInfo_.bStartPos;
|
||||
sliceInfo_.dealedSeqCnt += nextAlignSIdx - sliceInfo_.sIdx;
|
||||
uint32_t tcGap = CeilDivT(static_cast<int32_t>(seqLength - nextAlignSIdx),
|
||||
static_cast<int32_t>(cmpRatio));
|
||||
if (sliceInfo_.bSeqUsed == 0 && nextAlignSIdx > sliceInfo_.sIdx) {
|
||||
// 此时bseqused所在压缩块未被纳入计算
|
||||
tcGap++;
|
||||
}
|
||||
sliceInfo_.sIdx = nextAlignSIdx;
|
||||
if (needDealTcSize_ < tcGap) {
|
||||
sliceInfo_.dealedSeqCnt += needDealTcSize_ * cmpRatio;
|
||||
sliceInfo_.sIdx += needDealTcSize_ * cmpRatio;
|
||||
needDealTcSize_ = 0;
|
||||
break;
|
||||
}
|
||||
sliceInfo_.dealedSeqCnt += seqLength - sliceInfo_.sIdx;
|
||||
needDealTcSize_ -= tcGap;
|
||||
}
|
||||
sliceInfo_.bIdx++;
|
||||
if (sliceInfo_.bIdx == batch_size_) {
|
||||
sliceInfo_.bIdx = 0;
|
||||
}
|
||||
sliceInfo_.sIdx = 0;
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
} while (sliceInfo_.bSeqUsed == 0);
|
||||
sliceInfo_.bSeqLength = tools_.GetSeqLength(sliceInfo_.bIdx);
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
}
|
||||
if (isFirst_) {
|
||||
isFirst_ = false;
|
||||
}
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorVec1SliceIterator<COMP>::GetNeedDealTcSize()
|
||||
{
|
||||
return needDealTcSize_;
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorVec1SliceIterator<COMP>::IsEnd()
|
||||
{
|
||||
return (needDealTcSize_ == 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline Vec1SliceInfo& CompressorVec1SliceIterator<COMP>::GetSlice()
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
if (sliceInfo_.bSeqUsed < sliceInfo_.sIdx) {
|
||||
sliceInfo_.headHolderSeqCnt = 0;
|
||||
sliceInfo_.validSeqCnt = 0;
|
||||
sliceInfo_.tailHolderSeqCnt = 0;
|
||||
sliceInfo_.dealTcSize = 0;
|
||||
sliceInfo_.compressTcSize = 0;
|
||||
} else {
|
||||
// 计算头部占位行数、有效数据行数、尾部占位行数
|
||||
sliceInfo_.headHolderSeqCnt = (sliceInfo_.bStartPos + sliceInfo_.sIdx) % cmpRatio;
|
||||
sliceInfo_.validSeqCnt = sliceInfo_.bSeqUsed - sliceInfo_.sIdx;
|
||||
if (CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt, cmpRatio) > needDealTcSize_) {
|
||||
sliceInfo_.validSeqCnt = needDealTcSize_ * cmpRatio - sliceInfo_.headHolderSeqCnt;
|
||||
}
|
||||
uint32_t globalTotalSeqCnt = sliceInfo_.bStartPos + sliceInfo_.sIdx + sliceInfo_.validSeqCnt;
|
||||
sliceInfo_.tailHolderSeqCnt = Align(globalTotalSeqCnt, cmpRatio) - globalTotalSeqCnt;
|
||||
|
||||
// 计算本次可以处理的Tc个数
|
||||
sliceInfo_.dealTcSize =
|
||||
(sliceInfo_.headHolderSeqCnt + sliceInfo_.validSeqCnt + sliceInfo_.tailHolderSeqCnt) / cmpRatio;
|
||||
|
||||
sliceInfo_.compressTcSize =
|
||||
(sliceInfo_.headHolderSeqCnt + min(sliceInfo_.validSeqCnt, sliceInfo_.bSeqUsed - sliceInfo_.sIdx)) /
|
||||
cmpRatio;
|
||||
}
|
||||
|
||||
sliceInfo_.isFirst = isFirst_;
|
||||
sliceInfo_.isLast =
|
||||
sliceInfo_.bSeqUsed > sliceInfo_.sIdx &&
|
||||
CeilDivT(sliceInfo_.headHolderSeqCnt + sliceInfo_.bSeqUsed - sliceInfo_.sIdx, cmpRatio) >= needDealTcSize_;
|
||||
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
template <bool IS_STATISTIC>
|
||||
__aicore__ inline StatisticInfo& CompressorVec1SliceIterator<COMP>::FullIteratorSlice()
|
||||
{
|
||||
if constexpr (IS_STATISTIC) {
|
||||
statisticInfo_ = {0U, 0U, 0U};
|
||||
Vec1SliceInfo tempSliceInfo = GetSlice();
|
||||
while (!IsEnd()) {
|
||||
GetSlice();
|
||||
IteratorSlice<IS_STATISTIC>();
|
||||
}
|
||||
Vec1SliceInfo sliceInfo = GetSlice();
|
||||
statisticInfo_.dealSeqCnt = sliceInfo.dealedSeqCnt - tempSliceInfo.dealedSeqCnt;
|
||||
} else {
|
||||
while (!IsEnd()) {
|
||||
GetSlice();
|
||||
IteratorSlice<IS_STATISTIC>();
|
||||
}
|
||||
}
|
||||
return statisticInfo_;
|
||||
}
|
||||
|
||||
struct Vec2SliceInfo{
|
||||
__aicore__ inline Vec2SliceInfo(){};
|
||||
__aicore__ inline Vec2SliceInfo(uint32_t bIdx, uint32_t scIdx) : bIdx(bIdx), scIdx(scIdx)
|
||||
{
|
||||
}
|
||||
|
||||
uint32_t bIdx = 0U;
|
||||
uint32_t scIdx = 0U;
|
||||
uint32_t scNum = 0U;
|
||||
uint32_t remainScCnt = 0U; // 当前batch剩余sc数量
|
||||
uint32_t bStartPos = 0U;
|
||||
uint32_t bSeqUsed = 0U;
|
||||
uint32_t bSeqLength = 0U;
|
||||
uint32_t dealedScCnt = 0U; // 全局的dealedScCnt(Reset刷新)
|
||||
uint32_t curDealScNum = 0U; // 当前循环处理的sc数量(IteratorSlice刷新)
|
||||
uint32_t bOutputScLen = 0U; // BSH场景每个batch填充后的输出长度
|
||||
uint32_t padScIdx = 0U; // 当前sc输出位置,TH场景为全局的dealedScCnt,BSH场景则为填充后全局的索引(Reset刷新)
|
||||
uint32_t loopDealedScCnt = 0U; // 当前迭代已处理的sc数量(Reset刷新)
|
||||
};
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
class CompressorVec2SliceIterator {
|
||||
public:
|
||||
__aicore__ inline CompressorVec2SliceIterator(CompressorTools<COMP> &tools) : tools_(tools)
|
||||
{
|
||||
}
|
||||
__aicore__ inline void Reset(uint32_t bIdx, uint32_t scIdx, uint32_t dealedScCnt);
|
||||
__aicore__ inline void SetMaxBatchSize(uint32_t batch_size);
|
||||
__aicore__ inline void SetNeedDealScSize(uint32_t needDealScSize);
|
||||
__aicore__ inline void ResetLoopDealedScCnt();
|
||||
__aicore__ inline uint32_t GetNeedDealScSize();
|
||||
__aicore__ inline bool IsEnd();
|
||||
__aicore__ inline void IteratorSlice();
|
||||
__aicore__ inline Vec2SliceInfo &GetSlice();
|
||||
private:
|
||||
CompressorTools<COMP> &tools_;
|
||||
|
||||
Vec2SliceInfo sliceInfo_{};
|
||||
uint32_t needDealScSize_ = 0U;
|
||||
uint32_t batch_size_ = 0U;
|
||||
};
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec2SliceIterator<COMP>::Reset(uint32_t bIdx, uint32_t scIdx, uint32_t dealedScCnt)
|
||||
{
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
sliceInfo_.bIdx = bIdx;
|
||||
sliceInfo_.scIdx = scIdx;
|
||||
sliceInfo_.dealedScCnt = dealedScCnt;
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::BSH) {
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.scNum = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) / cmpRatio - sliceInfo_.bStartPos / cmpRatio;
|
||||
sliceInfo_.remainScCnt = sliceInfo_.scNum - sliceInfo_.scIdx;
|
||||
sliceInfo_.bOutputScLen = CeilDivT(tools_.GetSeqLength(sliceInfo_.bIdx), cmpRatio);
|
||||
sliceInfo_.padScIdx = sliceInfo_.bIdx * sliceInfo_.bOutputScLen + sliceInfo_.scIdx;
|
||||
} else {
|
||||
sliceInfo_.padScIdx = sliceInfo_.dealedScCnt;
|
||||
}
|
||||
sliceInfo_.loopDealedScCnt = 0U;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec2SliceIterator<COMP>::ResetLoopDealedScCnt()
|
||||
{
|
||||
sliceInfo_.loopDealedScCnt = 0U;
|
||||
}
|
||||
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec2SliceIterator<COMP>::SetMaxBatchSize(uint32_t batch_size)
|
||||
{
|
||||
this->batch_size_ = batch_size;
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec2SliceIterator<COMP>::SetNeedDealScSize(uint32_t needDealScSize)
|
||||
{
|
||||
this->needDealScSize_ = needDealScSize;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline void CompressorVec2SliceIterator<COMP>::IteratorSlice()
|
||||
{
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
sliceInfo_.padScIdx += sliceInfo_.curDealScNum;
|
||||
} else {
|
||||
if (needDealScSize_ <= sliceInfo_.remainScCnt) {
|
||||
sliceInfo_.scIdx += sliceInfo_.curDealScNum;
|
||||
sliceInfo_.padScIdx += sliceInfo_.curDealScNum;
|
||||
} else {
|
||||
uint32_t cmpRatio = tools_.toolParams_.cmpRatio;
|
||||
sliceInfo_.padScIdx += sliceInfo_.bOutputScLen - sliceInfo_.scIdx;
|
||||
sliceInfo_.bIdx++;
|
||||
sliceInfo_.scIdx = 0;
|
||||
sliceInfo_.bStartPos = tools_.GetStartPos(sliceInfo_.bIdx);
|
||||
sliceInfo_.bSeqUsed = tools_.GetSeqUsed(sliceInfo_.bIdx);
|
||||
sliceInfo_.scNum = (sliceInfo_.bStartPos + sliceInfo_.bSeqUsed) / cmpRatio - sliceInfo_.bStartPos / cmpRatio;
|
||||
}
|
||||
sliceInfo_.remainScCnt = sliceInfo_.scNum - sliceInfo_.scIdx;
|
||||
}
|
||||
sliceInfo_.dealedScCnt += sliceInfo_.curDealScNum;
|
||||
needDealScSize_ -= sliceInfo_.curDealScNum;
|
||||
sliceInfo_.loopDealedScCnt += sliceInfo_.curDealScNum;
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline uint32_t CompressorVec2SliceIterator<COMP>::GetNeedDealScSize()
|
||||
{
|
||||
return needDealScSize_;
|
||||
}
|
||||
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline bool CompressorVec2SliceIterator<COMP>::IsEnd()
|
||||
{
|
||||
return (needDealScSize_ == 0);
|
||||
}
|
||||
|
||||
template <typename COMP>
|
||||
__aicore__ inline Vec2SliceInfo &CompressorVec2SliceIterator<COMP>::GetSlice()
|
||||
{
|
||||
if constexpr (COMP::xLayout == X_LAYOUT::TH) {
|
||||
sliceInfo_.curDealScNum = needDealScSize_;
|
||||
} else {
|
||||
sliceInfo_.curDealScNum = min(sliceInfo_.remainScCnt, needDealScSize_);
|
||||
}
|
||||
return sliceInfo_;
|
||||
}
|
||||
|
||||
|
||||
|
||||
} // namespace Compressor
|
||||
|
||||
#endif
|
||||
378
csrc/attention/compressor/op_kernel/arch35/vf/vf_add.h
Normal file
378
csrc/attention/compressor/op_kernel/arch35/vf/vf_add.h
Normal file
@@ -0,0 +1,378 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_add.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef VF_ADD_H
|
||||
#define VF_ADD_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include <cstdint>
|
||||
using namespace AscendC;
|
||||
constexpr uint32_t FLOAT_REP_SIZE = 64;
|
||||
constexpr uint32_t BTYEALIGNSIZE = 32;
|
||||
constexpr uint32_t REGSIZE = 256;
|
||||
constexpr uint32_t HALFCORED = 128;
|
||||
|
||||
template <typename T>
|
||||
struct AddRegList {
|
||||
MicroAPI::RegTensor<T> vreg;
|
||||
MicroAPI::RegTensor<T> vregape;
|
||||
};
|
||||
|
||||
|
||||
template <typename T>
|
||||
__simd_callee__ void AddVFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, AddRegList<T> ®List, uint32_t row,
|
||||
uint32_t col, uint64_t offset0, uint64_t offset1)
|
||||
{
|
||||
uint32_t maskValue = col;
|
||||
MicroAPI::MaskReg mask = MicroAPI::UpdateMask<T>(maskValue);
|
||||
MicroAPI::LoadAlign(regList.vreg, inputAddr + offset0);
|
||||
MicroAPI::LoadAlign(regList.vregape, apeAddr + offset1);
|
||||
MicroAPI::Add(regList.vreg, regList.vreg, regList.vregape, mask);
|
||||
MicroAPI::StoreAlign(inputAddr + offset0, regList.vreg, mask);
|
||||
}
|
||||
|
||||
template <bool IS_FIRST, typename T>
|
||||
__simd_callee__ void MultiAddVFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, AddRegList<T> ®List, uint32_t row,
|
||||
uint32_t col, uint64_t offset, uint32_t repeatNum, uint64_t repeatOffset)
|
||||
{
|
||||
uint32_t maskValue = col;
|
||||
uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0;
|
||||
__ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr;
|
||||
MicroAPI::MaskReg mask = MicroAPI::UpdateMask<T>(maskValue);
|
||||
MicroAPI::LoadAlign(regList.vreg, initialAddr + offset);
|
||||
for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) {
|
||||
uint64_t addOffset = offset + repeatIdx * repeatOffset;
|
||||
MicroAPI::LoadAlign(regList.vregape, inputAddr + addOffset);
|
||||
MicroAPI::Add(regList.vreg, regList.vreg, regList.vregape, mask);
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + offset, regList.vreg, mask);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void Add64VFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, uint32_t row, uint32_t col, uint32_t actualCol0, uint32_t actualCol1)
|
||||
{
|
||||
AddRegList<T> regList[4];
|
||||
uint32_t loopTimes = row / 4;
|
||||
for (uint32_t idx = 0; idx < loopTimes; idx++) {
|
||||
uint64_t offset0 = idx * 4 * actualCol0;
|
||||
uint64_t offset1 = idx * 4 * actualCol1;
|
||||
AddVFImpl(inputAddr, apeAddr, regList[0], row, col, offset0, offset1);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[1], row, col, offset0 + actualCol0, offset1 + actualCol1);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[2], row, col, offset0 + 2 * actualCol0, offset1 + 2 * actualCol1);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[3], row, col, offset0 + 3 * actualCol0, offset1 + 3 * actualCol1);
|
||||
}
|
||||
|
||||
if (row % 4 > 0) {
|
||||
AddVFImpl(inputAddr, apeAddr, regList[0], row, col, loopTimes * 4 * actualCol0, loopTimes * 4 * actualCol1);
|
||||
}
|
||||
|
||||
if (row % 4 > 1) {
|
||||
AddVFImpl(inputAddr, apeAddr, regList[1], row, col, (loopTimes * 4 + 1) * actualCol0, (loopTimes * 4 + 1) * actualCol1);
|
||||
}
|
||||
|
||||
if (row % 4 > 2) {
|
||||
AddVFImpl(inputAddr, apeAddr, regList[2], row, col, (loopTimes * 4 + 2) * actualCol0, (loopTimes * 4 + 2) * actualCol1);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void Add128VFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, uint32_t row, uint32_t actualCol0, uint32_t actualCol1)
|
||||
{
|
||||
AddRegList<T> regList[4];
|
||||
uint32_t loopTimes = row / 2;
|
||||
for (uint32_t idx = 0; idx < loopTimes; idx++) {
|
||||
uint64_t offset0 = idx * 2 * actualCol0;
|
||||
uint64_t offset1 = idx * 2 * actualCol1;
|
||||
AddVFImpl(inputAddr, apeAddr, regList[0], row, FLOAT_REP_SIZE, offset0, offset1);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[1], row, FLOAT_REP_SIZE, offset0 + FLOAT_REP_SIZE, offset1 + FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[2], row, FLOAT_REP_SIZE, offset0 + actualCol0, offset1 + actualCol1);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[3], row, FLOAT_REP_SIZE, offset0 + actualCol0 + FLOAT_REP_SIZE, offset1 + actualCol1 + FLOAT_REP_SIZE);
|
||||
}
|
||||
|
||||
if (row % 2 > 0) {
|
||||
AddVFImpl(inputAddr, apeAddr, regList[0], row, FLOAT_REP_SIZE, loopTimes * 2 * actualCol0, loopTimes * 2 * actualCol1);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[1], row, FLOAT_REP_SIZE, loopTimes * 2 * actualCol0 + FLOAT_REP_SIZE, loopTimes * 2 * actualCol1 + FLOAT_REP_SIZE);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void Add256VFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, uint32_t row, uint32_t actualCol0, uint32_t actualCol1)
|
||||
{
|
||||
AddRegList<T> regList[4];
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
for (uint32_t idx = 0; idx < row; idx++) {
|
||||
uint64_t offset0 = idx * actualCol0;
|
||||
uint64_t offset1 = idx * actualCol1;
|
||||
AddVFImpl(inputAddr, apeAddr, regList[0], row, FLOAT_REP_SIZE, offset0, offset1);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[1], row, FLOAT_REP_SIZE, offset0 + FLOAT_REP_SIZE, offset1 + FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[2], row, FLOAT_REP_SIZE, offset0 + 2 * FLOAT_REP_SIZE, offset1 + 2 * FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[3], row, FLOAT_REP_SIZE, offset0 + 3 * FLOAT_REP_SIZE, offset1 + 3 * FLOAT_REP_SIZE);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void Add512VFImpl(__ubuf__ T *inputAddr, __ubuf__ T *apeAddr, uint32_t row, uint32_t actualCol0, uint32_t actualCol1)
|
||||
{
|
||||
AddRegList<T> regList[8];
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
for (uint32_t idx = 0; idx < row; idx++) {
|
||||
uint64_t offset0 = idx * actualCol0;
|
||||
uint64_t offset1 = idx * actualCol1;
|
||||
AddVFImpl(inputAddr, apeAddr, regList[0], row, FLOAT_REP_SIZE, offset0, offset1);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[1], row, FLOAT_REP_SIZE, offset0 + FLOAT_REP_SIZE, offset1 + FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[2], row, FLOAT_REP_SIZE, offset0 + 2 * FLOAT_REP_SIZE, offset1 + 2 * FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[3], row, FLOAT_REP_SIZE, offset0 + 3 * FLOAT_REP_SIZE, offset1 + 3 * FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[4], row, FLOAT_REP_SIZE, offset0 + 4 * FLOAT_REP_SIZE, offset1 + 4 * FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[5], row, FLOAT_REP_SIZE, offset0 + 5 * FLOAT_REP_SIZE, offset1 + 5 * FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[6], row, FLOAT_REP_SIZE, offset0 + 6 * FLOAT_REP_SIZE, offset1 + 6 * FLOAT_REP_SIZE);
|
||||
AddVFImpl(inputAddr, apeAddr, regList[7], row, FLOAT_REP_SIZE, offset0 + 7 * FLOAT_REP_SIZE, offset1 + 7 * FLOAT_REP_SIZE);
|
||||
}
|
||||
}
|
||||
|
||||
template <bool IS_FIRST, typename T>
|
||||
__simd_vf__ void MultiAdd64VFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, uint32_t row, uint32_t col,
|
||||
uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset)
|
||||
{
|
||||
AddRegList<T> regList[4];
|
||||
uint32_t loopTimes = row / 4;
|
||||
uint32_t maskValue = col;
|
||||
uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0;
|
||||
__ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr;
|
||||
MicroAPI::MaskReg mask = MicroAPI::UpdateMask<T>(maskValue);
|
||||
for (uint32_t idx = 0; idx < loopTimes; idx++) {
|
||||
uint64_t offset = idx * 4 * actualCol;
|
||||
MicroAPI::LoadAlign(regList[0].vreg, initialAddr + offset);
|
||||
MicroAPI::LoadAlign(regList[1].vreg, initialAddr + offset + actualCol);
|
||||
MicroAPI::LoadAlign(regList[2].vreg, initialAddr + offset + 2 * actualCol);
|
||||
MicroAPI::LoadAlign(regList[3].vreg, initialAddr + offset + 3 * actualCol);
|
||||
for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) {
|
||||
uint64_t addOffset = offset + repeatIdx * repeatOffset;
|
||||
MicroAPI::LoadAlign(regList[0].vregape, inputAddr + addOffset);
|
||||
MicroAPI::LoadAlign(regList[1].vregape, inputAddr + addOffset + actualCol);
|
||||
MicroAPI::LoadAlign(regList[2].vregape, inputAddr + addOffset + 2 * actualCol);
|
||||
MicroAPI::LoadAlign(regList[3].vregape, inputAddr + addOffset + 3 * actualCol);
|
||||
MicroAPI::Add(regList[0].vreg, regList[0].vreg, regList[0].vregape, mask);
|
||||
MicroAPI::Add(regList[1].vreg, regList[1].vreg, regList[1].vregape, mask);
|
||||
MicroAPI::Add(regList[2].vreg, regList[2].vreg, regList[2].vregape, mask);
|
||||
MicroAPI::Add(regList[3].vreg, regList[3].vreg, regList[3].vregape, mask);
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + offset, regList[0].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + actualCol, regList[1].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 2 * actualCol, regList[2].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 3 * actualCol, regList[3].vreg, mask);
|
||||
}
|
||||
|
||||
if (row % 4 > 0) {
|
||||
MultiAddVFImpl<IS_FIRST, T>(outputAddr, inputAddr, regList[0], row, col, loopTimes * 4 * actualCol, repeatNum,
|
||||
repeatOffset);
|
||||
}
|
||||
|
||||
if (row % 4 > 1) {
|
||||
MultiAddVFImpl<IS_FIRST, T>(outputAddr, inputAddr, regList[1], row, col, (loopTimes * 4 + 1) * actualCol,
|
||||
repeatNum, repeatOffset);
|
||||
}
|
||||
|
||||
if (row % 4 > 2) {
|
||||
MultiAddVFImpl<IS_FIRST, T>(outputAddr, inputAddr, regList[2], row, col, (loopTimes * 4 + 2) * actualCol,
|
||||
repeatNum, repeatOffset);
|
||||
}
|
||||
}
|
||||
|
||||
template <bool IS_FIRST, typename T>
|
||||
__simd_vf__ void MultiAdd128VFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, uint32_t row, uint32_t col,
|
||||
uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset)
|
||||
{
|
||||
AddRegList<T> regList[4];
|
||||
uint32_t loopTimes = row / 2;
|
||||
uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0;
|
||||
__ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr;
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
for (uint32_t idx = 0; idx < loopTimes; idx++) {
|
||||
uint64_t offset = idx * actualCol * 2;
|
||||
MicroAPI::LoadAlign(regList[0].vreg, initialAddr + offset);
|
||||
MicroAPI::LoadAlign(regList[1].vreg, initialAddr + offset + FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[2].vreg, initialAddr + offset + actualCol);
|
||||
MicroAPI::LoadAlign(regList[3].vreg, initialAddr + offset + actualCol + FLOAT_REP_SIZE);
|
||||
for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) {
|
||||
uint64_t addOffset = offset + repeatIdx * repeatOffset;
|
||||
MicroAPI::LoadAlign(regList[0].vregape, inputAddr + addOffset);
|
||||
MicroAPI::LoadAlign(regList[1].vregape, inputAddr + addOffset + FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[2].vregape, inputAddr + addOffset + actualCol);
|
||||
MicroAPI::LoadAlign(regList[3].vregape, inputAddr + addOffset + actualCol + FLOAT_REP_SIZE);
|
||||
MicroAPI::Add(regList[0].vreg, regList[0].vreg, regList[0].vregape, mask);
|
||||
MicroAPI::Add(regList[1].vreg, regList[1].vreg, regList[1].vregape, mask);
|
||||
MicroAPI::Add(regList[2].vreg, regList[2].vreg, regList[2].vregape, mask);
|
||||
MicroAPI::Add(regList[3].vreg, regList[3].vreg, regList[3].vregape, mask);
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + offset, regList[0].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + FLOAT_REP_SIZE, regList[1].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + actualCol, regList[2].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + actualCol + FLOAT_REP_SIZE, regList[3].vreg, mask);
|
||||
}
|
||||
|
||||
if (row % 2 > 0) {
|
||||
MultiAddVFImpl<IS_FIRST, T>(outputAddr, inputAddr, regList[0], row, col, loopTimes * 2 * actualCol, repeatNum,
|
||||
repeatOffset);
|
||||
MultiAddVFImpl<IS_FIRST, T>(outputAddr, inputAddr, regList[1], row, col,
|
||||
loopTimes * 2 * actualCol + FLOAT_REP_SIZE, repeatNum, repeatOffset);
|
||||
}
|
||||
}
|
||||
|
||||
template <bool IS_FIRST, typename T>
|
||||
__simd_vf__ void MultiAdd256VFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, uint32_t row,
|
||||
uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset)
|
||||
{
|
||||
AddRegList<T> regList[4];
|
||||
uint32_t loopTimes = row;
|
||||
uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0;
|
||||
__ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr;
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
for (uint32_t idx = 0; idx < loopTimes; idx++) {
|
||||
uint64_t offset = idx * actualCol;
|
||||
MicroAPI::LoadAlign(regList[0].vreg, initialAddr + offset);
|
||||
MicroAPI::LoadAlign(regList[1].vreg, initialAddr + offset + FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[2].vreg, initialAddr + offset + 2 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[3].vreg, initialAddr + offset + 3 * FLOAT_REP_SIZE);
|
||||
for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) {
|
||||
uint64_t addOffset = offset + repeatIdx * repeatOffset;
|
||||
MicroAPI::LoadAlign(regList[0].vregape, inputAddr + addOffset);
|
||||
MicroAPI::LoadAlign(regList[1].vregape, inputAddr + addOffset + FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[2].vregape, inputAddr + addOffset + 2 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[3].vregape, inputAddr + addOffset + 3 * FLOAT_REP_SIZE);
|
||||
MicroAPI::Add(regList[0].vreg, regList[0].vreg, regList[0].vregape, mask);
|
||||
MicroAPI::Add(regList[1].vreg, regList[1].vreg, regList[1].vregape, mask);
|
||||
MicroAPI::Add(regList[2].vreg, regList[2].vreg, regList[2].vregape, mask);
|
||||
MicroAPI::Add(regList[3].vreg, regList[3].vreg, regList[3].vregape, mask);
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + offset, regList[0].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + FLOAT_REP_SIZE, regList[1].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 2 * FLOAT_REP_SIZE, regList[2].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 3 * FLOAT_REP_SIZE, regList[3].vreg, mask);
|
||||
}
|
||||
}
|
||||
|
||||
template <bool IS_FIRST, typename T>
|
||||
__simd_vf__ void MultiAdd512VFImpl(__ubuf__ T *outputAddr, __ubuf__ T *inputAddr, uint32_t row,
|
||||
uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset)
|
||||
{
|
||||
AddRegList<T> regList[8];
|
||||
uint32_t loopTimes = row;
|
||||
uint32_t initialRepeatIdx = IS_FIRST ? 1 : 0;
|
||||
__ubuf__ T *initialAddr = IS_FIRST ? inputAddr : outputAddr;
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
for (uint32_t idx = 0; idx < loopTimes; idx++) {
|
||||
uint64_t offset = idx * actualCol;
|
||||
MicroAPI::LoadAlign(regList[0].vreg, initialAddr + offset);
|
||||
MicroAPI::LoadAlign(regList[1].vreg, initialAddr + offset + FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[2].vreg, initialAddr + offset + 2 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[3].vreg, initialAddr + offset + 3 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[4].vreg, initialAddr + offset + 4 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[5].vreg, initialAddr + offset + 5 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[6].vreg, initialAddr + offset + 6 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[7].vreg, initialAddr + offset + 7 * FLOAT_REP_SIZE);
|
||||
for (uint32_t repeatIdx = initialRepeatIdx; repeatIdx < repeatNum; repeatIdx++) {
|
||||
uint64_t addOffset = offset + repeatIdx * row * actualCol;
|
||||
MicroAPI::LoadAlign(regList[0].vregape, inputAddr + addOffset);
|
||||
MicroAPI::LoadAlign(regList[1].vregape, inputAddr + addOffset + FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[2].vregape, inputAddr + addOffset + 2 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[3].vregape, inputAddr + addOffset + 3 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[4].vregape, inputAddr + addOffset + 4 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[5].vregape, inputAddr + addOffset + 5 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[6].vregape, inputAddr + addOffset + 6 * FLOAT_REP_SIZE);
|
||||
MicroAPI::LoadAlign(regList[7].vregape, inputAddr + addOffset + 7 * FLOAT_REP_SIZE);
|
||||
MicroAPI::Add(regList[0].vreg, regList[0].vreg, regList[0].vregape, mask);
|
||||
MicroAPI::Add(regList[1].vreg, regList[1].vreg, regList[1].vregape, mask);
|
||||
MicroAPI::Add(regList[2].vreg, regList[2].vreg, regList[2].vregape, mask);
|
||||
MicroAPI::Add(regList[3].vreg, regList[3].vreg, regList[3].vregape, mask);
|
||||
MicroAPI::Add(regList[4].vreg, regList[4].vreg, regList[4].vregape, mask);
|
||||
MicroAPI::Add(regList[5].vreg, regList[5].vreg, regList[5].vregape, mask);
|
||||
MicroAPI::Add(regList[6].vreg, regList[6].vreg, regList[6].vregape, mask);
|
||||
MicroAPI::Add(regList[7].vreg, regList[7].vreg, regList[7].vregape, mask);
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + offset, regList[0].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + FLOAT_REP_SIZE, regList[1].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 2 * FLOAT_REP_SIZE, regList[2].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 3 * FLOAT_REP_SIZE, regList[3].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 4 * FLOAT_REP_SIZE, regList[4].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 5 * FLOAT_REP_SIZE, regList[5].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 6 * FLOAT_REP_SIZE, regList[6].vreg, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + offset + 7 * FLOAT_REP_SIZE, regList[7].vreg, mask);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief AddVF 输入与apt相加
|
||||
* @param rightLocal 输出tensor []
|
||||
* @param leftLocal 输入tensor [row, col]
|
||||
* @param aptLocal apt输入tensor [r]
|
||||
* @param apeIdx ape起始位置
|
||||
* @param d coff*d为ape的D轴大小
|
||||
* @param coreSplitD scoreleft大小,coff*coreSplitD为总大小
|
||||
* @param coreSplitS 核间d轴切分大小
|
||||
*/
|
||||
template <typename T>
|
||||
__aicore__ inline void AddVF(const LocalTensor<T> &scoreLocal, const LocalTensor<T> &apeLocal, uint32_t row,
|
||||
uint32_t col, uint32_t actualCol0, uint32_t actualCol1)
|
||||
{
|
||||
__ubuf__ T *scoreAddr = (__ubuf__ T *)scoreLocal.GetPhyAddr();
|
||||
__ubuf__ T *apeAddr = (__ubuf__ T *)apeLocal.GetPhyAddr();
|
||||
|
||||
if (col <= 64) {
|
||||
Add64VFImpl<T>(scoreAddr, apeAddr, row, col, actualCol0, actualCol1);
|
||||
} else if (col == 128) {
|
||||
Add128VFImpl<T>(scoreAddr, apeAddr, row, actualCol0, actualCol1);
|
||||
} else if (col == 256) {
|
||||
Add256VFImpl<T>(scoreAddr, apeAddr, row, actualCol0, actualCol1);
|
||||
} else if (col == 512) {
|
||||
Add512VFImpl<T>(scoreAddr, apeAddr, row, actualCol0, actualCol1);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__aicore__ inline void AddVF(const LocalTensor<T> &scoreLocal, const LocalTensor<T> &apeLocal, uint32_t row,
|
||||
uint32_t col, uint32_t actualCol)
|
||||
{
|
||||
__ubuf__ T *scoreAddr = (__ubuf__ T *)scoreLocal.GetPhyAddr();
|
||||
__ubuf__ T *apeAddr = (__ubuf__ T *)apeLocal.GetPhyAddr();
|
||||
|
||||
if (col <= 64) {
|
||||
Add64VFImpl<T>(scoreAddr, apeAddr, row, col, actualCol, actualCol);
|
||||
} else if (col == 128) {
|
||||
Add128VFImpl<T>(scoreAddr, apeAddr, row, actualCol, actualCol);
|
||||
} else if (col == 256) {
|
||||
Add256VFImpl<T>(scoreAddr, apeAddr, row, actualCol, actualCol);
|
||||
} else if (col == 512) {
|
||||
Add512VFImpl<T>(scoreAddr, apeAddr, row, actualCol, actualCol);
|
||||
}
|
||||
}
|
||||
|
||||
template <bool IS_FIRST, typename T>
|
||||
__aicore__ inline void MultiAddVF(const LocalTensor<T> &outputLocal, const LocalTensor<T> &inputLocal, uint32_t row,
|
||||
uint32_t col, uint32_t actualCol, uint32_t repeatNum, uint64_t repeatOffset)
|
||||
{
|
||||
__ubuf__ T *outputAddr = (__ubuf__ T *)outputLocal.GetPhyAddr();
|
||||
__ubuf__ T *inputAddr = (__ubuf__ T *)inputLocal.GetPhyAddr();
|
||||
if (col <= 64) {
|
||||
MultiAdd64VFImpl<IS_FIRST, T>(outputAddr, inputAddr, row, col, actualCol, repeatNum, repeatOffset);
|
||||
} else if (col == 128) {
|
||||
MultiAdd128VFImpl<IS_FIRST, T>(outputAddr, inputAddr, row, col, actualCol, repeatNum, repeatOffset);
|
||||
} else if (col == 256) {
|
||||
MultiAdd256VFImpl<IS_FIRST, T>(outputAddr, inputAddr, row, actualCol, repeatNum, repeatOffset);
|
||||
} else if (col == 512) {
|
||||
MultiAdd512VFImpl<IS_FIRST, T>(outputAddr, inputAddr, row, actualCol, repeatNum, repeatOffset);
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
318
csrc/attention/compressor/op_kernel/arch35/vf/vf_mul.h
Normal file
318
csrc/attention/compressor/op_kernel/arch35/vf/vf_mul.h
Normal file
@@ -0,0 +1,318 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_mul.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef VF_MUL_H
|
||||
#define VF_MUL_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include <cstdint>
|
||||
using namespace AscendC;
|
||||
|
||||
constexpr uint32_t FLOATBYTE = 4;
|
||||
constexpr uint32_t baseD8 = 8;
|
||||
constexpr uint32_t baseD16 = 16;
|
||||
constexpr uint32_t baseD32 = 32;
|
||||
constexpr uint32_t baseD64 = 64;
|
||||
constexpr uint32_t baseD128 = 128;
|
||||
constexpr uint32_t baseD256 = 256;
|
||||
constexpr uint32_t baseD512 = 512;
|
||||
|
||||
|
||||
template <typename T>
|
||||
__simd_callee__ inline T SimdCeilDivT(T num1, T num2)
|
||||
{
|
||||
if (num2 == 0) {
|
||||
return static_cast<T>(0);
|
||||
}
|
||||
return (num1 + num2 - 1) / num2;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
struct ReduceMulRegList {
|
||||
MicroAPI::RegTensor<T> vreg0;
|
||||
MicroAPI::RegTensor<T> vreg1;
|
||||
MicroAPI::RegTensor<T> vregMul;
|
||||
MicroAPI::RegTensor<T> vregSum;
|
||||
};
|
||||
|
||||
|
||||
template <typename T>
|
||||
__simd_callee__ void LoadMulAddVFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, ReduceMulRegList<T> ®List, uint64_t offset, uint32_t maskValue)
|
||||
{
|
||||
MicroAPI::MaskReg mask = MicroAPI::UpdateMask<T>(maskValue);
|
||||
MicroAPI::LoadAlign(regList.vreg0, kvAddr + offset);
|
||||
MicroAPI::LoadAlign(regList.vreg1, scoreAddr + offset);
|
||||
MicroAPI::Mul(regList.vregMul, regList.vreg0, regList.vreg1, mask);
|
||||
MicroAPI::Add(regList.vregSum, regList.vregSum, regList.vregMul, mask);
|
||||
}
|
||||
|
||||
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void MulReduceSumbase8VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr,
|
||||
const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt,
|
||||
const uint32_t baseD)
|
||||
{
|
||||
ReduceMulRegList<T> regList;
|
||||
MicroAPI::RegTensor<T> vregSum0;
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg maskL32 = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::VL32>();
|
||||
MicroAPI::MaskReg maskL16 = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::VL16>();
|
||||
MicroAPI::MaskReg maskL8 = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::VL8>();
|
||||
MicroAPI::MaskReg maskH32;
|
||||
MicroAPI::MaskReg maskH48;
|
||||
MicroAPI::MaskReg maskH56;
|
||||
MicroAPI::Not(maskH48, maskL16, mask);
|
||||
MicroAPI::Not(maskH32, maskL32, mask);
|
||||
MicroAPI::Not(maskH56, maskL8, mask);
|
||||
uint32_t offset = 0;
|
||||
uint32_t rCnt = coff * cmpRatio;
|
||||
for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) {
|
||||
MicroAPI::Duplicate(regList.vregSum, 0, mask);
|
||||
// 当前仅支持coff * cmpRatio为2的幂的情况
|
||||
for (uint32_t rLoop = 0; rLoop < SimdCeilDivT(rCnt, 8U); rLoop++) {
|
||||
uint32_t dealLen = min((rCnt - rLoop * 8) * baseD, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList, offset, dealLen);
|
||||
offset += dealLen;
|
||||
}
|
||||
// 64 -> 32
|
||||
MicroAPI::Squeeze<T, AscendC::MicroAPI::GatherMaskMode::NO_STORE_REG>(vregSum0, regList.vregSum, maskH32);
|
||||
MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL32);
|
||||
|
||||
// 32 -> 16
|
||||
MicroAPI::Squeeze<T, AscendC::MicroAPI::GatherMaskMode::NO_STORE_REG>(vregSum0, regList.vregSum, maskH48);
|
||||
MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL16);
|
||||
|
||||
// 16 -> 8
|
||||
MicroAPI::Squeeze<T, AscendC::MicroAPI::GatherMaskMode::NO_STORE_REG>(vregSum0, regList.vregSum, maskH56);
|
||||
MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL8);
|
||||
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList.vregSum, maskL8);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void MulReduceSumbase16VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr,
|
||||
const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt,
|
||||
const uint32_t baseD)
|
||||
{
|
||||
ReduceMulRegList<T> regList;
|
||||
MicroAPI::RegTensor<T> vregSum0;
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg maskL32 = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::VL32>();
|
||||
MicroAPI::MaskReg maskL16 = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::VL16>();
|
||||
MicroAPI::MaskReg maskH32;
|
||||
MicroAPI::MaskReg maskH48;
|
||||
MicroAPI::Not(maskH48, maskL16, mask);
|
||||
MicroAPI::Not(maskH32, maskL32, mask);
|
||||
uint32_t offset = 0;
|
||||
uint32_t rCnt = coff * cmpRatio;
|
||||
for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) {
|
||||
MicroAPI::Duplicate(regList.vregSum, 0, mask);
|
||||
// 当前仅支持coff * cmpRatio为2的幂的情况
|
||||
for (uint32_t rLoop = 0; rLoop < SimdCeilDivT(rCnt, 4U); rLoop++) {
|
||||
uint32_t dealLen = min((rCnt - rLoop * 4) * baseD, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList, offset, dealLen);
|
||||
offset += dealLen;
|
||||
}
|
||||
// 64 -> 32
|
||||
MicroAPI::Squeeze<T, AscendC::MicroAPI::GatherMaskMode::NO_STORE_REG>(vregSum0, regList.vregSum, maskH32);
|
||||
MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL32);
|
||||
|
||||
// 32 -> 16
|
||||
MicroAPI::Squeeze<T, AscendC::MicroAPI::GatherMaskMode::NO_STORE_REG>(vregSum0, regList.vregSum, maskH48);
|
||||
MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL16);
|
||||
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList.vregSum, maskL16);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void MulReduceSumbase32VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr,
|
||||
const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt,
|
||||
const uint32_t baseD)
|
||||
{
|
||||
ReduceMulRegList<T> regList;
|
||||
MicroAPI::RegTensor<T> vregSum0;
|
||||
MicroAPI::RegTensor<T> vregSum1;
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg maskL32 = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::VL32>();
|
||||
MicroAPI::MaskReg maskH32;
|
||||
MicroAPI::Not(maskH32, maskL32, mask);
|
||||
uint32_t offset = 0;
|
||||
uint32_t rCnt = coff * cmpRatio;
|
||||
for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) {
|
||||
MicroAPI::Duplicate(regList.vregSum, 0, mask);
|
||||
// 当前仅支持coff * cmpRatio为2的幂的情况
|
||||
for (uint32_t rLoop = 0; rLoop < SimdCeilDivT(rCnt, 2U); rLoop++) {
|
||||
uint32_t dealLen = min((rCnt - rLoop * 2) * baseD, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList, offset, dealLen);
|
||||
offset += dealLen;
|
||||
}
|
||||
// 64 -> 32
|
||||
MicroAPI::Squeeze<T, AscendC::MicroAPI::GatherMaskMode::NO_STORE_REG>(vregSum0, regList.vregSum, maskH32);
|
||||
MicroAPI::Add(regList.vregSum, regList.vregSum, vregSum0, maskL32);
|
||||
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList.vregSum, maskL32);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void MulReduceSumbase64VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr,
|
||||
const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt,
|
||||
const uint32_t baseD)
|
||||
{
|
||||
ReduceMulRegList<T> regList;
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
uint32_t offset = 0;
|
||||
uint32_t rCnt = coff * cmpRatio;
|
||||
for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) {
|
||||
MicroAPI::Duplicate(regList.vregSum, 0, mask);
|
||||
for (uint32_t rLoop = 0; rLoop < rCnt; rLoop++) {
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList, offset, baseD64);
|
||||
offset += baseD;
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList.vregSum, mask);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void MulReduceSumbase128VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr,
|
||||
const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt,
|
||||
const uint32_t baseD)
|
||||
{
|
||||
ReduceMulRegList<T> regList[2];
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
uint32_t offset = 0;
|
||||
uint32_t rCnt = coff * cmpRatio;
|
||||
for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) {
|
||||
MicroAPI::Duplicate(regList[0].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[1].vregSum, 0, mask);
|
||||
for (uint32_t rLoop = 0; rLoop < rCnt; rLoop++) {
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[0], offset, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[1], offset + baseD64, baseD64);
|
||||
offset += baseD;
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList[0].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + baseD64, regList[1].vregSum, mask);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void MulReduceSumbase256VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr,
|
||||
const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt,
|
||||
const uint32_t baseD)
|
||||
{
|
||||
ReduceMulRegList<T> regList[4];
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
uint32_t offset = 0;
|
||||
uint32_t rCnt = coff * cmpRatio;
|
||||
for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) {
|
||||
MicroAPI::Duplicate(regList[0].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[1].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[2].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[3].vregSum, 0, mask);
|
||||
for (uint32_t rLoop = 0; rLoop < rCnt; rLoop++) {
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[0], offset, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[1], offset + baseD64, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[2], offset + 2 * baseD64, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[3], offset + 3 * baseD64, baseD64);
|
||||
offset += baseD;
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList[0].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + baseD64, regList[1].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 2 * baseD64, regList[2].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 3 * baseD64, regList[3].vregSum, mask);
|
||||
}
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__simd_vf__ void MulReduceSumbase512VFImpl(__ubuf__ T *kvAddr, __ubuf__ T *scoreAddr, __ubuf__ T *outputAddr,
|
||||
const uint32_t coff, const uint32_t cmpRatio, const uint32_t scLoopCnt,
|
||||
const uint32_t baseD)
|
||||
{
|
||||
ReduceMulRegList<T> regList[8];
|
||||
MicroAPI::MaskReg mask = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
uint32_t offset = 0;
|
||||
uint32_t rCnt = coff * cmpRatio;
|
||||
for (uint32_t scLoop = 0; scLoop < scLoopCnt; scLoop++) {
|
||||
MicroAPI::Duplicate(regList[0].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[1].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[2].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[3].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[4].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[5].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[6].vregSum, 0, mask);
|
||||
MicroAPI::Duplicate(regList[7].vregSum, 0, mask);
|
||||
for (uint32_t rLoop = 0; rLoop < rCnt; rLoop++) {
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[0], offset, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[1], offset + baseD64, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[2], offset + 2 * baseD64, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[3], offset + 3 * baseD64, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[4], offset + 4 * baseD64, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[5], offset + 5 * baseD64, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[6], offset + 6 * baseD64, baseD64);
|
||||
LoadMulAddVFImpl(kvAddr, scoreAddr, regList[7], offset + 7 * baseD64, baseD64);
|
||||
offset += baseD;
|
||||
}
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD, regList[0].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + baseD64, regList[1].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 2 * baseD64, regList[2].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 3 * baseD64, regList[3].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 4 * baseD64, regList[4].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 5 * baseD64, regList[5].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 6 * baseD64, regList[6].vregSum, mask);
|
||||
MicroAPI::StoreAlign(outputAddr + scLoop * baseD + 7 * baseD64, regList[7].vregSum, mask);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief MulReduceSumbaseVF 包含mul和reducesum
|
||||
* @param outputLocal 输出tensor []
|
||||
* @param coff
|
||||
* @param cmpRatio 压缩块大小
|
||||
* @param baseD 核内d轴切分大小
|
||||
* @param scLoopCnt sc数,
|
||||
*/
|
||||
|
||||
// 当前仅支持coff * cmpRatio为2的幂的情况
|
||||
template <typename T>
|
||||
__aicore__ inline void MulReduceSumbaseVF(const LocalTensor<T> &kvLocal, const LocalTensor<T> &scoreLocal,
|
||||
const LocalTensor<T> &outputLocal, const uint32_t coff, const uint32_t cmpRatio,
|
||||
const uint32_t baseD, const uint32_t scLoopCnt)
|
||||
{
|
||||
|
||||
__ubuf__ T *kvAddr = (__ubuf__ T *)kvLocal.GetPhyAddr();
|
||||
__ubuf__ T *scoreAddr = (__ubuf__ T *)scoreLocal.GetPhyAddr();
|
||||
__ubuf__ T *outputAddr = (__ubuf__ T *)outputLocal.GetPhyAddr();
|
||||
if (baseD == baseD8) {
|
||||
MulReduceSumbase8VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD);
|
||||
} else if (baseD == baseD16) {
|
||||
MulReduceSumbase16VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD);
|
||||
} else if (baseD == baseD32) {
|
||||
MulReduceSumbase32VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD);
|
||||
} else if (baseD == baseD64) {
|
||||
MulReduceSumbase64VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD);
|
||||
} else if (baseD == baseD128) {
|
||||
MulReduceSumbase128VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD);
|
||||
} else if (baseD == baseD256) {
|
||||
MulReduceSumbase256VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD);
|
||||
} else if (baseD == baseD512) {
|
||||
MulReduceSumbase512VFImpl(kvAddr, scoreAddr, outputAddr, coff, cmpRatio, scLoopCnt, baseD);
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
#endif
|
||||
95
csrc/attention/compressor/op_kernel/arch35/vf/vf_rms_norm.h
Normal file
95
csrc/attention/compressor/op_kernel/arch35/vf/vf_rms_norm.h
Normal file
@@ -0,0 +1,95 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_rms_norm.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef VF_RMS_NORM_H
|
||||
#define VF_RMS_NORM_H
|
||||
#include "kernel_tensor.h"
|
||||
|
||||
//repeatTimes——D轴的分块数
|
||||
template <typename T, typename GammaType>
|
||||
__simd_vf__ void RmsNormVFImpl(__ubuf__ T * inputBuf, __ubuf__ GammaType * gammaBuf, __ubuf__ T * outputBuf,
|
||||
uint32_t repeatTimes, float reciprocal, float epsilon)
|
||||
{
|
||||
MicroAPI::RegTensor<T> vregSum;
|
||||
MicroAPI::RegTensor<T> vregSumReduce;
|
||||
MicroAPI::RegTensor<T> vregDiv;
|
||||
MicroAPI::RegTensor<T> vregSquareRoot;
|
||||
|
||||
MicroAPI::MaskReg maskAll = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::ALL>();
|
||||
MicroAPI::MaskReg maskFirst = MicroAPI::CreateMask<T, MicroAPI::MaskPattern::VL1>();
|
||||
|
||||
static constexpr MicroAPI::CastTrait castTraitB162B32 = {MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::UNKNOWN, MicroAPI::MaskMergeMode::ZEROING, RoundMode::UNKNOWN};
|
||||
|
||||
MicroAPI::Duplicate<T,T>(vregSum, 0.0f);
|
||||
|
||||
for(uint32_t i = 0; i < repeatTimes; ++i){
|
||||
MicroAPI::RegTensor<T> vregX;
|
||||
MicroAPI::RegTensor<T> vregXSquare;
|
||||
uint64_t loopOffset = i * FLOAT_REP_SIZE;
|
||||
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregX, inputBuf + loopOffset);
|
||||
MicroAPI::Mul(vregXSquare, vregX, vregX, maskAll);
|
||||
MicroAPI::Add(vregSum, vregXSquare, vregSum, maskAll);
|
||||
}
|
||||
|
||||
MicroAPI::Reduce<MicroAPI::ReduceType::SUM, T, T, MicroAPI::MaskMergeMode::ZEROING>(vregSumReduce, vregSum, maskAll);
|
||||
MicroAPI::Muls<T, T, MicroAPI::MaskMergeMode::ZEROING>(vregSumReduce, vregSumReduce, reciprocal, maskFirst);
|
||||
MicroAPI::Adds<T, T, MicroAPI::MaskMergeMode::ZEROING>(vregSumReduce, vregSumReduce, epsilon, maskFirst);
|
||||
MicroAPI::Sqrt(vregSquareRoot, vregSumReduce, maskFirst);
|
||||
MicroAPI::Duplicate<T, MicroAPI::HighLowPart::LOWEST, MicroAPI::MaskMergeMode::ZEROING>(vregDiv, vregSquareRoot, maskAll);
|
||||
|
||||
for(uint32_t i = 0; i < repeatTimes; ++i){
|
||||
MicroAPI::RegTensor<T> vregX;
|
||||
MicroAPI::RegTensor<T> vregGammaCast;
|
||||
uint16_t loopOffset = i * FLOAT_REP_SIZE;
|
||||
|
||||
MicroAPI::LoadAlign<T, MicroAPI::LoadDist::DIST_NORM>(vregX, inputBuf + loopOffset);
|
||||
MicroAPI::LoadAlign<GammaType, MicroAPI::LoadDist::DIST_NORM>(vregGammaCast, gammaBuf + loopOffset);
|
||||
|
||||
MicroAPI::Div(vregX, vregX, vregDiv, maskAll);
|
||||
MicroAPI::Mul(vregX, vregX, vregGammaCast, maskAll);
|
||||
|
||||
MicroAPI::StoreAlign<T, MicroAPI::StoreDist::DIST_NORM>(outputBuf + loopOffset, vregX, maskAll);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* @brief RmsNormVF 对一行进行rmsnorm
|
||||
* @param outputLocal 输出tensor [row, col],row目前均为1
|
||||
* @param inputLocal 输入tensor [row, col]
|
||||
* @param gammaLocal gamma参数tensor [row, col]
|
||||
* @param rmsNormParams rmsNrom计算所需系数,包括
|
||||
row 行数 1
|
||||
col 列数,对应headSizeCq或headSizeCkv
|
||||
reciprocal ,1/N
|
||||
epsilon,防止除零极小数
|
||||
*/
|
||||
template <typename T, typename GammaType>
|
||||
__aicore__ inline void RmsNormVF(const LocalTensor<T> outputLocal, const LocalTensor<T> inputLocal, const LocalTensor<GammaType> gammaLocal,
|
||||
float reciprocal, float epsilon, uint32_t row, uint32_t col)
|
||||
{
|
||||
uint32_t cnt = row * col;
|
||||
uint32_t repeatTimes = (cnt + FLOAT_REP_SIZE - 1) / FLOAT_REP_SIZE;
|
||||
|
||||
__ubuf__ T * inputBuf = (__ubuf__ T *)inputLocal.GetPhyAddr();
|
||||
__ubuf__ GammaType * gammaBuf = (__ubuf__ GammaType *)gammaLocal.GetPhyAddr();
|
||||
__ubuf__ T * outputBuf = (__ubuf__ T *)outputLocal.GetPhyAddr();
|
||||
|
||||
RmsNormVFImpl<T, GammaType>(inputBuf, gammaBuf, outputBuf, repeatTimes, reciprocal, epsilon);
|
||||
}
|
||||
|
||||
|
||||
#endif
|
||||
158
csrc/attention/compressor/op_kernel/arch35/vf/vf_rope.h
Normal file
158
csrc/attention/compressor/op_kernel/arch35/vf/vf_rope.h
Normal file
@@ -0,0 +1,158 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file vf_rope.h
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#ifndef VF_ROPE_H
|
||||
#define VF_ROPE_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
#include "../compressor_comm.h"
|
||||
|
||||
using namespace AscendC;
|
||||
|
||||
constexpr MicroAPI::CastTrait castTraitB162B32 = {
|
||||
MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::UNKNOWN,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::UNKNOWN,
|
||||
};
|
||||
|
||||
constexpr MicroAPI::CastTrait castTraitB322B16 = {
|
||||
MicroAPI::RegLayout::ZERO,
|
||||
MicroAPI::SatMode::NO_SAT,
|
||||
MicroAPI::MaskMergeMode::ZEROING,
|
||||
RoundMode::CAST_RINT,
|
||||
};
|
||||
|
||||
|
||||
template <typename T, typename ROPET>
|
||||
__simd_vf__ void HalfModeRopeVF(__ubuf__ T *sinUb, __ubuf__ T *cosUb, __ubuf__ T *inUb, __ubuf__ ROPET *outUb,
|
||||
uint32_t row, uint32_t col, uint32_t actualCol, uint64_t baseAddr)
|
||||
{
|
||||
MicroAPI::RegTensor<T> vregCos;
|
||||
MicroAPI::RegTensor<T> vregHalfCos;
|
||||
MicroAPI::RegTensor<T> vregSin;
|
||||
MicroAPI::RegTensor<T> vregHalfSin;
|
||||
MicroAPI::RegTensor<T> vregIn;
|
||||
MicroAPI::RegTensor<T> vregHalfIn;
|
||||
MicroAPI::RegTensor<T> vregOut;
|
||||
MicroAPI::RegTensor<T> vregHalfOut;
|
||||
MicroAPI::RegTensor<T> vregCastIn;
|
||||
MicroAPI::RegTensor<ROPET> vregOutBf16;
|
||||
MicroAPI::RegTensor<ROPET> vregOutHalfBf16;
|
||||
MicroAPI::RegTensor<ROPET> vregCastOut;
|
||||
uint32_t maskValue = col / 2;
|
||||
MicroAPI::MaskReg mask = MicroAPI::UpdateMask<T>(maskValue);
|
||||
uint32_t halfCol = col / 2;
|
||||
|
||||
|
||||
for (uint32_t rIdx = 0; rIdx < row; rIdx++) {
|
||||
__ubuf__ T *curSinUb = sinUb + rIdx * col;
|
||||
__ubuf__ T *curCosUb = cosUb + rIdx * col;
|
||||
__ubuf__ T *curInUb = inUb + rIdx * actualCol;
|
||||
__ubuf__ ROPET *curOutUb = outUb + rIdx * actualCol;
|
||||
|
||||
MicroAPI::DataCopy(vregIn, curInUb + baseAddr);
|
||||
MicroAPI::DataCopy(vregHalfIn, curInUb + baseAddr + halfCol);
|
||||
MicroAPI::DataCopy(vregCos, curCosUb);
|
||||
MicroAPI::DataCopy(vregHalfCos, curCosUb + halfCol);
|
||||
MicroAPI::DataCopy(vregSin, curSinUb);
|
||||
MicroAPI::DataCopy(vregHalfSin, curSinUb + halfCol);
|
||||
MicroAPI::Mul(vregSin, vregSin, vregHalfIn, mask);
|
||||
MicroAPI::Mul(vregHalfSin, vregHalfSin, vregIn, mask);
|
||||
MicroAPI::Mul(vregCos, vregCos, vregIn, mask);
|
||||
MicroAPI::Sub(vregOut, vregCos, vregSin, mask);
|
||||
MicroAPI::Mul(vregHalfCos, vregHalfCos, vregHalfIn, mask);
|
||||
MicroAPI::Add(vregHalfOut, vregHalfSin, vregHalfCos, mask);
|
||||
MicroAPI::Cast<ROPET, T, castTraitB322B16>(vregOutBf16, vregOut, mask);
|
||||
MicroAPI::DataCopy<ROPET, MicroAPI::StoreDist::DIST_PACK_B32>(curOutUb + baseAddr, vregOutBf16, mask);
|
||||
MicroAPI::Cast<ROPET, T, castTraitB322B16>(vregOutHalfBf16, vregHalfOut, mask);
|
||||
MicroAPI::DataCopy<ROPET, MicroAPI::StoreDist::DIST_PACK_B32>(curOutUb + baseAddr + halfCol, vregOutHalfBf16,
|
||||
mask);
|
||||
|
||||
for (uint64_t dOffset = 0; dOffset < baseAddr; dOffset += 64) {
|
||||
uint32_t castMaskValue = min(baseAddr - dOffset, static_cast<uint64_t>(64));
|
||||
MicroAPI::MaskReg castMask = MicroAPI::UpdateMask<T>(castMaskValue);
|
||||
MicroAPI::DataCopy(vregCastIn, curInUb + dOffset);
|
||||
MicroAPI::Cast<ROPET, T, castTraitB322B16>(vregCastOut, vregCastIn, castMask);
|
||||
MicroAPI::DataCopy<ROPET, MicroAPI::StoreDist::DIST_PACK_B32>(curOutUb + dOffset, vregCastOut, castMask);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename T, typename ROPET>
|
||||
__simd_vf__ void InterleaveModeRopeVF(__ubuf__ T *sinUb, __ubuf__ T *cosUb, __ubuf__ T *inUb, __ubuf__ ROPET *outUb,
|
||||
uint32_t row, uint32_t col, uint32_t actualCol, uint64_t baseAddr)
|
||||
{
|
||||
MicroAPI::RegTensor<T> vregCos;
|
||||
MicroAPI::RegTensor<T> vregSin;
|
||||
MicroAPI::RegTensor<T> vregIn;
|
||||
MicroAPI::RegTensor<T> vregOdd;
|
||||
MicroAPI::RegTensor<T> vregEven;
|
||||
MicroAPI::RegTensor<T> vregOut;
|
||||
MicroAPI::RegTensor<T> vregTemp;
|
||||
MicroAPI::RegTensor<T> vregCastIn;
|
||||
MicroAPI::RegTensor<ROPET> vregOutBf16;
|
||||
MicroAPI::RegTensor<ROPET> vregCastOut;
|
||||
uint32_t maskValue = col;
|
||||
MicroAPI::MaskReg mask = MicroAPI::UpdateMask<T>(maskValue);
|
||||
|
||||
|
||||
for (uint32_t rIdx = 0; rIdx < row; rIdx++) {
|
||||
__ubuf__ T *curSinUb = sinUb + rIdx * col;
|
||||
__ubuf__ T *curCosUb = cosUb + rIdx * col;
|
||||
__ubuf__ T *curInUb = inUb + rIdx * actualCol;
|
||||
__ubuf__ ROPET *curOutUb = outUb + rIdx * actualCol;
|
||||
|
||||
MicroAPI::DataCopy(vregIn, curInUb + baseAddr);
|
||||
MicroAPI::DataCopy(vregCos, curCosUb);
|
||||
MicroAPI::DataCopy(vregSin, curSinUb);
|
||||
MicroAPI::Mul(vregCos, vregCos, vregIn, mask);
|
||||
MicroAPI::DeInterleave<T>(vregEven, vregOdd, vregIn, vregTemp);
|
||||
MicroAPI::Muls(vregOdd, vregOdd, static_cast<T>(-1.0), mask);
|
||||
MicroAPI::Interleave<T>(vregIn, vregTemp, vregOdd, vregEven);
|
||||
MicroAPI::Mul(vregSin, vregSin, vregIn, mask);
|
||||
MicroAPI::Add(vregOut, vregCos, vregSin, mask);
|
||||
MicroAPI::Cast<ROPET, T, castTraitB322B16>(vregOutBf16, vregOut, mask);
|
||||
MicroAPI::DataCopy<ROPET, MicroAPI::StoreDist::DIST_PACK_B32>(curOutUb + baseAddr, vregOutBf16, mask);
|
||||
for (uint64_t dOffset = 0; dOffset < baseAddr; dOffset += 64) {
|
||||
uint32_t castMaskValue = min(baseAddr - dOffset, static_cast<uint64_t>(64));
|
||||
MicroAPI::MaskReg castMask = MicroAPI::UpdateMask<T>(castMaskValue);
|
||||
MicroAPI::DataCopy(vregCastIn, curInUb + dOffset);
|
||||
MicroAPI::Cast<ROPET, T, castTraitB322B16>(vregCastOut, vregCastIn, castMask);
|
||||
MicroAPI::DataCopy<ROPET, MicroAPI::StoreDist::DIST_PACK_B32>(curOutUb + dOffset, vregCastOut, castMask);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <Compressor::ROTARY_MODE MODE, typename T, typename ROPET>
|
||||
__aicore__ inline void RopeVF(const LocalTensor<T> &sinTensor, const LocalTensor<T> &cosTensor,
|
||||
const LocalTensor<T> &inTensor, const LocalTensor<ROPET> &outTensor, uint32_t row,
|
||||
uint32_t col, uint32_t actualCol, uint64_t baseAddr)
|
||||
{
|
||||
__ubuf__ T *sinUb = (__ubuf__ T *)sinTensor.GetPhyAddr();
|
||||
__ubuf__ T *cosUb = (__ubuf__ T *)cosTensor.GetPhyAddr();
|
||||
__ubuf__ T *inUb = (__ubuf__ T *)inTensor.GetPhyAddr();
|
||||
__ubuf__ ROPET *outUb = (__ubuf__ ROPET *)outTensor.GetPhyAddr();
|
||||
|
||||
if constexpr (MODE == Compressor::ROTARY_MODE::HALF) {
|
||||
HalfModeRopeVF(sinUb, cosUb, inUb, outUb, row, col, actualCol, baseAddr);
|
||||
} else {
|
||||
InterleaveModeRopeVF(sinUb, cosUb, inUb, outUb, row, col, actualCol, baseAddr);
|
||||
}
|
||||
}
|
||||
|
||||
#endif
|
||||
1592
csrc/attention/compressor/op_kernel/arch35/vf/vf_softmax.h
Normal file
1592
csrc/attention/compressor/op_kernel/arch35/vf/vf_softmax.h
Normal file
File diff suppressed because it is too large
Load Diff
87
csrc/attention/compressor/op_kernel/compressor.cpp
Normal file
87
csrc/attention/compressor/op_kernel/compressor.cpp
Normal file
@@ -0,0 +1,87 @@
|
||||
/**
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
* CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
* Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
* See LICENSE in the root of the software repository for the full text of the License.
|
||||
*/
|
||||
|
||||
/*!
|
||||
* \file compressor.cpp
|
||||
* \brief
|
||||
*/
|
||||
|
||||
#if (__CCE_AICORE__ == 220)
|
||||
#include "arch32/compressor_kernel.h"
|
||||
#include "arch32/compressor_kernel_perf.h"
|
||||
#else
|
||||
#include "arch35/compressor_kernel.h"
|
||||
#include "arch35/compressor_kernel_full_load.h"
|
||||
#endif
|
||||
|
||||
using namespace Compressor;
|
||||
|
||||
#define INVOKE_COMPRESSOR_GENERAL_OP_IMPL(templateClass, ...) \
|
||||
do { \
|
||||
templateClass<COMPType<__VA_ARGS__>> op(&pipe, tilingData); \
|
||||
op.Init(x, wKv, wGate, stateCache, ape, normWeight, ropeSin, ropeCos, stateBlockTable, \
|
||||
cuSeqlens, seqUsed, startPos, cmpKvOut, workspace); \
|
||||
op.Process(); \
|
||||
} while (0)
|
||||
|
||||
#if (__CCE_AICORE__ == 220)
|
||||
template<uint8_t XLayout, uint8_t XDType, uint8_t Coff, uint8_t RotaryMode, uint8_t CacheMode, uint8_t TemplateId, uint8_t RopeDType>
|
||||
#else
|
||||
template<uint8_t XLayout, uint8_t XDType, uint8_t Coff, uint8_t RotaryMode, uint8_t CacheMode, uint8_t TemplateId>
|
||||
#endif
|
||||
__global__ __aicore__ void compressor(
|
||||
__gm__ uint8_t *x,
|
||||
__gm__ uint8_t *wKv,
|
||||
__gm__ uint8_t *wGate,
|
||||
__gm__ uint8_t *stateCache,
|
||||
__gm__ uint8_t *ape,
|
||||
__gm__ uint8_t *normWeight,
|
||||
__gm__ uint8_t *ropeSin,
|
||||
__gm__ uint8_t *ropeCos,
|
||||
__gm__ uint8_t *stateBlockTable,
|
||||
__gm__ uint8_t *cuSeqlens,
|
||||
__gm__ uint8_t *seqUsed,
|
||||
__gm__ uint8_t *startPos,
|
||||
__gm__ uint8_t *cmpKvOut,
|
||||
__gm__ uint8_t *stateCacheOut,
|
||||
__gm__ uint8_t *workspace,
|
||||
__gm__ uint8_t *tiling) {
|
||||
REGISTER_TILING_DEFAULT(optiling::CompressorTilingData);
|
||||
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
|
||||
GET_TILING_DATA_WITH_STRUCT(optiling::CompressorTilingData, tilingDataIn, tiling);
|
||||
if constexpr (static_cast<TEMPLATE_ID>(TemplateId) == TEMPLATE_ID::EMPTY_X) {
|
||||
return;
|
||||
}
|
||||
const optiling::CompressorTilingData *__restrict tilingData = &tilingDataIn;
|
||||
TPipe pipe;
|
||||
constexpr auto xLayout = static_cast<X_LAYOUT>(XLayout);
|
||||
constexpr auto xDtype = static_cast<X_DTYPE>(XDType);
|
||||
#if (__CCE_AICORE__ == 220)
|
||||
constexpr auto ropeDtype = static_cast<ROPE_DTYPE>(RopeDType);
|
||||
#endif
|
||||
constexpr auto coff = static_cast<COFF>(Coff);
|
||||
constexpr auto rotaryMode = static_cast<ROTARY_MODE>(RotaryMode);
|
||||
#if (__CCE_AICORE__ != 220)
|
||||
constexpr auto cacheMode = static_cast<CACHE_MODE>(CacheMode);
|
||||
#endif
|
||||
#if (__CCE_AICORE__ == 220)
|
||||
if constexpr (static_cast<TEMPLATE_ID>(TemplateId) == TEMPLATE_ID::PERF) {
|
||||
INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernelPerf, xLayout, xDtype, ropeDtype, coff, rotaryMode);
|
||||
} else {
|
||||
INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernel, xLayout, xDtype, ropeDtype, coff, rotaryMode);
|
||||
}
|
||||
#else
|
||||
if constexpr (static_cast<TEMPLATE_ID>(TemplateId) == TEMPLATE_ID::FULL_LOAD) {
|
||||
INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernelFullLoad, xLayout, xDtype, coff, rotaryMode, cacheMode);
|
||||
} else {
|
||||
INVOKE_COMPRESSOR_GENERAL_OP_IMPL(CompressorKernel, xLayout, xDtype, coff, rotaryMode, cacheMode);
|
||||
}
|
||||
#endif
|
||||
}
|
||||
19
csrc/attention/compressor_metadata/CMakeLists.txt
Normal file
19
csrc/attention/compressor_metadata/CMakeLists.txt
Normal file
@@ -0,0 +1,19 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
|
||||
file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
|
||||
if(NOT ENABLE_TEST AND NOT BENCHMARK)
|
||||
list(REMOVE_ITEM CURRENT_DIRS tests)
|
||||
endif()
|
||||
foreach(SUB_DIR ${CURRENT_DIRS})
|
||||
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
|
||||
add_subdirectory(${SUB_DIR})
|
||||
endif()
|
||||
endforeach()
|
||||
33
csrc/attention/compressor_metadata/op_host/CMakeLists.txt
Normal file
33
csrc/attention/compressor_metadata/op_host/CMakeLists.txt
Normal file
@@ -0,0 +1,33 @@
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
# Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
|
||||
# CANN Open Software License Agreement Version 2.0 (the "License").
|
||||
# Please refer to the License for details. You may not use this file except in compliance with the License.
|
||||
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
|
||||
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
|
||||
# See LICENSE in the root of the software repository for the full text of the License.
|
||||
# -----------------------------------------------------------------------------------------------------------
|
||||
|
||||
add_op_to_compiled_list()
|
||||
|
||||
if (BUILD_OPEN_PROJECT)
|
||||
target_sources(op_host_aclnn PRIVATE
|
||||
compressor_metadata_def.cpp
|
||||
)
|
||||
endif()
|
||||
|
||||
add_ops_compile_options(
|
||||
OP_NAME CompressorMetadata
|
||||
OPTIONS --cce-auto-sync=off
|
||||
-Wno-deprecated-declarations
|
||||
-mllvm -cce-aicore-hoist-movemask=false
|
||||
--op_relocatable_kernel_binary=true
|
||||
)
|
||||
|
||||
if (NOT BUILD_OPS_RTY_KERNEL)
|
||||
add_modules_sources(OPTYPE compressor_metadata ACLNNTYPE aclnn)
|
||||
target_include_directories(${OPHOST_NAME}_tiling_obj PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}
|
||||
)
|
||||
endif()
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
/*
|
||||
* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
|
||||
*/
|
||||
|
||||
#include "register/op_def_registry.h"
|
||||
|
||||
namespace ops {
|
||||
class CompressorMetadata : public OpDef {
|
||||
public:
|
||||
explicit CompressorMetadata(const char* name) : OpDef(name)
|
||||
{
|
||||
this->Input("ropeCos")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("ropeSin")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("cuSeqlens")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("startPos")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Input("kvBlockTable")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.AutoContiguous();
|
||||
this->Output("compressCos")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Output("compressSin")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_BF16})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
this->Output("slotMapping")
|
||||
.ParamType(REQUIRED)
|
||||
.DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
|
||||
.Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
|
||||
.UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
|
||||
|
||||
this->Attr("kvBlockSize").Int();
|
||||
this->Attr("slotMappingFormat").Int();
|
||||
this->Attr("cmpRatio").Int();
|
||||
this->Attr("actualNumReqs").Int();
|
||||
|
||||
this->AICore().AddConfig("ascend910b");
|
||||
this->AICore().AddConfig("ascend910_93");
|
||||
this->AICore().AddConfig("ascend950");
|
||||
}
|
||||
};
|
||||
|
||||
OP_ADD(CompressorMetadata);
|
||||
} // namespace ops
|
||||
@@ -0,0 +1,289 @@
|
||||
/*
|
||||
* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
|
||||
*/
|
||||
|
||||
#include "compressor_metadata_tiling.h"
|
||||
|
||||
#include <algorithm>
|
||||
|
||||
#include "register/op_def_registry.h"
|
||||
#include "tiling/platform/platform_ascendc.h"
|
||||
#include "tiling_base/error_log.h"
|
||||
|
||||
namespace optiling {
|
||||
namespace {
|
||||
constexpr uint32_t ROPE_COS_INDEX = 0;
|
||||
constexpr uint32_t ROPE_SIN_INDEX = 1;
|
||||
constexpr uint32_t CU_SEQLENS_INDEX = 2;
|
||||
constexpr uint32_t START_POS_INDEX = 3;
|
||||
constexpr uint32_t KV_BLOCK_TABLE_INDEX = 4;
|
||||
constexpr uint32_t COMPRESS_COS_INDEX = 0;
|
||||
constexpr uint32_t COMPRESS_SIN_INDEX = 1;
|
||||
constexpr uint32_t SLOT_MAPPING_INDEX = 2;
|
||||
constexpr uint32_t SLOT_MAPPING_FLAT = 1;
|
||||
constexpr uint32_t SLOT_MAPPING_BLOCK_OFFSET = 2;
|
||||
constexpr int64_t MAX_UINT32_VALUE = 0xFFFFFFFFLL;
|
||||
constexpr int64_t MAX_INT32_VALUE = 0x7FFFFFFFLL;
|
||||
|
||||
constexpr uint32_t TILING_KEY_FLOAT = 1;
|
||||
constexpr uint32_t TILING_KEY_FLOAT16 = 2;
|
||||
constexpr uint32_t TILING_KEY_BF16 = 3;
|
||||
constexpr uint32_t ALIGN_BYTES = 32;
|
||||
constexpr uint32_t BUFFER_NUM = 2;
|
||||
constexpr uint32_t MAX_TILE_ROWS = 512;
|
||||
constexpr uint32_t MAX_DATACOPY_BLOCK_COUNT = 4095;
|
||||
constexpr uint32_t ROWS_PER_CORE_TARGET = 64;
|
||||
constexpr uint32_t UB_RESERVED_BYTES = 16 * 1024;
|
||||
|
||||
uint32_t AlignUp(uint64_t value, uint32_t align)
|
||||
{
|
||||
return static_cast<uint32_t>((value + align - 1) / align * align);
|
||||
}
|
||||
|
||||
uint32_t CeilDiv(uint64_t lhs, uint64_t rhs)
|
||||
{
|
||||
return static_cast<uint32_t>((lhs + rhs - 1) / rhs);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
static ge::graphStatus CompressorMetadataTilingFunc(gert::TilingContext* context)
|
||||
{
|
||||
auto platformInfo = context->GetPlatformInfo();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, platformInfo);
|
||||
auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
|
||||
uint32_t aivCoreNum = ascendcPlatform.GetCoreNumAiv();
|
||||
if (aivCoreNum == 0) {
|
||||
aivCoreNum = ascendcPlatform.GetCoreNum();
|
||||
}
|
||||
if (aivCoreNum == 0) {
|
||||
OP_LOGE(context->GetNodeName(), "Failed to get AIV core num.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
uint64_t ubSize = 0;
|
||||
ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
|
||||
if (ubSize == 0) {
|
||||
OP_LOGE(context->GetNodeName(), "Failed to get UB size.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
|
||||
auto outputShape = context->GetOutputShape(COMPRESS_COS_INDEX);
|
||||
auto compressSinShape = context->GetOutputShape(COMPRESS_SIN_INDEX);
|
||||
auto slotMappingShape = context->GetOutputShape(SLOT_MAPPING_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, outputShape);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, compressSinShape);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, slotMappingShape);
|
||||
auto outputDimNum = outputShape->GetStorageShape().GetDimNum();
|
||||
if (outputDimNum < 2) {
|
||||
OP_LOGE(context->GetNodeName(), "compressCos dim num should be at least 2.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
if (compressSinShape->GetStorageShape().GetDimNum() != outputDimNum) {
|
||||
OP_LOGE(context->GetNodeName(), "compressCos and compressSin dim num mismatch.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
for (size_t dimIdx = 0; dimIdx < outputDimNum; ++dimIdx) {
|
||||
if (compressSinShape->GetStorageShape().GetDim(dimIdx) != outputShape->GetStorageShape().GetDim(dimIdx)) {
|
||||
OP_LOGE(context->GetNodeName(), "compressCos and compressSin shape mismatch.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
}
|
||||
int64_t numRows = outputShape->GetStorageShape().GetDim(0);
|
||||
int64_t ropeDim = outputShape->GetStorageShape().GetDim(outputDimNum - 1);
|
||||
if (numRows <= 0 || ropeDim <= 0 || numRows > MAX_UINT32_VALUE || ropeDim > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "compressCos shape is invalid.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
|
||||
auto ropeCosShape = context->GetInputShape(ROPE_COS_INDEX);
|
||||
auto ropeSinShape = context->GetInputShape(ROPE_SIN_INDEX);
|
||||
auto cuSeqlensShape = context->GetInputShape(CU_SEQLENS_INDEX);
|
||||
auto startPosShape = context->GetInputShape(START_POS_INDEX);
|
||||
auto kvBlockTableShape = context->GetInputShape(KV_BLOCK_TABLE_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, ropeCosShape);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, ropeSinShape);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, cuSeqlensShape);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, startPosShape);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, kvBlockTableShape);
|
||||
if (ropeCosShape->GetStorageShape().GetDimNum() != 2 || ropeSinShape->GetStorageShape().GetDimNum() != 2) {
|
||||
OP_LOGE(context->GetNodeName(), "ropeCos and ropeSin should be 2D tensors.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
int64_t ropeRows = ropeCosShape->GetStorageShape().GetDim(0);
|
||||
int64_t ropeCosDim = ropeCosShape->GetStorageShape().GetDim(1);
|
||||
if (ropeRows <= 0 || ropeCosDim <= 0 || ropeRows > MAX_UINT32_VALUE || ropeCosDim != ropeDim ||
|
||||
ropeSinShape->GetStorageShape().GetDim(0) != ropeRows ||
|
||||
ropeSinShape->GetStorageShape().GetDim(1) != ropeCosDim) {
|
||||
OP_LOGE(context->GetNodeName(), "ropeCos and ropeSin shape mismatch.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
int64_t cuSeqlensDim0 = cuSeqlensShape->GetStorageShape().GetDim(0);
|
||||
if (cuSeqlensDim0 < 2 || cuSeqlensDim0 > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "cuSeqlens dim0 should be at least 2.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
if (startPosShape->GetStorageShape().GetDimNum() != 1 ||
|
||||
startPosShape->GetStorageShape().GetDim(0) <= 0 ||
|
||||
startPosShape->GetStorageShape().GetDim(0) > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "startPos should be a non-empty 1D tensor.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
if (kvBlockTableShape->GetStorageShape().GetDimNum() != 2) {
|
||||
OP_LOGE(context->GetNodeName(), "kvBlockTable should be a 2D tensor.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
int64_t kvBlockTableRows = kvBlockTableShape->GetStorageShape().GetDim(0);
|
||||
int64_t kvBlockTableStride = kvBlockTableShape->GetStorageShape().GetDim(1);
|
||||
if (kvBlockTableRows <= 0 || kvBlockTableStride <= 0 || kvBlockTableRows > MAX_UINT32_VALUE ||
|
||||
kvBlockTableStride > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "kvBlockTable shape is invalid.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
|
||||
auto attrs = context->GetAttrs();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
|
||||
const int64_t* kvBlockSizePtr = attrs->GetInt(0);
|
||||
const int64_t* slotMappingFormatPtr = attrs->GetInt(1);
|
||||
const int64_t* cmpRatioPtr = attrs->GetInt(2);
|
||||
const int64_t* actualNumReqsPtr = attrs->GetInt(3);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, kvBlockSizePtr);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, slotMappingFormatPtr);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, cmpRatioPtr);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, actualNumReqsPtr);
|
||||
if (*kvBlockSizePtr <= 0 || *kvBlockSizePtr > MAX_INT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "kvBlockSize should be in (0, INT32_MAX].");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
if (*cmpRatioPtr <= 0 || *cmpRatioPtr > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "cmpRatio should be in (0, UINT32_MAX].");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
if (*slotMappingFormatPtr != SLOT_MAPPING_FLAT && *slotMappingFormatPtr != SLOT_MAPPING_BLOCK_OFFSET) {
|
||||
OP_LOGE(context->GetNodeName(), "slotMappingFormat should be 1(flat) or 2(block_offset).");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
auto slotMappingDimNum = slotMappingShape->GetStorageShape().GetDimNum();
|
||||
if ((*slotMappingFormatPtr == SLOT_MAPPING_FLAT &&
|
||||
(slotMappingDimNum != 1 || slotMappingShape->GetStorageShape().GetDim(0) != numRows)) ||
|
||||
(*slotMappingFormatPtr == SLOT_MAPPING_BLOCK_OFFSET &&
|
||||
(slotMappingDimNum != 2 || slotMappingShape->GetStorageShape().GetDim(0) != numRows ||
|
||||
slotMappingShape->GetStorageShape().GetDim(1) != 2))) {
|
||||
OP_LOGE(context->GetNodeName(), "slotMapping shape does not match slotMappingFormat.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
if (*actualNumReqsPtr <= 0 || *actualNumReqsPtr >= cuSeqlensDim0 ||
|
||||
*actualNumReqsPtr > startPosShape->GetStorageShape().GetDim(0) ||
|
||||
*actualNumReqsPtr > kvBlockTableRows ||
|
||||
*actualNumReqsPtr > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "actualNumReqs is invalid.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
|
||||
CompressorMetadataTilingData tilingData;
|
||||
tilingData.set_numRows(static_cast<uint32_t>(numRows));
|
||||
tilingData.set_numReqs(static_cast<uint32_t>(cuSeqlensDim0 - 1));
|
||||
tilingData.set_actualNumReqs(static_cast<uint32_t>(*actualNumReqsPtr));
|
||||
tilingData.set_ropeRows(static_cast<uint32_t>(ropeRows));
|
||||
tilingData.set_ropeDim(static_cast<uint32_t>(ropeDim));
|
||||
tilingData.set_kvBlockTableStride(static_cast<uint32_t>(kvBlockTableStride));
|
||||
tilingData.set_kvBlockSize(static_cast<uint32_t>(*kvBlockSizePtr));
|
||||
tilingData.set_slotMappingFormat(static_cast<uint32_t>(*slotMappingFormatPtr));
|
||||
tilingData.set_cmpRatio(static_cast<uint32_t>(*cmpRatioPtr));
|
||||
|
||||
auto ropeDesc = context->GetInputDesc(ROPE_COS_INDEX);
|
||||
auto ropeSinDesc = context->GetInputDesc(ROPE_SIN_INDEX);
|
||||
auto compressCosDesc = context->GetOutputDesc(COMPRESS_COS_INDEX);
|
||||
auto compressSinDesc = context->GetOutputDesc(COMPRESS_SIN_INDEX);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, ropeDesc);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, ropeSinDesc);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, compressCosDesc);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, compressSinDesc);
|
||||
auto ropeDtype = ropeDesc->GetDataType();
|
||||
if (ropeSinDesc->GetDataType() != ropeDtype ||
|
||||
compressCosDesc->GetDataType() != ropeDtype ||
|
||||
compressSinDesc->GetDataType() != ropeDtype) {
|
||||
OP_LOGE(context->GetNodeName(), "rope and compress output dtypes should match.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
|
||||
uint64_t tilingKey = 0;
|
||||
uint32_t dtypeSize = 0;
|
||||
if (ropeDtype == ge::DataType::DT_FLOAT) {
|
||||
tilingKey = TILING_KEY_FLOAT;
|
||||
dtypeSize = sizeof(float);
|
||||
} else if (ropeDtype == ge::DataType::DT_FLOAT16) {
|
||||
tilingKey = TILING_KEY_FLOAT16;
|
||||
dtypeSize = sizeof(uint16_t);
|
||||
} else if (ropeDtype == ge::DataType::DT_BF16) {
|
||||
tilingKey = TILING_KEY_BF16;
|
||||
dtypeSize = sizeof(uint16_t);
|
||||
} else {
|
||||
OP_LOGE(context->GetNodeName(), "Unsupported rope dtype.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
|
||||
uint32_t actualNumReqs = static_cast<uint32_t>(*actualNumReqsPtr);
|
||||
uint32_t cmpRatio = static_cast<uint32_t>(*cmpRatioPtr);
|
||||
if (static_cast<uint64_t>(ropeDim) * dtypeSize > MAX_UINT32_VALUE ||
|
||||
(static_cast<uint64_t>(actualNumReqs) + 1) * sizeof(int32_t) > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "tiling byte size exceeds UINT32_MAX.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
uint32_t ropeRowBytes = static_cast<uint32_t>(ropeDim) * dtypeSize;
|
||||
if (static_cast<uint64_t>(cmpRatio - 1) * ropeRowBytes > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "rope stride exceeds UINT32_MAX.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
uint32_t ropeRowBytesAligned = AlignUp(ropeRowBytes, ALIGN_BYTES);
|
||||
uint32_t slotCols = (*slotMappingFormatPtr == SLOT_MAPPING_FLAT) ? 1U : 2U;
|
||||
uint32_t reqTableBytes = AlignUp((static_cast<uint64_t>(actualNumReqs) + 1) * sizeof(int32_t), ALIGN_BYTES);
|
||||
uint64_t fixedUbBytes = static_cast<uint64_t>(reqTableBytes) * 3 + ALIGN_BYTES + UB_RESERVED_BYTES;
|
||||
uint64_t rowUbBytes =
|
||||
static_cast<uint64_t>(BUFFER_NUM) * ropeRowBytesAligned * 2 + slotCols * sizeof(int32_t) + sizeof(int32_t);
|
||||
if (rowUbBytes > MAX_UINT32_VALUE) {
|
||||
OP_LOGE(context->GetNodeName(), "row UB footprint exceeds UINT32_MAX.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
uint64_t minUbBytes = static_cast<uint64_t>(reqTableBytes) * 3 + ALIGN_BYTES + rowUbBytes;
|
||||
if (ubSize <= minUbBytes) {
|
||||
OP_LOGE(context->GetNodeName(), "UB size is insufficient for compressor metadata.");
|
||||
return ge::GRAPH_FAILED;
|
||||
}
|
||||
uint32_t tileRows = 1;
|
||||
if (ubSize > fixedUbBytes && rowUbBytes > 0) {
|
||||
tileRows = static_cast<uint32_t>((ubSize - fixedUbBytes) / rowUbBytes);
|
||||
tileRows = std::max(tileRows, 1U);
|
||||
}
|
||||
tileRows = std::min(tileRows, MAX_TILE_ROWS);
|
||||
tileRows = std::min(tileRows, MAX_DATACOPY_BLOCK_COUNT);
|
||||
|
||||
uint32_t usedCoreNum =
|
||||
std::min(aivCoreNum, std::max(1U, CeilDiv(static_cast<uint64_t>(numRows), ROWS_PER_CORE_TARGET)));
|
||||
tilingData.set_usedCoreNum(usedCoreNum);
|
||||
tilingData.set_tileRows(tileRows);
|
||||
tilingData.set_ropeRowBytes(ropeRowBytes);
|
||||
tilingData.set_ropeRowBytesAligned(ropeRowBytesAligned);
|
||||
tilingData.set_slotCols(slotCols);
|
||||
|
||||
size_t* workspaceSize = context->GetWorkspaceSizes(1);
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, workspaceSize);
|
||||
*workspaceSize = 0;
|
||||
context->SetBlockDim(usedCoreNum);
|
||||
context->SetTilingKey(tilingKey);
|
||||
|
||||
auto rawTilingData = context->GetRawTilingData();
|
||||
OP_CHECK_NULL_WITH_CONTEXT(context, rawTilingData);
|
||||
tilingData.SaveToBuffer(rawTilingData->GetData(), rawTilingData->GetCapacity());
|
||||
rawTilingData->SetDataSize(tilingData.GetDataSize());
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
static ge::graphStatus TilingParseForCompressorMetadata(gert::TilingParseContext* context)
|
||||
{
|
||||
return ge::GRAPH_SUCCESS;
|
||||
}
|
||||
|
||||
IMPL_OP_OPTILING(CompressorMetadata)
|
||||
.Tiling(CompressorMetadataTilingFunc)
|
||||
.TilingParse<CompressorMetadataCompileInfo>(TilingParseForCompressorMetadata);
|
||||
|
||||
} // namespace optiling
|
||||
@@ -0,0 +1,32 @@
|
||||
#ifndef COMPRESSOR_METADATA_TILING_H
|
||||
#define COMPRESSOR_METADATA_TILING_H
|
||||
|
||||
#include "register/tilingdata_base.h"
|
||||
|
||||
namespace optiling {
|
||||
BEGIN_TILING_DATA_DEF(CompressorMetadataTilingData)
|
||||
TILING_DATA_FIELD_DEF(uint32_t, numRows);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, numReqs);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, actualNumReqs);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, ropeRows);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, ropeDim);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, kvBlockTableStride);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, kvBlockSize);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, slotMappingFormat);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, cmpRatio);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, usedCoreNum);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, tileRows);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, ropeRowBytes);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, ropeRowBytesAligned);
|
||||
TILING_DATA_FIELD_DEF(uint32_t, slotCols);
|
||||
END_TILING_DATA_DEF;
|
||||
|
||||
REGISTER_TILING_DATA_CLASS(CompressorMetadata, CompressorMetadataTilingData)
|
||||
|
||||
struct CompressorMetadataCompileInfo {
|
||||
uint32_t coreNum;
|
||||
uint64_t ubSizePlatForm;
|
||||
};
|
||||
} // namespace optiling
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,38 @@
|
||||
/*
|
||||
* Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
|
||||
*/
|
||||
|
||||
#include "compressor_metadata.h"
|
||||
|
||||
extern "C" __global__ __aicore__ void compressor_metadata(
|
||||
GM_ADDR ropeCos,
|
||||
GM_ADDR ropeSin,
|
||||
GM_ADDR cuSeqlens,
|
||||
GM_ADDR startPos,
|
||||
GM_ADDR kvBlockTable,
|
||||
GM_ADDR compressCos,
|
||||
GM_ADDR compressSin,
|
||||
GM_ADDR slotMapping,
|
||||
GM_ADDR workspace,
|
||||
GM_ADDR tiling)
|
||||
{
|
||||
REGISTER_TILING_DEFAULT(CompressorMetadata::CompressorMetadataTilingData);
|
||||
GET_TILING_DATA(tilingData, tiling);
|
||||
KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_AIV_ONLY);
|
||||
|
||||
AscendC::TPipe pipe;
|
||||
|
||||
if (TILING_KEY_IS(1)) {
|
||||
CompressorMetadata::CompressorMetadataKernel<float> op;
|
||||
op.Init(&tilingData, &pipe);
|
||||
op.Process(ropeCos, ropeSin, cuSeqlens, startPos, kvBlockTable, compressCos, compressSin, slotMapping, workspace);
|
||||
} else if (TILING_KEY_IS(2)) {
|
||||
CompressorMetadata::CompressorMetadataKernel<half> op;
|
||||
op.Init(&tilingData, &pipe);
|
||||
op.Process(ropeCos, ropeSin, cuSeqlens, startPos, kvBlockTable, compressCos, compressSin, slotMapping, workspace);
|
||||
} else if (TILING_KEY_IS(3)) {
|
||||
CompressorMetadata::CompressorMetadataKernel<bfloat16_t> op;
|
||||
op.Init(&tilingData, &pipe);
|
||||
op.Process(ropeCos, ropeSin, cuSeqlens, startPos, kvBlockTable, compressCos, compressSin, slotMapping, workspace);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,428 @@
|
||||
#ifndef COMPRESSOR_METADATA_H
|
||||
#define COMPRESSOR_METADATA_H
|
||||
|
||||
#include "kernel_operator.h"
|
||||
|
||||
namespace CompressorMetadata {
|
||||
using namespace AscendC;
|
||||
|
||||
constexpr uint32_t SLOT_MAPPING_FLAT = 1;
|
||||
constexpr uint32_t ALIGN_BYTES = 32;
|
||||
constexpr uint32_t BUFFER_NUM = 2;
|
||||
constexpr int64_t MAX_INT32_VALUE = 0x7FFFFFFFLL;
|
||||
|
||||
__aicore__ inline uint32_t MinU32(uint32_t lhs, uint32_t rhs)
|
||||
{
|
||||
return lhs < rhs ? lhs : rhs;
|
||||
}
|
||||
|
||||
__aicore__ inline uint32_t MaxU32(uint32_t lhs, uint32_t rhs)
|
||||
{
|
||||
return lhs > rhs ? lhs : rhs;
|
||||
}
|
||||
|
||||
__aicore__ inline uint32_t AlignUpU32(uint32_t value, uint32_t align)
|
||||
{
|
||||
return (value + align - 1) / align * align;
|
||||
}
|
||||
|
||||
__aicore__ inline uint32_t Int32BytesU32(uint32_t elems)
|
||||
{
|
||||
return elems * static_cast<uint32_t>(sizeof(int32_t));
|
||||
}
|
||||
|
||||
__aicore__ inline void PipeMte2ToS()
|
||||
{
|
||||
event_t eventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE2_S));
|
||||
SetFlag<HardEvent::MTE2_S>(eventID);
|
||||
WaitFlag<HardEvent::MTE2_S>(eventID);
|
||||
}
|
||||
|
||||
__aicore__ inline void PipeMte3ToS()
|
||||
{
|
||||
event_t eventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::MTE3_S));
|
||||
SetFlag<HardEvent::MTE3_S>(eventID);
|
||||
WaitFlag<HardEvent::MTE3_S>(eventID);
|
||||
}
|
||||
|
||||
__aicore__ inline void PipeSToMte3()
|
||||
{
|
||||
event_t eventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::S_MTE3));
|
||||
SetFlag<HardEvent::S_MTE3>(eventID);
|
||||
WaitFlag<HardEvent::S_MTE3>(eventID);
|
||||
}
|
||||
|
||||
__aicore__ inline void PipeVToMte3()
|
||||
{
|
||||
event_t eventID = static_cast<event_t>(GetTPipePtr()->FetchEventID(HardEvent::V_MTE3));
|
||||
SetFlag<HardEvent::V_MTE3>(eventID);
|
||||
WaitFlag<HardEvent::V_MTE3>(eventID);
|
||||
}
|
||||
|
||||
struct CompressorMetadataTilingData {
|
||||
uint32_t numRows;
|
||||
uint32_t numReqs;
|
||||
uint32_t actualNumReqs;
|
||||
uint32_t ropeRows;
|
||||
uint32_t ropeDim;
|
||||
uint32_t kvBlockTableStride;
|
||||
uint32_t kvBlockSize;
|
||||
uint32_t slotMappingFormat;
|
||||
uint32_t cmpRatio;
|
||||
uint32_t usedCoreNum;
|
||||
uint32_t tileRows;
|
||||
uint32_t ropeRowBytes;
|
||||
uint32_t ropeRowBytesAligned;
|
||||
uint32_t slotCols;
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class CompressorMetadataKernel {
|
||||
public:
|
||||
__aicore__ inline CompressorMetadataKernel() {}
|
||||
|
||||
__aicore__ inline void Init(CompressorMetadataTilingData* tilingData, TPipe* pipe)
|
||||
{
|
||||
numRows_ = tilingData->numRows;
|
||||
actualNumReqs_ = tilingData->actualNumReqs;
|
||||
ropeRows_ = tilingData->ropeRows;
|
||||
ropeDim_ = tilingData->ropeDim;
|
||||
kvBlockTableStride_ = tilingData->kvBlockTableStride;
|
||||
kvBlockSize_ = tilingData->kvBlockSize;
|
||||
slotMappingFormat_ = tilingData->slotMappingFormat;
|
||||
cmpRatio_ = tilingData->cmpRatio;
|
||||
tileRows_ = tilingData->tileRows;
|
||||
ropeRowBytes_ = tilingData->ropeRowBytes;
|
||||
ropeRowBytesAligned_ = tilingData->ropeRowBytesAligned;
|
||||
slotCols_ = tilingData->slotCols;
|
||||
reqTableBytes_ = AlignUpU32(Int32BytesU32(actualNumReqs_ + 1), ALIGN_BYTES);
|
||||
ropeDimAligned_ = ropeRowBytesAligned_ / sizeof(T);
|
||||
ropePadElems_ = ropeDimAligned_ - ropeDim_;
|
||||
slotTileBytes_ = AlignUpU32(Int32BytesU32(tileRows_ * slotCols_), ALIGN_BYTES);
|
||||
blockTableTileBytes_ = AlignUpU32(Int32BytesU32(tileRows_), ALIGN_BYTES);
|
||||
|
||||
pipe->InitBuffer(prefixBuf_, reqTableBytes_);
|
||||
pipe->InitBuffer(startPosBuf_, reqTableBytes_);
|
||||
pipe->InitBuffer(cuSeqlensBuf_, reqTableBytes_);
|
||||
pipe->InitBuffer(blockTableBuf_, blockTableTileBytes_);
|
||||
pipe->InitBuffer(slotBuf_, slotTileBytes_);
|
||||
pipe->InitBuffer(cosQueue_, BUFFER_NUM, tileRows_ * ropeRowBytesAligned_);
|
||||
pipe->InitBuffer(sinQueue_, BUFFER_NUM, tileRows_ * ropeRowBytesAligned_);
|
||||
}
|
||||
|
||||
__aicore__ inline void Process(
|
||||
GM_ADDR ropeCos,
|
||||
GM_ADDR ropeSin,
|
||||
GM_ADDR cuSeqlens,
|
||||
GM_ADDR startPos,
|
||||
GM_ADDR kvBlockTable,
|
||||
GM_ADDR compressCos,
|
||||
GM_ADDR compressSin,
|
||||
GM_ADDR slotMapping,
|
||||
GM_ADDR)
|
||||
{
|
||||
ropeCosGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(ropeCos));
|
||||
ropeSinGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(ropeSin));
|
||||
cuSeqlensGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(cuSeqlens));
|
||||
startPosGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(startPos));
|
||||
kvBlockTableGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(kvBlockTable));
|
||||
compressCosGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(compressCos));
|
||||
compressSinGm_.SetGlobalBuffer(reinterpret_cast<__gm__ T*>(compressSin));
|
||||
slotMappingGm_.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t*>(slotMapping));
|
||||
|
||||
LocalTensor<int32_t> prefixLocal = prefixBuf_.Get<int32_t>();
|
||||
LocalTensor<int32_t> startPosLocal = startPosBuf_.Get<int32_t>();
|
||||
LocalTensor<int32_t> cuSeqlensLocal = cuSeqlensBuf_.Get<int32_t>();
|
||||
BuildCompressedPrefix(prefixLocal, startPosLocal, cuSeqlensLocal);
|
||||
|
||||
uint32_t validRows = static_cast<uint32_t>(prefixLocal.GetValue(actualNumReqs_));
|
||||
validRows = MinU32(validRows, numRows_);
|
||||
ProcessValidRows(prefixLocal, startPosLocal, validRows);
|
||||
ProcessPaddingRows(validRows);
|
||||
}
|
||||
|
||||
private:
|
||||
__aicore__ inline void BuildCompressedPrefix(
|
||||
LocalTensor<int32_t>& prefixLocal,
|
||||
LocalTensor<int32_t>& startPosLocal,
|
||||
LocalTensor<int32_t>& cuSeqlensLocal)
|
||||
{
|
||||
DataCopyExtParams startCopyParams{1, Int32BytesU32(actualNumReqs_), 0, 0, 0};
|
||||
DataCopyExtParams cuCopyParams{1, Int32BytesU32(actualNumReqs_ + 1), 0, 0, 0};
|
||||
DataCopyPadExtParams<int32_t> padParams{true, 0, 0, 0};
|
||||
DataCopyPad(startPosLocal, startPosGm_, startCopyParams, padParams);
|
||||
DataCopyPad(cuSeqlensLocal, cuSeqlensGm_, cuCopyParams, padParams);
|
||||
PipeMte2ToS();
|
||||
|
||||
uint32_t prefix = 0;
|
||||
prefixLocal.SetValue(0, 0);
|
||||
for (uint32_t reqIdx = 0; reqIdx < actualNumReqs_; ++reqIdx) {
|
||||
int64_t startPos = static_cast<int64_t>(startPosLocal.GetValue(reqIdx));
|
||||
int64_t seqLen = static_cast<int64_t>(cuSeqlensLocal.GetValue(reqIdx + 1)) -
|
||||
static_cast<int64_t>(cuSeqlensLocal.GetValue(reqIdx));
|
||||
uint32_t compressedRows = 0;
|
||||
if (startPos >= 0 && seqLen > 0) {
|
||||
compressedRows = static_cast<uint32_t>(((startPos + seqLen) / cmpRatio_) - (startPos / cmpRatio_));
|
||||
}
|
||||
prefix += compressedRows;
|
||||
prefixLocal.SetValue(reqIdx + 1, static_cast<int32_t>(prefix));
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void SplitRange(uint32_t totalRows, uint32_t& begin, uint32_t& end)
|
||||
{
|
||||
uint32_t blockIdx = GetBlockIdx();
|
||||
uint32_t blockNum = MaxU32(GetBlockNum(), 1);
|
||||
uint32_t rowsPerBlock = (totalRows + blockNum - 1) / blockNum;
|
||||
begin = MinU32(blockIdx * rowsPerBlock, totalRows);
|
||||
end = MinU32(begin + rowsPerBlock, totalRows);
|
||||
}
|
||||
|
||||
__aicore__ inline uint32_t FindRequest(LocalTensor<int32_t>& prefixLocal, uint32_t row)
|
||||
{
|
||||
uint32_t reqIdx = 0;
|
||||
while (reqIdx < actualNumReqs_ && static_cast<uint32_t>(prefixLocal.GetValue(reqIdx + 1)) <= row) {
|
||||
++reqIdx;
|
||||
}
|
||||
return reqIdx;
|
||||
}
|
||||
|
||||
__aicore__ inline void ProcessValidRows(
|
||||
LocalTensor<int32_t>& prefixLocal,
|
||||
LocalTensor<int32_t>& startPosLocal,
|
||||
uint32_t validRows)
|
||||
{
|
||||
uint32_t begin = 0;
|
||||
uint32_t end = 0;
|
||||
SplitRange(validRows, begin, end);
|
||||
if (begin >= end) {
|
||||
return;
|
||||
}
|
||||
|
||||
uint32_t reqIdx = FindRequest(prefixLocal, begin);
|
||||
uint32_t row = begin;
|
||||
while (row < end && reqIdx < actualNumReqs_) {
|
||||
uint32_t reqBegin = static_cast<uint32_t>(prefixLocal.GetValue(reqIdx));
|
||||
uint32_t reqEnd = static_cast<uint32_t>(prefixLocal.GetValue(reqIdx + 1));
|
||||
if (row >= reqEnd) {
|
||||
++reqIdx;
|
||||
continue;
|
||||
}
|
||||
uint32_t rowsInReq = MinU32(end - row, reqEnd - row);
|
||||
int64_t startPos = static_cast<int64_t>(startPosLocal.GetValue(reqIdx));
|
||||
uint32_t localCompressedIdx = row - reqBegin;
|
||||
// KV slot uses compressed position; RoPE uses the original group-start position.
|
||||
uint32_t compressedPos = static_cast<uint32_t>(startPos / cmpRatio_) + localCompressedIdx;
|
||||
ProcessRequestRows(reqIdx, row, compressedPos, rowsInReq);
|
||||
row += rowsInReq;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void ProcessRequestRows(
|
||||
uint32_t reqIdx,
|
||||
uint32_t outputRow,
|
||||
uint32_t compressedPos,
|
||||
uint32_t rows)
|
||||
{
|
||||
while (rows > 0) {
|
||||
uint32_t blockOffset = compressedPos % kvBlockSize_;
|
||||
uint32_t rowsToBlockEnd = kvBlockSize_ - blockOffset;
|
||||
uint32_t curRows = MinU32(rows, tileRows_);
|
||||
curRows = MinU32(curRows, rowsToBlockEnd);
|
||||
ProcessTile(reqIdx, outputRow, compressedPos, curRows);
|
||||
outputRow += curRows;
|
||||
compressedPos += curRows;
|
||||
rows -= curRows;
|
||||
}
|
||||
}
|
||||
|
||||
__aicore__ inline void ProcessTile(
|
||||
uint32_t reqIdx,
|
||||
uint32_t outputRow,
|
||||
uint32_t compressedPos,
|
||||
uint32_t rows)
|
||||
{
|
||||
uint32_t blockIdOffset = compressedPos / kvBlockSize_;
|
||||
if (blockIdOffset >= kvBlockTableStride_) {
|
||||
WriteInvalidTile(outputRow, rows);
|
||||
return;
|
||||
}
|
||||
|
||||
LocalTensor<int32_t> blockTableLocal = blockTableBuf_.Get<int32_t>();
|
||||
DataCopyExtParams blockCopyParams{1, Int32BytesU32(1), 0, 0, 0};
|
||||
DataCopyPadExtParams<int32_t> padParams{true, 0, 0, 0};
|
||||
uint64_t blockTableGmOffset = static_cast<uint64_t>(reqIdx) * kvBlockTableStride_ + blockIdOffset;
|
||||
DataCopyPad(blockTableLocal, kvBlockTableGm_[blockTableGmOffset], blockCopyParams, padParams);
|
||||
PipeMte2ToS();
|
||||
|
||||
int32_t blockId = blockTableLocal.GetValue(0);
|
||||
if (blockId < 0) {
|
||||
WriteInvalidTile(outputRow, rows);
|
||||
return;
|
||||
}
|
||||
|
||||
uint32_t blockOffset = compressedPos % kvBlockSize_;
|
||||
if (slotMappingFormat_ == SLOT_MAPPING_FLAT) {
|
||||
int64_t maxSlot = static_cast<int64_t>(blockId) * kvBlockSize_ + blockOffset + rows - 1;
|
||||
if (maxSlot > MAX_INT32_VALUE) {
|
||||
WriteInvalidTile(outputRow, rows);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
uint64_t lastRopePos = (static_cast<uint64_t>(compressedPos) + rows - 1) * cmpRatio_;
|
||||
if (lastRopePos >= ropeRows_) {
|
||||
WriteInvalidTile(outputRow, rows);
|
||||
return;
|
||||
}
|
||||
|
||||
CopyRopeTile(outputRow, compressedPos, rows);
|
||||
WriteSlotTile(outputRow, compressedPos, rows, blockId);
|
||||
}
|
||||
|
||||
__aicore__ inline void CopyRopeTile(uint32_t outputRow, uint32_t compressedPos, uint32_t rows)
|
||||
{
|
||||
LocalTensor<T> cosLocal = cosQueue_.AllocTensor<T>();
|
||||
LocalTensor<T> sinLocal = sinQueue_.AllocTensor<T>();
|
||||
uint64_t ropePos = static_cast<uint64_t>(compressedPos) * cmpRatio_;
|
||||
uint32_t srcStride = (cmpRatio_ - 1) * ropeRowBytes_;
|
||||
|
||||
DataCopyExtParams copyInParams{
|
||||
static_cast<uint16_t>(rows), ropeRowBytes_, srcStride, 0, 0};
|
||||
DataCopyPadExtParams<T> padParams{true, 0, static_cast<uint8_t>(ropePadElems_), 0};
|
||||
DataCopyPad(cosLocal, ropeCosGm_[ropePos * ropeDim_], copyInParams, padParams);
|
||||
DataCopyPad(sinLocal, ropeSinGm_[ropePos * ropeDim_], copyInParams, padParams);
|
||||
PipeMte2ToS();
|
||||
|
||||
DataCopyExtParams copyOutParams{
|
||||
static_cast<uint16_t>(rows), ropeRowBytes_, 0, 0, 0};
|
||||
uint64_t outputBase = static_cast<uint64_t>(outputRow) * ropeDim_;
|
||||
DataCopyPad(compressCosGm_[outputBase], cosLocal, copyOutParams);
|
||||
DataCopyPad(compressSinGm_[outputBase], sinLocal, copyOutParams);
|
||||
PipeMte3ToS();
|
||||
|
||||
cosQueue_.FreeTensor<T>(cosLocal);
|
||||
sinQueue_.FreeTensor<T>(sinLocal);
|
||||
}
|
||||
|
||||
__aicore__ inline void WriteSlotTile(
|
||||
uint32_t outputRow,
|
||||
uint32_t compressedPos,
|
||||
uint32_t rows,
|
||||
int32_t blockId)
|
||||
{
|
||||
LocalTensor<int32_t> slotLocal = slotBuf_.Get<int32_t>();
|
||||
int32_t blockOffset = static_cast<int32_t>(compressedPos % kvBlockSize_);
|
||||
if (slotMappingFormat_ == SLOT_MAPPING_FLAT) {
|
||||
int32_t slotBase = blockId * static_cast<int32_t>(kvBlockSize_) + blockOffset;
|
||||
for (uint32_t row = 0; row < rows; ++row) {
|
||||
slotLocal.SetValue(row, slotBase + static_cast<int32_t>(row));
|
||||
}
|
||||
} else {
|
||||
for (uint32_t row = 0; row < rows; ++row) {
|
||||
uint32_t slotOffset = row * slotCols_;
|
||||
slotLocal.SetValue(slotOffset, blockId);
|
||||
slotLocal.SetValue(slotOffset + 1, blockOffset + static_cast<int32_t>(row));
|
||||
}
|
||||
}
|
||||
|
||||
DataCopyExtParams slotCopyParams{1, Int32BytesU32(rows * slotCols_), 0, 0, 0};
|
||||
PipeSToMte3();
|
||||
DataCopyPad(slotMappingGm_[static_cast<uint64_t>(outputRow) * slotCols_], slotLocal, slotCopyParams);
|
||||
PipeMte3ToS();
|
||||
}
|
||||
|
||||
__aicore__ inline void WriteInvalidTile(uint32_t outputRow, uint32_t rows)
|
||||
{
|
||||
LocalTensor<T> cosLocal = cosQueue_.AllocTensor<T>();
|
||||
LocalTensor<T> sinLocal = sinQueue_.AllocTensor<T>();
|
||||
|
||||
Duplicate<T>(cosLocal, static_cast<T>(1.0f), rows * ropeDimAligned_);
|
||||
Duplicate<T>(sinLocal, static_cast<T>(0.0f), rows * ropeDimAligned_);
|
||||
PipeVToMte3();
|
||||
|
||||
DataCopyExtParams ropeCopyParams{
|
||||
static_cast<uint16_t>(rows), ropeRowBytes_, 0, 0, 0};
|
||||
uint64_t outputBase = static_cast<uint64_t>(outputRow) * ropeDim_;
|
||||
DataCopyPad(compressCosGm_[outputBase], cosLocal, ropeCopyParams);
|
||||
DataCopyPad(compressSinGm_[outputBase], sinLocal, ropeCopyParams);
|
||||
PipeMte3ToS();
|
||||
|
||||
cosQueue_.FreeTensor<T>(cosLocal);
|
||||
sinQueue_.FreeTensor<T>(sinLocal);
|
||||
|
||||
LocalTensor<int32_t> slotLocal = slotBuf_.Get<int32_t>();
|
||||
if (slotMappingFormat_ == SLOT_MAPPING_FLAT) {
|
||||
for (uint32_t row = 0; row < rows; ++row) {
|
||||
slotLocal.SetValue(row, -1);
|
||||
}
|
||||
} else {
|
||||
int32_t padOffset = static_cast<int32_t>(kvBlockSize_ - 1);
|
||||
for (uint32_t row = 0; row < rows; ++row) {
|
||||
uint32_t slotOffset = row * slotCols_;
|
||||
slotLocal.SetValue(slotOffset, -1);
|
||||
slotLocal.SetValue(slotOffset + 1, padOffset);
|
||||
}
|
||||
}
|
||||
PipeSToMte3();
|
||||
DataCopyExtParams slotCopyParams{1, Int32BytesU32(rows * slotCols_), 0, 0, 0};
|
||||
DataCopyPad(slotMappingGm_[static_cast<uint64_t>(outputRow) * slotCols_], slotLocal, slotCopyParams);
|
||||
PipeMte3ToS();
|
||||
}
|
||||
|
||||
__aicore__ inline void ProcessPaddingRows(uint32_t validRows)
|
||||
{
|
||||
if (validRows >= numRows_) {
|
||||
return;
|
||||
}
|
||||
uint32_t padRows = numRows_ - validRows;
|
||||
uint32_t begin = 0;
|
||||
uint32_t end = 0;
|
||||
SplitRange(padRows, begin, end);
|
||||
uint32_t row = validRows + begin;
|
||||
uint32_t padEnd = validRows + end;
|
||||
while (row < padEnd) {
|
||||
uint32_t curRows = MinU32(tileRows_, padEnd - row);
|
||||
WriteInvalidTile(row, curRows);
|
||||
row += curRows;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t numRows_{0};
|
||||
uint32_t actualNumReqs_{0};
|
||||
uint32_t ropeRows_{0};
|
||||
uint32_t ropeDim_{0};
|
||||
uint32_t kvBlockTableStride_{0};
|
||||
uint32_t kvBlockSize_{0};
|
||||
uint32_t slotMappingFormat_{0};
|
||||
uint32_t cmpRatio_{1};
|
||||
uint32_t tileRows_{1};
|
||||
uint32_t ropeRowBytes_{0};
|
||||
uint32_t ropeRowBytesAligned_{0};
|
||||
uint32_t slotCols_{1};
|
||||
uint32_t reqTableBytes_{0};
|
||||
uint32_t ropeDimAligned_{0};
|
||||
uint32_t ropePadElems_{0};
|
||||
uint32_t slotTileBytes_{0};
|
||||
uint32_t blockTableTileBytes_{0};
|
||||
|
||||
TBuf<TPosition::VECCALC> prefixBuf_;
|
||||
TBuf<TPosition::VECCALC> startPosBuf_;
|
||||
TBuf<TPosition::VECCALC> cuSeqlensBuf_;
|
||||
TBuf<TPosition::VECCALC> blockTableBuf_;
|
||||
TBuf<TPosition::VECCALC> slotBuf_;
|
||||
TQue<TPosition::VECOUT, BUFFER_NUM> cosQueue_;
|
||||
TQue<TPosition::VECOUT, BUFFER_NUM> sinQueue_;
|
||||
|
||||
GlobalTensor<T> ropeCosGm_;
|
||||
GlobalTensor<T> ropeSinGm_;
|
||||
GlobalTensor<T> compressCosGm_;
|
||||
GlobalTensor<T> compressSinGm_;
|
||||
GlobalTensor<int32_t> cuSeqlensGm_;
|
||||
GlobalTensor<int32_t> startPosGm_;
|
||||
GlobalTensor<int32_t> kvBlockTableGm_;
|
||||
GlobalTensor<int32_t> slotMappingGm_;
|
||||
};
|
||||
} // namespace CompressorMetadata
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,59 @@
|
||||
/*
|
||||
* Copyright (c) 2026 Huawei Technologies Co., Ltd.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
* SPDX-FileCopyrightText: Copyright contributors to the vllm-ascend project
|
||||
*/
|
||||
|
||||
#ifndef FUSED_GDN_GATING_TORCH_ADPT_H
|
||||
#define FUSED_GDN_GATING_TORCH_ADPT_H
|
||||
|
||||
#include <tuple>
|
||||
|
||||
namespace vllm_ascend {
|
||||
|
||||
std::tuple<at::Tensor, at::Tensor> npu_fused_gdn_gating(
|
||||
const at::Tensor& A_log,
|
||||
const at::Tensor& a,
|
||||
const at::Tensor& b,
|
||||
const at::Tensor& dt_bias,
|
||||
double beta = 1.0,
|
||||
double threshold = 20.0)
|
||||
{
|
||||
TORCH_CHECK(A_log.dim() == 1, "A_log should be 1-D [num_heads], got ", A_log.dim(), "D");
|
||||
TORCH_CHECK(dt_bias.dim() == 1, "dt_bias should be 1-D [num_heads], got ", dt_bias.dim(), "D");
|
||||
TORCH_CHECK(a.dim() == 2, "a should be 2-D [batch, num_heads], got ", a.dim(), "D");
|
||||
TORCH_CHECK(b.dim() == 2, "b should be 2-D [batch, num_heads], got ", b.dim(), "D");
|
||||
TORCH_CHECK(b.size(0) == a.size(0) && b.size(1) == a.size(1),
|
||||
"a and b must have the same shape, got a=", a.sizes(), " b=", b.sizes());
|
||||
TORCH_CHECK(a.scalar_type() == b.scalar_type(),
|
||||
"a and b must have the same dtype, got a=", a.scalar_type(),
|
||||
" b=", b.scalar_type());
|
||||
TORCH_CHECK(A_log.scalar_type() == dt_bias.scalar_type(),
|
||||
"A_log and dt_bias must have the same dtype, got A_log=",
|
||||
A_log.scalar_type(), " dt_bias=", dt_bias.scalar_type());
|
||||
TORCH_CHECK(a.size(1) == A_log.size(0),
|
||||
"a second dim (num_heads) must equal A_log first dim, got a.size(1)=",
|
||||
a.size(1), " A_log.size(0)=", A_log.size(0));
|
||||
|
||||
int64_t batch = a.size(0);
|
||||
int64_t num_heads = a.size(1);
|
||||
|
||||
at::Tensor g = at::empty({1, batch, num_heads},
|
||||
a.options().dtype(c10::kFloat));
|
||||
at::Tensor beta_output = at::empty({1, batch, num_heads}, b.options());
|
||||
|
||||
float beta_val = static_cast<float>(beta);
|
||||
float threshold_val = static_cast<float>(threshold);
|
||||
|
||||
EXEC_NPU_CMD(aclnnFusedGdnGating,
|
||||
A_log, a, b, dt_bias,
|
||||
beta_val,
|
||||
threshold_val,
|
||||
g, beta_output);
|
||||
|
||||
return std::make_tuple(g, beta_output);
|
||||
}
|
||||
|
||||
} // namespace vllm_ascend
|
||||
|
||||
#endif // FUSED_GDN_GATING_TORCH_ADPT_H
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user